SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 2: (A) Consensus motif of the first 4 and last 3 residues for the engineered TCRs and known TCRs; the height of the alphabet indicates the frequency of the amino acid occurring at the position. (B) (top) Average attention patterns of the functional encoder for the positive TCRs; (middle) average attention patterns of the structural encoder for the positive TCRs;(bottom) mutation frequency of the optimized TCRs compared to the templates; averaged across sequences of length 15. (C) Pairwise Miyazawa-Jernigan energy for the positive optimized sequence and its template with respect to the target peptide (left: SSYRRPVGI; right: CTPYDINQM).
$100\%$  ) are unique and all are novel (Appendix Table 10). In general, TCR-dWAE-based methods generate more valid and positive sequences compared to other methods. We also include the results on McPAS-TCR in the Appendix Table 9, where the observations are similar.
One advantage of TCR-dWAE is that  $\mathbf{z}_s$  implicitly constrains the sequence backbone, ensuring validity. Methods like genetic and MCTS, on the contrary, generate much fewer valid sequences without explicit control on validity. As a result, they could produce high binding score sequences because they are guided to do so, but most of them are invalid TCRs, or out-of-distribution samples that ERGO cannot correctly predict. Also, they require calling of an external evaluation model during generation at every iteration. Adding a regularizer to enforce validity could potentially improve these baselines' generation quality, but the two objectives would be difficult to balance and will further increase the computational burden. TCR-dWAE, on the contrary, can perform sequence engineering in one pass, requiring 10x less time (Appendix Table 10). Compared to VAE-based models like TCR-dVAE, TCR-dWAE is a deterministic model with simpler objectives, circumventing some well-known challenges in the training of VAE which we've found to be rather sensitive to hyperparameters (see Appendix C.3 and E.5 for more comparisons between model architectures and generation modes).
3.2.4 Analysis of Engineered TCRs
We define the first four and the last three residues as the conservative region and the rest as the hypervariable region. On average, the sequences generated by TCR-dAE have a similar conservative region motif as known TCR CDR3 $\beta$ s [41] (Fig. 2A). The average attention score of the functional encoder concentrates on the hypervariable part (Fig. 2B, top), which has a similar distribution as the mutation frequency of the engineered sequences compared to their templates (Fig. 2B, bottom). On the other hand, we can also observe some, though less frequent, changes in the conservative regions through modifying the  $z_{f}$ , while some residues in the hypervariable region are maintained (Fig. 2C). Furthermore, the attention patterns of the functional and structural encoders, despite having their own preferences, overlap at some positions and do not have a clear-cut separation. Neither pattern
%positive valid	zf↑	zs↓
full	0.23±0.02	0.163±0.015	0.013±0.004
Wass-	0.21±0.02	0.040±0.026	0.013±0.006
L_s+	0.18±0.01	0.238±0.029	0.008±0.000
L_i-	0.13±0.04	0.087±0.067	0.011±0.004
Table 2: Ablation study. Left: %positive valid for TCR engineering; Right: sample-based MMD between the embeddings of positive and negative samples. All means and standard deviations are calculated across five random seeds.
Embedding	Average reconstruction accuracy
Original (zf, zs)	0.9044 ± 0.0020
Random zf	0.7535 ± 0.0085
Fully random	0.6016 ± 0.0030
Table 3: Reconstruction accuracy with altered embeddings (VDJDB).
aligns with the hypervariable/conservative separation (Fig. 2B top and middle). These results indicate that both  $z_{f}$  and  $z_{s}$  learn to encode patterns from the entire sequence, and do not fit into the manual separation of a "functional" hypervariable region and a "structural" conservative region.
We also show some examples where residues introduced by the positive  $\mathbf{z}_f$  to the engineered sequence have the potential of forming lower-energy (i.e. more stable) interactions with the target peptide (Fig. 2C), using the classic Miyazawa-Jernigan contact energy [42], which is a knowledge-based pairwise energy matrix for amino acid interactions. These results demonstrate that the positive patterns carried by  $\mathbf{z}_f$  have learned some biologically relevant information such as the favored binding energy of the TCR-pMHC complex, and are successfully introduced to the engineered TCRs.
3.3 Ablation Study
For a quantitative comparison of the disentanglement, we calculate a sample-based MMD between the embeddings of positive and negative samples. Ideally, a disentangled latent space would make  $\mathbf{z}_f$  of the positive and negative groups very different, and  $\mathbf{z}_s$  indistinguishable. As shown in Table 2, removing the Wasserstein loss (Wass-) or the auxiliary classifier  $(\Psi -)$  leads to both poorer performance and poorer disentanglement.
We also attempt to explicitly control the information content of  $\mathbf{z}_s$  by adding another decoder using only  $\mathbf{z}_s(L_s+)$ . This in practice results in similar, even slightly better, disentanglement compared to the original model, but not necessarily better performance. These observations agree with our conclusions in Section 2.3.
3.4 Analysis of the Embedding Space
We select a balanced subset from the test set (332 positives and 332 negatives) and obtain their  $\mathbf{z}_f$  and  $\mathbf{z}_s$  embeddings. T-SNE [43] visualization in Fig. 3A shows that for each peptide, the binding and non-binding TCRs can be well separated by  $\mathbf{z}_f$  but not  $\mathbf{z}_s$  (more examples in Appendix Fig. 5). Also, among the positive samples,  $\mathbf{z}_f$  show strong clustering patterns corresponding to their peptide targets (Fig. 2B). As a result of the Wasserstein loss constraint, there is minimal correlation between  $\mathbf{z}_f$  and  $\mathbf{z}_s$  (Fig. 3C).
Furthermore, as shown in Table 3, the reconstruction accuracy is the highest when both  $(\mathbf{z}_f,\mathbf{z}_s)$  are provided, and significantly impaired when either embedding is missing. These results imply that the embedding space is properly disentangled as intended:  $\mathbf{z}_f$  and  $\mathbf{z}_s$  each encodes separate parts of the sequential patterns, and both are required for faithful reconstruction.
4 Conclusions and Outlook
In this work, we present TCR-dWAE, a disentangled Wasserstein autoencoder-based framework for massive TCR engineering. The disentanglement constraints separate key patterns related to binding
CTPYDINOM
Figure 3: (A) T-SNE patterns of  $\mathbf{z}_f$  and  $\mathbf{z}_s$  for selected peptides, colored by ground-truth labels. (B) t-SNE of the  $\mathbf{z}_f$  of the positive TCRs. Colors correspond to their binding peptide. (C) Correlation between dimensions of  $\mathbf{z}_f$  and  $\mathbf{z}_s$ , where the orange color corresponds to  $\mathbf{z}_f$  and green to  $\mathbf{z}_s$ .
("functional" embedding) from generic structural backbones ("structural" embedding). By modifying only the functional embedding, we are able to generate new TCRs with desired binding properties while preserving the backbone.
We are aware that our experiment is performed on a smaller subset selected by ERGO, whose limitations would bias the results. Also, our model assumes there is some "natural", yet unknown, disentanglement between features within the data. Thus, it could potentially be applied to several protein engineering tasks but would fail if the definition of functional terms is vague. With more high-throughput experimental data on TCR-peptide interaction, we will be able to train more comprehensive models on full TCR sequences. To our knowledge, ours is the first work that approaches protein sequence engineering from a style transfer perspective. As our framework mimics a vast amount of biological knowledge on function versus structure, we believe it can be further extended to a broader definition of protein functions or other molecular contexts.
TCR-peptide binding prediction Most human TCRs are heterodimers comprised of a  $\mathrm{TCR}\alpha$  chain and a  $\mathrm{TCR}\beta$  chain [21]. Each chain contains three complementarity-determining regions (CDRs), CDR1, 2 and 3, which are highly variable due to gene recombination. A large amount of literature investigates the prediction of the highly specific interaction between TCRs and peptides. Early methods use consensus motifs or simple structural patterns like k-mers as predictive features [44- 46, 24]. Other methods assess the binding energy for the TCR-peptide pair using physical or structural modeling [47, 48]. In recent years, several deep neural network architectures have been applied to this task [49, 26, 50, 23, 24, 51, 27], using both TCR and peptide sequences as input. Most studies
focus on the CDR3 $\beta$  sequence and some also include other information like the CDR3 $\alpha$ , depending on the availability of the data.
Computational biological sequence design Traditional methods of computational biological sequence design uses search algorithms, where random mutations are introduced and then selected by specified criteria [52-54]. Deep generative models like variational autoencoder (VAE) [55] or generative adversarial network (GAN) [6] have also been successfully applied to biological sequences such as DNA. Other recent approaches include reinforcement learning (RL) [7], and iterative refinement [56]. It is also possible to co-design sequences along with the 3D structure, such as in motif scaffolding problems [28-32].
Disentangled representation learning The formulation of DRL tasks depends on its nature, such as separating the "style" and the "content" of a piece of text [16] or a picture [8], or static (time-invariant) and dynamic (time-variant) elements of a video [9, 10]. The disentangled representations can be used for style transfer [11, 8, 16] and conditional generation [12, 13]. Several regularization methods have been used to achieve disentanglement with or without a priori knowledge, including controls on the latent space capacity [57, 58], adversarial loss [12], cyclic reconstruction [11], mutual information constraints [59, 16, 9], and self-supervising auxiliary tasks [9].
[1] Peter K Robinson. Enzymes: principles and biotechnological applications. Essays in biochemistry, 59:1, 2015.
[2] P Bork, L Holm, and C Sander. The immunoglobulin fold: structural classification, sequence patterns and common core. Journal of molecular biology, 242(4):309-320, 1994.
[3] Bissan Al-Lazikani, Arthur M Lesk, and Cyrus Chothia. Standard conformations for the canonical structures of immunoglobulins. Journal of molecular biology, 273(4):927-948, 1997.
[4] Yongcheng Jing, Yezhou Yang, Zunlei Feng, Jingwen Ye, Yizhou Yu, and Mingli Song. Neural style transfer: A review. IEEE transactions on visualization and computer graphics, 26(11):3365-3385, 2019.
[5] Martina Toshevka and Sonja Gievska. A review of text style transfer using deep learning. IEEE Transactions on Artificial Intelligence, 2021.
[6] Anvita Gupta and James Zou. Feedback gan for dna optimizes protein functions. Nature Machine Intelligence, 1(2):105-111, 2019.
[7] Christof Angermueller, David Dohan, David Belanger, Ramya Deshpande, Kevin Murphy, and Lucy Colwell. Model-based reinforcement learning for biological sequence design. In International conference on learning representations, 2019.
[8] Dmytro Kotovenko, Artsiom Sanakoyeu, Sabine Lang, and Bjorn Ommer. Content and style disentangle-ment for artistic style transfer. In Proceedings of the IEEE/CVF international conference on computer vision, pages 4422-4431, 2019.
[9] Yizhe Zhu, Martin Renqiang Min, Asim Kadav, and Hans Peter Graf. S3vae: Self-supervised sequential vae for representation disentanglement and data generation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 6538-6547, 2020.
[10] Jun Han, Martin Renqiang Min, Ligong Han, Li Erran Li, and Xuan Zhang. Disentangled recurrent Wasserstein autoencoder. In International Conference on Learning Representations, 2021.
[11] Hsin-Ying Lee, Hung-Yu Tseng, Jia-Bin Huang, Maneesh Singh, and Ming-Hsuan Yang. Diverse image-to-image translation via disentangled representations. In Proceedings of the European conference on computer vision (ECCV), pages 35-51, 2018.
[12] Emily L Denton et al. Unsupervised learning of disentangled representations from video. Advances in neural information processing systems, 30, 2017.
[13] Jun-Yan Zhu, Zhoutong Zhang, Chengkai Zhang, Jiajun Wu, Antonio Torralba, Josh Tenenbaum, and Bill Freeman. Visual object networks: Image generation with disentangled 3d representations. Advances in neural information processing systems, 31, 2018.
[14] Guillaume Lample, Sandeep Subramanian, Eric Smith, Ludovic Denoyer, Marc'Aurelio Ranzato, and Y-Lan Boureau. Multiple-attribute text rewriting. In International Conference on Learning Representations, 2019.
[53] Alexander J Hartemink, David K Gifford, and Julia Khodor. Automated constraint-based nucleotide sequence selection for dna computation. Biosystems, 52(1-3):227-235, 1999.
[54] Soo-Yong Shin, In-Hee Lee, Dongmin Kim, and Byoung-Tak Zhang. Multiobjective evolutionary optimization of dna sequences for reliable dna computing. IEEE transactions on evolutionary computation, 9 (2):143-158, 2005.
[55] Nathan Killoran, Leo J Lee, Andrew Delong, David Duvenaud, and Brendan J Frey. Generating and designing dna with deep generative models. arXiv preprint arXiv:1712.06148, 2017.
[56] Wengong Jin, Jeremy Wohlwend, Regina Barzilay, and Tommi Jaakkola. Iterative refinement graph neural network for antibody sequence-structure co-design. arXiv preprint arXiv:2110.04624, 2021.
[57] Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with a constrained variational framework. In International conference on learning representations, 2017.
[58] Christopher P Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner. Understanding disentangling in  $\beta$ -vae. arXiv preprint arXiv:1804.03599, 2018.
[59] Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, and Pieter Abbeel. Infogan: Interpretable representation learning by information maximizing generative adversarial nets. Advances in neural information processing systems, 29, 2016.
[60] Steven Henikoff and Jorja G Henikoff. Amino acid substitution matrices from protein blocks. Proceedings of the National Academy of Sciences, 89(22):10915-10919, 1992.
[61] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[62] Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. Advances in neural information processing systems, 28, 2015.
[63] Uniprot: the universal protein knowledgebase in 2021. Nucleic acids research, 49(D1):D480-D489, 2021.
[64] Moksh Jain, Emmanuel Bengio, Alex Hernandez-Garcia, Jarrid Rector-Brooks, Bonaventure FP Dossou, Chanakya Ajit Ekbote, Jie Fu, Tianyu Zhang, Michael Kilgour, Dinghuai Zhang, et al. Biological sequence design with gflownets. In International Conference on Machine Learning, pages 9786-9801. PMLR, 2022.
A Data preparation
A.1 Combination of data sources
TCR-peptide interaction data are obtained from VDJDB [36] and MCPAS [37]. Only peptides with  $>10$  observed pairs are used for downstream filtering. Because VDJDB and MCPAS only report interacting pairs, we first combine the dataset with the training set from NETTCR [26] which contains experimentally validated non-interacting pairs. Conflicting records are removed.
A.2 Filtering by ERGO performance
Since ERGO [23] trains two separate models for VDJDB and MCPAS, the following filtering process is also performed separately on the two datasets. For this and all subsequent ERGO-based predictions, we use the pre-trained weights from https://github.com/louzounlab/ERGO.
Additional negative samples are generated as follows: a random TCR sequence is first selected from the dataset and is paired with all existing peptides in the dataset. Any unobserved pair is treated as negative. We repeat this process until the size of the negative set is  $5\mathrm{x}$  that of the positive set. The expanded dataset is then provided to the respective ERGO model. Predictive performance is evaluated for each peptide. We keep the peptides with AUROC and AUPR  $>0.9$  and select those among top 10 positive sample counts (Table 5).
To ensure the specificity of TCR recognition in the following study, we did a second round of filtering of both the TCRs and the peptides. We pair all TCRs with at least one positive binding event and all peptides in the filtered dataset. Any unobserved pair is treated as negative. This dataset is then provided to ERGO. Performance is shown in Table 5. We discard peptides with AUPR  $< 0.7$  and TCRs that have more than one positive prediction or have at least one wrong prediction.
After that, we downsample all peptides to at most 400 positive TCRs. This number is chosen so that the resulting dataset is more balanced across peptides. The final number of samples for each peptide can be found in Table 6. To make sure the model captures peptide-specific information, for every TCR in the positive set, we add its unobserved pairings with other peptides to the negative set. We then split the TCRs into train/test/Validation sets with a ratio of 8:1:1, and put all pairings of each TCR to the respective subset, to ensure all TCRs in the test and validation sets are not seen in the training. For the training set, the positive samples are up-sampled by the negative/positive ratio of the original dataset.
B Model details
Notation	Meaning
Θf	functional encoder
Θs	structural encoder
Γ	decoder
Ψ	auxiliary functional classifier
{x, u, y}	a data point with TCR x, peptide u and binding label y
zf	functional embedding
zs	structural embedding
z	concatenation of {zf, zs}
x'	reconstructed/generated sequence from the decoder
x(i)	the probability distribution over amino acids at the i-th position in x
concat(x1, ...,xn)	concatenation of vectors {x1, ...,xn}
Table 4: Notations used for this paper. Sequences are represented as  $l \times |V|$  matrices where  $l$  is the length  $|V|$  is the number of amino acids.
We use density functions for simplicity. Let  $q_{\theta}(\mathbf{z} \mid \mathbf{x})$  be the encoder and  $p_{\gamma}(\mathbf{x} \mid \mathbf{z})$  be the decoder. We have the joint generative distribution:
$$
p (\mathbf {x}, \mathbf {z}) = p _ {\gamma} (\mathbf {x} \mid \mathbf {z}) p (\mathbf {z}),
$$
where  $p(\mathbf{z})$  is the prior. Also, we have the joint inference distribution:
$$
q (\mathbf {x}, \mathbf {z}) = q _ {\theta} (\mathbf {z} \mid \mathbf {x}) p _ {D} (\mathbf {x}),
$$
where  $p_D(\mathbf{x})$  is the data distribution.
$$
\begin{array}{l} I (\mathbf {X}; \mathbf {Z}) = \mathbb {E} _ {q (\mathbf {x}, \mathbf {z})} \log \frac {q (\mathbf {x} , \mathbf {z})}{p _ {D} (\mathbf {x}) q (\mathbf {z})} \\ = \mathbb {E} _ {p _ {D}} \sum_ {\mathbf {z}} p _ {D} (\mathbf {x}) q _ {\theta} (\mathbf {z} \mid \mathbf {x}) \log \frac {q _ {\theta} (\mathbf {z} \mid \mathbf {x}) p _ {D} (\mathbf {x})}{p _ {D} (\mathbf {x}) q (\mathbf {z})} \\ = \mathbb {E} _ {p _ {D}} \sum_ {\mathbf {z}} q _ {\theta} (\mathbf {z} \mid \mathbf {x}) \log \frac {q _ {\theta} (\mathbf {z} \mid \mathbf {x})}{q (\mathbf {z})} \\ = \mathbb {E} _ {p _ {D}} \sum_ {\mathbf {z}} q _ {\theta} (\mathbf {z} \mid \mathbf {x}) \log \frac {q _ {\theta} (\mathbf {z} \mid \mathbf {x})}{p (\mathbf {z})} - \mathbb {E} _ {p _ {D}} \sum_ {\mathbf {z}} q _ {\theta} (\mathbf {z} \mid \mathbf {x}) \log \frac {q (\mathbf {z})}{p (\mathbf {z})} \\ = \mathbb {E} _ {p _ {D}} \sum_ {\mathbf {z}} q _ {\theta} (\mathbf {z} \mid \mathbf {x}) \log \frac {q _ {\theta} (\mathbf {z} \mid \mathbf {x})}{p (\mathbf {z})} - \sum_ {\mathbf {z}} q (\mathbf {z}) \log \frac {q (\mathbf {z})}{p (\mathbf {z})} \\ = \mathbb {E} _ {p _ {D}} \left[ \mathbb {D} _ {\mathrm {K L}} \left(Q _ {\theta} (\mathbf {Z} \mid \mathbf {X}) \| P (\mathbf {Z})\right) \right] - \mathbb {D} _ {\mathrm {K L}} \left(Q (\mathbf {Z}) \| P (\mathbf {Z})\right). \\ \end{array}
$$
$$
\mathbb {D} _ {\mathrm {K L}} (Q (\mathbf {Z}) \| P (\mathbf {Z})) = \mathbb {E} _ {p _ {D}} \left[ \mathbb {D} _ {\mathrm {K L}} \left(Q _ {\theta} (\mathbf {Z} | \mathbf {X}) \| P (\mathbf {Z})\right) \right] - I (\mathbf {X}; \mathbf {Z}).
$$
B.2 Implementation and training details
All input sequences are padded to the same length (25). The peptide  $\mathbf{u}$  is represented as the average BLOSUM50 score [60] for all its amino acids. The model is trained from end to end using the Adam optimizer [61]. The first layer of the model is an embedding layer that transforms the one-hot encoded sequence  $\mathbf{x}$  into continuous-valued vectors of 128 dimensions:
$$
\mathbf {e} = W ^ {e m b} \mathbf {x}.
$$
Both  $\mathbf{z}_f$  and  $\mathbf{z}_s$  encoders are 1-layer transformer encoders with 8 attention heads and an intermediate size of 128. The transformer layer utilizes the multi-head self-attention mechanism. For each attention head  $i$ :
$$
Q _ {i} = W _ {i} ^ {Q} \mathbf {e}, K _ {i} = W _ {i} ^ {K} \mathbf {e}, V _ {i} = W _ {i} ^ {V} \mathbf {e}
$$
$$
\operatorname {A t t n} _ {i} (\mathbf {e}) = \operatorname {s o f t m a x} \left(\frac {Q _ {i} K _ {i} ^ {T}}{\sqrt {d _ {k}}}\right) V _ {i},
$$
where  $d_{k}$  is the dimension of  $Q_{i}$  and  $K_{i}$ . The outputs of the attention heads are then aggregated as follows:
$$
\operatorname {M u l t i h e a d} (\mathbf {e}) = \operatorname {c o n c a t} \left(\operatorname {A t t n} _ {1} (\mathbf {e}), \operatorname {A t t n} _ {2} (\mathbf {e}), \dots\right) W ^ {O}.
$$
A 2-layer MLP with a 128-dimension hidden layer is then built on top of the transformer (which has the same dimension as the input embeddings) to transform the output to the dimensions of  $\mathbf{z}_f$  and  $\mathbf{z}_s$ , respectively. The functional classifier is a 2-layer MLP with a 32-dimension hidden layer. The decoder is a 2-layer LSTM with 256 hidden dimensions.
The hyperparameters are selected with grid search and models with the best generation results are reported. Specifically, weights of all losses are selected from [1.0, 0.1]. The dimension of  $\mathbf{z}_f$  is fixed to 8 and  $\mathbf{z}_s$  to 32. We train each model with 200 epochs and a learning rate of  $1e - 4$ , and evaluate the checkpoint of every 50 epochs. We find the variance of the RBF kernel (for the calculation of the Wasserstein loss) does not have a strong impact on the results significantly, so the value is fixed to 1.0.
The model is trained with five different random seeds (42, 123, 456, 789, 987). We report the hyperparameter setting with the best average performance (i.e. one that generates the highest average number of qualified positive sequences for the well-classified peptides).
The hyperparameter settings of the models for comparison and visualization are:
$$
[ \beta_ {1} = 1. 0, \beta_ {2} = 0. 1, \text {e p o c h} = 2 0 0 ]
$$
where  $\beta$  's are weights of the losses:
For the visualization and analysis of the model trained on VDJDB, we use random seed  $= 789$ .
We use the scheduled sampling technique [62] for the LSTM decoder during training, where for each position in the input sequence, there is a 0.5 probability of using the previous predicted token, instead of the original token, to calculate the hidden state for the next position. This is employed to avoid the discrepancy between the training and the generation, as the former uses the original sequence to calculate the hidden states and the latter uses predicted tokens.
The model is trained on 2 rtx3090 GPUs with a batch size of 256 (128 per GPU). Training with 200 epochs typically takes  $\sim 4$  hours.
C Baseline methods