Title: Disentangled Wasserstein Autoencoder for T-Cell Receptor Engineering
Venue: NeurIPS 2023
Accepted: yes
Keywords: protein engineering

disentangled representation

T cell receptor
Abstract: In protein biophysics, the separation between the functionally important residues (forming the active site or binding surface) and those that create the overall structure (the fold) is a well-established and fundamental concept. Identifying and modifying those functional sites is critical for protein engineering but computationally non-trivial, and requires significant domain knowledge. To automate this process from a data-driven perspective, we propose a disentangled Wasserstein autoencoder with an auxiliary classifier, which isolates the function-related patterns from the rest with theoretical guarantees. This enables one-pass protein sequence editing and improves the understanding of the resulting sequences and editing actions involved. To demonstrate its effectiveness, we apply it to T-cell receptors (TCRs), a well-studied structure-function case. We show that our method can be used to alter the function of TCRs without changing the structural backbone, outperforming several competing methods in generation quality and efficiency, and requiring only 10\% of the running time needed by baseline models. To our knowledge, this is the first approach that utilizes disentangled representations for TCR engineering.

Disentangled Wasserstein Autoencoder for T-Cell Receptor Engineering

Abstract

In protein biophysics, the separation between the functionally important residues (forming the active site or binding surface) and those that create the overall structure (the fold) is a well-established and fundamental concept. Identifying and modifying those functional sites is critical for protein engineering but computationally nontrivial, and requires significant domain knowledge. To automate this process from a data-driven perspective, we propose a disentangled Wasserstein autoencoder with an auxiliary classifier, which isolates the function-related patterns from the rest with theoretical guarantees. This enables one-pass protein sequence editing and improves the understanding of the resulting sequences and editing actions involved. To demonstrate its effectiveness, we apply it to T-cell receptors (TCRs), a well-studied structure-function case. We show that our method can be used to alter the function of TCRs without changing the structural backbone, outperforming several competing methods in generation quality and efficiency, and requiring only 10% of the running time needed by baseline models. To our knowledge, this is the first approach that utilizes disentangled representations for TCR engineering.

1 Introduction

Decades of work in protein biology have shown the separation of the overall structure and the smaller 'functional' site, such as the generic structure versus the active site in enzymes [1], and the characteristic immunoglobulin fold versus the antigen-binding complementarity-determining region (CDR) in immunoproteins [2, 3]. The latter usually defines the protein's key function, but cannot work on its own without the stabilizing effect of the former. This dichotomy is similar to the contentstyle separation in computer vision [4] and natural language processing [5]. For efficient protein engineering, it is often desired that the overall structure is preserved while only the functionally relevant sites are modified. Traditional methods for this task require significant domain knowledge and are usually limited to specific scenarios. Several recent studies make use of deep generative models [6] or reinforcement learning [7] to learn from large-scale data the implicit generation and editing policies to alter proteins. Here, we tackle the problem by utilizing explicit functional features through the disentangled representation learning (DRL), where the protein sequence is separately embedded into a 'functional' embedding and a 'structural' embedding. This approach results in a more interpretable latent space and enables more efficient conditional generation and property manipulation for protein engineering.

DRL has been applied to the separation of 'style' and 'content' of images [8], or static and dynamic parts of videos [9, 10] for tasks such as style transfer [11, 8] and conditional generation [12, 13].

Attaining the aforementioned disentangled embeddings in discrete sequences such as protein sequences, however, is challenging because the functional residues can vary greatly across different proteins. To this end, several recent works on discrete sequences such as natural languages use adversarial objectives to achieve disentangled embeddings [14, 15]. Cheng, et al. [16] improves the disentanglement with a mutual information (MI) upper bound on the embedding space of a variational autoencoder (VAE). However, this approach relies on a complicated implementation of multiple losses that are approximated through various neural networks, and requires finding a dedicated trade-off among them, making the model difficult to train. To address these challenges, we propose a Wasserstein autoencoder (WAE) [17] framework that achieves disentangled embeddings with a theoretical guarantee, using a simpler loss function. Also, WAE could be trained deterministically, avoiding several practical challenges of VAE in general, especially on sequences [18, 19]. Our approach is proven to simultaneously maximize the mutual information (MI) between the data and the latent embedding space while minimizing the MI between the different parts of the embeddings, through minimizing the Wasserstein loss.

To demonstrate the effectiveness and utility of our method, we apply it to the engineering of T cell receptors (TCRs), which uses a similar structure fold as the immunoglobulin, one of the best-studied protein structures and a good example of separation of structure and functions. TCRs play an important role in the adaptive immune response [20] by specifically binding to peptide antigens [21] (Fig. 1A). Designing TCRs with higher affinity to the target peptide is thus of high interest in immunotherapy [22]. Various data-driven methods have been proposed to enhance the accuracy of TCR binding prediction [23-25, 12, 13, 26, 27]. However, there has been limited research on leveraging machine learning for TCR engineering. A related application is the motif scaffolding problem [28-32] where a structural 'scaffold' is generated supporting a fixed functional 'motif'. Here our goal is the opposite: to modify the functional parts of the sequence instead by directly introducing mutations to it.

We focus on the CDR3 β region of TCRs where there is sufficient data and a clearly defined functional role (peptide binding). Using a large TCR-peptide binding dataset, we empirically demonstrate that our method successfully separates key patterns related to binding ('functional' embedding) from generic structural backbones ('structural' embedding). Due to the lack of 3D structural data, we assume that similarity between the 'structure-related' parts of the sequence indicates similarity in the structure. Furthermore, by modifying only the functional embedding, our approach is able to generate new TCRs with desired binding properties while preserving the structural backbone, requiring only 10% of the running time needed by baseline models in comparison. We also note from our TCR engineering results that mutations can be introduced throughout the sequence, which implies that the model learns higher-level functional and structural patterns that could span the whole sequence, instead of looking for a universal clear cut between a 'functional segment' and a 'structural segment'.

We summarize our main contributions as follows:

To our knowledge, we are the first to formulate computational protein design as a style transfer problem and leverage disentangled embeddings for TCR engineering, resulting in more interpretable and efficient conditional generation and property manipulation.

We propose a disentangled Wasserstein autoencoder with an auxiliary classifier, which effectively isolates the function-related patterns from the rest with theoretical guarantees.

We show that by modifying only the functional embedding, we can edit TCR sequences into desired properties while maintaining their backbones, running 10 times faster than baselines.

2 Methods

2.1 Problem Formulation

We define our problem of TCR engineering task as follows: given a TCR sequence and a peptide it could not bind to, introduce a minimal number of mutations to the TCR so that it gains the ability to bind to the peptide. In the meantime, the modified TCR should remain a valid TCR, with no major changes in the structural backbone. Based on the assumption that only certain amino acids within the TCR should be responsible for peptide interactions, we can define two kinds of patterns in the TCR sequence: functional patterns and structural patterns . The former comprises the amino acids

/g8

/g2

/g8

/g1

/g3/g15/g21/g10/g19/g7/g8/g21/g12/g16/g15

/g7

/g5

/g2

/g2

/g2/g7/g22/g20/g20/g12/g7/g15

/g3/g15/g17/g22/g21/g1

/g20/g10/g18/g22/g10/g15/g8/g10

/g6

/g2

/g1

/g1

/g3

/g6/g4

/g5/g10/g8/g16/g15/g20/g21/g19/g22/g8/g21/g10/g9/g1

/g20/g10/g18/g22/g10/g15/g8/g10

/g8

/g2

/g4/g16/g9/g12/g11/g12/g10/g9/g1

/g8/g4

/g1

/g5

/g2

/g2

/g6 /g10/g14/g17/g13/g7/g21/g10

/g20/g10/g18/g22/g10/g15/g8/g10

/g6

/g1

/g6/g4

/g4/g16/g9/g12/g11/g12/g10/g9/g1

/g20/g10/g18/g22/g10/g15/g8/g10

TCR

MHC

peptide

TCRα

TCRβ

peptide

Viral infection or tumor

(B)

(A)

(C)

that define the peptide binding property. TCRs that bind to the same peptide should have similar functional patterns. The latter refers to all other patterns that do not relate to the function but could affect the validity. Following [23-25], we here limit our modeling to the CDR3 β region since it is the most active region for TCR binding. In the rest of the paper, we use 'TCR' and 'CDR3 β ' interchangeably.

2.2 Disentangled Wasserstein Autoencoder

Our proposed framework, named TCR-dWAE , leverages a disentangled Wasserstein autoencoder that learns embeddings corresponding to the functional and structural patterns. In this setting, the input data sample for the model is a triplet { x , u , y } , where x is the TCR sequence, u is the peptide sequence, and y is the binary label indicating the interaction.

In detail, given an input triplet { x , u , y } , the embedding space of x is separated into two parts: z = concat ( z f , z s ) , where z f is the functional embedding, and z s is the structural embedding. The schematic of the model is shown in Fig. 1B.

2.2.1 Encoders and Auxiliary Classifier

We use two separate encoders for the embeddings, respectively:

where e ∈ { s, f } correspond to 'structure' and 'function'.

First, the functional embedding z f is encoded by the functional encoder Θ f ( x ) . In order to make sure z f carries information about binding to the given peptide u , we introduce an auxiliary classifier Ψ( z f , u ) that takes z f and the peptide u as input and predicts the probability of positive binding label q Ψ ( y | z f , u ) ,

We define the binding prediction loss as binary cross entropy:

Second, the structural embedding z s is encoded by the structural encoder Θ s ( x ) . To enforce z s to contain all the information other than the peptide binding-related patterns, we leverage a sequence reconstruction loss, as will be discussed in detail in Section 2.2.3.

2.2.2 Disentanglement of the Embeddings

To attain disentanglement between z f and z s , we introduce a Wasserstein autoencoder regularization term in our loss function following [17], by minimizing the maximum mean discrepancy (MMD) between the distribution of the embeddings Z ∼ Q Z where z = concat ( z f , z s ) and an isotropic multivariate Gaussian prior Z 0 ∼ P Z where P Z = N (0 , I d ) :

The MMD is estimated as follows: given the embeddings { z 1 , z 2 , ..., z n } of an input batch of size n , we randomly sample from the Gaussian prior { ˜ z 1 , ˜ z 2 , ..., ˜ z n } with the same sample size. We then use the linear time unbiased estimator proposed by [33] to estimate the MMD:

where h (( z i , ˜ z i ) , ( z j , ˜ z j )) = k ( z i , z j )+ k (˜ z i , ˜ z j ) -k ( z i , ˜ z j ) -k ( z j , ˜ z i ) and k is the kernel function. Here we use a radial basis function (RBF) [34] with σ = 1 as the kernel.

By minimizing this loss, the joint distribution of the embeddings matches N (0 , I d ) , so that z f and z s are independent. Also, the Gaussian shape of the latent space facilitates generation [35, 17].

2.2.3 Decoder and Overall Training Objective

The decoder Γ takes z f , z s and peptide u as input and reconstructs the original sequence as x ′ . It also acts as a regularizer to enforce the structural embedding z s to contain all the information other than the peptide binding-related patterns. The reconstruction loss is the position-wise binary cross entropy between x and x ′ averaged across all positions of the sequence:

where l is the length of the sequence and x ( i ) is the probability distribution over the amino acids at the i -th position.

Combining all these losses with weights β 1 and β 2 , we have the final objective function, which then can be optimized through gradient descent in an end-to-end fashion:

2.3 Disentanglement Guarantee

To show how our method can guarantee the disentangled embeddings, we provide a novel perspective on the latent space of Wasserstein autoencoders [17] utilizing the variation of information following [16]. We introduce a measurement of disentanglement as follows:

where V I is the variation of information, V I ( X ; Y ) = H ( X ) + H ( Y ) -2 I ( X ; Y ) , which is a measure of independence between two random variables. For simplicity, we omit the condition U (peptide) in the following parts.

This measurement reaches 0 when Z f and Z s are totally independent, i.e. disentangled. It could further be simplified as:

Note that H ( X ) is a constant. Also, according to data processing inequality, as z f → x → y forms a Markov chain, we have I ( x ; z f ) ≥ I ( y ; z f ) . Combining the results above, we have the upper bound of the disentanglement objective:

Next, we show how our framework could minimize each part of the upper bound in (2).

Maximizing I ( X ; Z s ) Similar to [10], we have the following theorem:

Theorem 2.1 Given the encoder Q θ ( Z | X ) , decoder P γ ( X | Z ) , prior P ( Z ) , and the data distribution P D

where Q ( Z ) is the marginal distribution of the encoder when X ∼ P D and Z ∼ Q θ ( Z | X ) .

Theorem 2.1 shows that by minimizing the KL divergence between the marginal Q ( Z ) and the prior P ( Z ) , we jointly maximize the mutual information between the data X and the embedding Z , and minimize the KL divergence between Q θ ( Z | X ) and the prior P ( Z ) . Detailed proof of the theorem can be found in the Appendix B.1. This also applies to the two separate parts of Z , Z f and Z s . In practice, because the marginal cannot be measured directly, we minimize the aforementioned kernel MMD(1) instead.

As a result, there is no need for additional constraints on the information content of Z s because I ( X ; Z f ) is automatically maximized by the objective. We also empirically verify in Section 3.2.3 that supervision on Z s does not improve the model performance.

Maximizing I ( Y ; Z f ) I ( Y ; Z f ) has an lower bound as follows:

where q Ψ ( Y | Z f ) is the predicted probability by the auxiliary classifier Ψ (note we omitted the condition U here). Thus, maximizing the performance of classifier Ψ would maximize I ( Y ; Z f ) .

Minimizing I ( Z f ; Z s ) Minimization of the Wasserstein loss (1) forces the distribution of the embedding space Z to approach an isotropic multivariate Gaussian prior P Z = N (0 , I d ) , where all the dimensions are independent. Thus, the dimensions of Z will be independent, which also minimizes the mutual information between the two parts of the embedding, Z f and Z s .

As a conclusion, our objective can jointly minimize I ( Z f ; Z s ) and maximize I ( X ; Z s ) and I ( Y ; Z f ) . The former ensures independence between the embeddings, and the latter enforces them to learn separate information, achieving disentanglement.

3 Experiments

3.1 Setup

Datasets Interacting TCR-peptide pairs are obtained from VDJDB [36], merged with experimentallyvalidated negative pairs from NetTCR [26]. We then expended the negatives to 5x the size of positives by adding randomly shuffled TCR-peptide pairs and only selected the peptides that could be well-classified by ERGO[23], a state-of-the-art method for TCR-binding prediction (Appendix A.2). The selected samples are then balanced resulting in 26262 positive and 26262 negative pairs with 2918 unique TCRs and 10 peptides. We also performed the same experiments on the McPAS-TCR dataset [37] (Appendix A.1). For each peptide, we select from VDJDB an additional set of 5000 TCR sequences that do not overlap with the training set. This set is used for the TCR engineering experiment.

Implementation Details The TCR-dWAE model uses two transformer encoders [38] for Θ s , Θ f and a long short-term memory (LSTM) recurrent neural network decoder [39] for Γ . The auxiliary classifier Ψ is a 2- layer perceptron. Hyperparameters are selected through grid search. All results are averaged across five random seeds. See Appendix B.2 for more details.

3.2 TCR Engineering

3.2.1 Manipulating TCR Binding via Functional Embeddings

As shown in Fig. 1C, given any TCR sequence template x , we combine its original z s and a new functional embedding z ′ f that is positive for the target peptide u which is then fed to the decoder to generate a new TCR x ′ that could potentially bind to u while maintaining the backbone of x . For each generation, we obtain the z ′ f from a random positive sample in the dataset.

3.2.2 Metrics and Baselines

We use the following metrics to evaluate whether the engineered sequence x ′ (1) is a valid TCR sequence and (2) binds to the given peptide, which we denote as validity score and binding score , respectively:

The validity score r v evaluates whether the generated TCR follows similar generic patterns as naturally observed TCRs from TCRdb, an independent and much larger dataset [40]. We train another LSTM-based autoencoder on TCRdb. If the generated sequence can be reconstructed successfully by the autoencoder and has a high likelihood in the distribution of the latent space, we consider it as a valid TCR from the same distribution as the known ones (Appendix D.2). We use the sum of the reconstruction accuracy and the likelihood as the validity score and show that this metric separates true TCRs from other protein segments and random sequences (Appendix D.3).

For the binding score , the engineered sequence x ′ and the peptide u are fed into the pre-trained ERGO classifier and binding probability r b = ERGO ( x ′ , u ) is calculated.

We compare TCR-dWAE with the following baselines (see Appendix C for more details):

Random mutation-based For each iteration, a new random mutation is added to the template. The best mutated sequence (one with the highest ERGO score) is selected on either a sample-level ( greedy ), population-level ( genetic ), or without any selection ( naive rm ), then goes into the next round.

Generation-based This includes Monte Carlo tree search ( MCTS ), where random residues are iteratively added until the sequence reaches a maximum length. For each generation process, the leaf with the highest ERGO score is added to the results.

TCR-dVAE We also include a VAE baseline inspired by the objectives of IDEL from [16], which minimizes a MI upper bound to achieve disentanglement in the embedding space (Appendix C.3).

3.2.3 Main Results

We use peptides with AUROC > 0 . 8 by the classifier Ψ for the engineering task (SSYRRPVGI, TTPESANL, FRDYVDRFYKTLRAEQASQE, CTPYDINQM). The results are listed in Table 1. The cutoff for 'valid' is r v ≥ 1 . 25 , as it can successfully separate known TCRs from other protein sequences. The cutoff for 'positive' TCR is r b > 0 . 5 as the prediction scores by ERGO is rather extreme, either 0 or 1 . The vast majority of the generated valid sequences (ranging from 70% ∼

(A)

(C)

(B)

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 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 wellknown challenges in the training of V AE 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 β 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

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 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 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 z f of the positive and negative groups very different, and z s indistinguishable. As shown in Table 2, removing the Wasserstein loss (Wass -) or the auxiliary classifier ( Ψ -) leads to both poorer performance and poorer disentanglement.

We also attempt to explicitly control the information content of z s by adding another decoder using only 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 z f and 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 z f but not z s (more examples in Appendix Fig. 5). Also, among the positive samples, 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 z f and z s (Fig. 3C).

Furthermore, as shown in Table 3, the reconstruction accuracy is the highest when both ( z f , z s ) are provided, and significantly impaired when either embedding is missing. These results imply that the embedding space is properly disentangled as intended: z f and 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

(A)

(B)

(C)

('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 highthroughput 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.

5 Related Works

TCR-peptide binding prediction Most human TCRs are heterodimers comprised of a TCR α chain and a TCR β 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 [4446, 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 β sequence and some also include other information like the CDR3 α , 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].

References

Peter K Robinson. Enzymes: principles and biotechnological applications. Essays in biochemistry , 59:1, 2015.

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.

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.

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.

Martina Toshevska and Sonja Gievska. A review of text style transfer using deep learning. IEEE Transactions on Artificial Intelligence , 2021.

Anvita Gupta and James Zou. Feedback gan for dna optimizes protein functions. Nature Machine Intelligence , 1(2):105-111, 2019.

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.

Dmytro Kotovenko, Artsiom Sanakoyeu, Sabine Lang, and Bjorn Ommer. Content and style disentanglement for artistic style transfer. In Proceedings of the IEEE/CVF international conference on computer vision , pages 4422-4431, 2019.

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.

Jun Han, Martin Renqiang Min, Ligong Han, Li Erran Li, and Xuan Zhang. Disentangled recurrent wasserstein autoencoder. In International Conference on Learning Representations , 2021.

Hsin-Ying Lee, Hung-Yu Tseng, Jia-Bin Huang, Maneesh Singh, and Ming-Hsuan Yang. Diverse imageto-image translation via disentangled representations. In Proceedings of the European conference on computer vision (ECCV) , pages 35-51, 2018.

Emily L Denton et al. Unsupervised learning of disentangled representations from video. Advances in neural information processing systems , 30, 2017.

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.

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.

Vineet John, Lili Mou, Hareesh Bahuleyan, and Olga Vechtomova. Disentangled representation learning for non-parallel text style transfer. arXiv preprint arXiv:1808.04339 , 2018.

Pengyu Cheng, Martin Renqiang Min, Dinghan Shen, Christopher Malon, Yizhe Zhang, Yitong Li, and Lawrence Carin. Improving disentangled text representation learning with information-theoretic guidance. arXiv preprint arXiv:2006.00693 , 2020.

Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein auto-encoders. arXiv preprint arXiv:1711.01558 , 2017.

Samuel R Bowman, Luke Vilnis, Oriol Vinyals, Andrew M Dai, Rafal Jozefowicz, and Samy Bengio. Generating sentences from a continuous space. arXiv preprint arXiv:1511.06349 , 2015.

Partha Ghosh, Mehdi SM Sajjadi, Antonio Vergari, Michael Black, and Bernhard Schölkopf. From variational to deterministic autoencoders. arXiv preprint arXiv:1903.12436 , 2019.

Balachandra K Gorentla and Xiao-Ping Zhong. T cell receptor signal transduction in t lymphocytes. Journal of clinical & cellular immunology , 2012(Suppl 12):005, 2012.

Kinjal Shah, Amr Al-Haidari, Jianmin Sun, and Julhash U Kazi. T cell receptor (tcr) signaling in health and disease. Signal transduction and targeted therapy , 6(1):1-26, 2021.

Nicholas P Restifo, Mark E Dudley, and Steven A Rosenberg. Adoptive immunotherapy for cancer: harnessing the t cell response. Nature Reviews Immunology , 12(4):269-281, 2012.

Ido Springer, Hanan Besser, Nili Tickotsky-Moskovitz, Shirit Dvorkin, and Yoram Louzoun. Prediction of specific tcr-peptide binding from large dictionaries of tcr-peptide pairs. Frontiers in immunology , page 1803, 2020.

Anna Weber, Jannis Born, and María Rodriguez Martínez. Titan: T-cell receptor specificity prediction with bimodal attention networks. Bioinformatics , 37(Supplement_1):i237-i244, 2021.

Yicheng Gao, Yuli Gao, Yuxiao Fan, Chengyu Zhu, Zhiting Wei, Chi Zhou, Guohui Chuai, Qinchang Chen, He Zhang, and Qi Liu. Pan-peptide meta learning for t-cell receptor-antigen binding recognition. Nature Machine Intelligence , pages 1-14, 2023.

Alessandro Montemurro, Viktoria Schuster, Helle Rus Povlsen, Amalie Kai Bentzen, Vanessa Jurtz, William D Chronister, Austin Crinklaw, Sine R Hadrup, Ole Winther, Bjoern Peters, et al. Nettcr-2.0 enables accurate prediction of tcr-peptide binding by using paired tcr α and β sequence data. Communications biology , 4(1):1-13, 2021.

Filippo Grazioli, Pierre Machart, Anja Mösch, Kai Li, Leonardo V Castorina, Nico Pfeifer, and Martin Renqiang Min. Attentive variational information bottleneck for tcr-peptide interaction prediction. Bioinformatics , 39(1):btac820, 2023.

Erik Procko, Geoffrey Y Berguig, Betty W Shen, Yifan Song, Shani Frayo, Anthony J Convertine, Daciana Margineantu, Garrett Booth, Bruno E Correia, Yuanhua Cheng, et al. A computationally designed inhibitor of an epstein-barr viral bcl-2 protein induces apoptosis in infected cells. Cell , 157(7):1644-1656, 2014.

Jue Wang, Sidney Lisanza, David Juergens, Doug Tischer, Joseph L Watson, Karla M Castro, Robert Ragotte, Amijai Saragovi, Lukas F Milles, Minkyung Baek, et al. Scaffolding protein functional sites using deep learning. Science , 377(6604):387-394, 2022.

Brian L Trippe, Jason Yim, Doug Tischer, David Baker, Tamara Broderick, Regina Barzilay, and Tommi Jaakkola. Diffusion probabilistic modeling of protein backbones in 3d for the motif-scaffolding problem. arXiv preprint arXiv:2206.04119 , 2022.

Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eisenach, Woody Ahern, Andrew J Borst, Robert J Ragotte, Lukas F Milles, et al. Broadly applicable and accurate protein design by integrating structure prediction networks and diffusion generative models. BioRxiv , pages 2022-12, 2022.

Jin Sub Lee, Jisun Kim, and Philip M Kim. Score-based generative modeling for de novo protein design. Nature Computational Science , pages 1-11, 2023.

Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. The Journal of Machine Learning Research , 13(1):723-773, 2012.

Arthur Gretton, Dino Sejdinovic, Heiko Strathmann, Sivaraman Balakrishnan, Massimiliano Pontil, Kenji Fukumizu, and Bharath K Sriperumbudur. Optimal kernel choice for large-scale two-sample tests. Advances in neural information processing systems , 25, 2012.

Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 , 2013.

Dmitry V Bagaev, Renske MA Vroomans, Jerome Samir, Ulrik Stervbo, Cristina Rius, Garry Dolton, Alexander Greenshields-Watson, Meriem Attaf, Evgeny S Egorov, Ivan V Zvyagin, et al. Vdjdb in 2019: database extension, new analysis infrastructure and a t-cell receptor motif compendium. Nucleic Acids Research , 48(D1):D1057-D1062, 2020.

Nili Tickotsky, Tal Sagiv, Jaime Prilusky, Eric Shifrut, and Nir Friedman. Mcpas-tcr: a manually curated catalogue of pathology-associated t cell receptor sequences. Bioinformatics , 33(18):2924-2929, 2017.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation , 9(8):1735-1780, 1997.

Si-Yi Chen, Tao Yue, Qian Lei, and An-Yuan Guo. Tcrdb: a comprehensive database for t-cell receptor sequences with powerful search function. Nucleic Acids Research , 49(D1):D468-D474, 2021.

J Douglas Freeman, René L Warren, John R Webb, Brad H Nelson, and Robert A Holt. Profiling the t-cell receptor beta-chain repertoire by massively parallel sequencing. Genome research , 19(10):1817-1824, 2009.

Sanzo Miyazawa and Robert L Jernigan. Estimation of effective interresidue contact energies from protein crystal structures: quasi-chemical approximation. Macromolecules , 18(3):534-552, 1985.

Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research , 9(11), 2008.

Pradyot Dash, Andrew J Fiore-Gartland, Tomer Hertz, George C Wang, Shalini Sharma, Aisha Souquette, Jeremy Chase Crawford, E Bridie Clemens, Thi HO Nguyen, Katherine Kedzierska, et al. Quantifiable predictive features define epitope-specific t cell receptor repertoires. Nature , 547(7661):89-93, 2017.

Yao Tong, Jiayin Wang, Tian Zheng, Xuanping Zhang, Xiao Xiao, Xiaoyan Zhu, Xin Lai, and Xiang Liu. Sete: Sequence-based ensemble learning approach for tcr epitope binding prediction. Computational Biology and Chemistry , 87:107281, 2020.

Nicolas De Neuter, Wout Bittremieux, Charlie Beirnaert, Bart Cuypers, Aida Mrzic, Pieter Moris, Arvid Suls, Viggo Van Tendeloo, Benson Ogunjimi, Kris Laukens, et al. On the feasibility of mining cd8+ t cell receptor patterns underlying immunogenic peptide recognition. Immunogenetics , 70(3):159-168, 2018.

Kamilla Kjærgaard Jensen, Vasileios Rantos, Emma Christine Jappe, Tobias Hegelund Olsen, Martin Closter Jespersen, Vanessa Jurtz, Leon Eyrich Jessen, Esteban Lanzarotti, Swapnil Mahajan, Bjoern Peters, et al. Tcrpmhcmodels: Structural modelling of tcr-pmhc class i complexes. Scientific reports , 9(1): 1-12, 2019.

I-Hsin Liu, Yu-Shu Lo, and Jinn-Moon Yang. Genome-wide structural modelling of tcr-pmhc interactions. BMC genomics , 14(5):1-13, 2013.

David S Fischer, Yihan Wu, Benjamin Schubert, and Fabian J Theis. Predicting antigen specificity of single t cells based on tcr cdr 3 regions. Molecular systems biology , 16(8):e9416, 2020.

Pieter Moris, Joey De Pauw, Anna Postovskaya, Sofie Gielis, Nicolas De Neuter, Wout Bittremieux, Benson Ogunjimi, Kris Laukens, and Pieter Meysman. Current challenges for unseen-epitope tcr interaction prediction and a new perspective derived from image classification. Briefings in Bioinformatics , 22(4): bbaa318, 2021.

Yiren Jian, Erik Kruus, and Martin Renqiang Min. T-cell receptor-peptide interaction prediction with physical model augmented pseudo-labeling. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining , pages 3090-3097, 2022.

Frances H Arnold. Design by directed evolution. Accounts of chemical research , 31(3):125-131, 1998.

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.

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.

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.

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.

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.

Christopher P Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner. Understanding disentangling in β -vae. arXiv preprint arXiv:1804.03599 , 2018.

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.

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.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

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.

Uniprot: the universal protein knowledgebase in 2021. Nucleic acids research , 49(D1):D480-D489, 2021.

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 5x 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

B.1 Proof of Theorem 1

We use density functions for simplicity. Let q θ ( z | x ) be the encoder and p γ ( x | z ) be the decoder. We have the joint generative distribution:

APPENDIX

where p ( z ) is the prior. Also, we have the joint inference distribution:

where p D ( x ) is the data distribution.

Thus,

B.2 Implementation and training details

All input sequences are padded to the same length (25). The peptide 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 x into continuous-valued vectors of 128 dimensions:

Both z f and 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 :

where d k is the dimension of Q i and K i . The outputs of the attention heads are then aggregated as follows:

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 z f and 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 z f is fixed to 8 and z s to 32. We train each model with 200 epochs and a learning rate of 1 e -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:

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 ∼ 4 hours.

where β 's are weights of the losses:

C Baseline methods

We compare our model with two types of methods for the generation of the optimized TCR x ′ : (1) mutationbased, which iteratively adds random mutations to the template sequence; and (2) generation-based, which generates novel sequences of the pre-determined length range. For both types of methods, the modified/generated sequences are selected by peptide binding scores from the respective pre-trained ERGO. The experiments are performed on each peptide in the dataset independently.

C.1 Mutation-based baselines

Random mutation ( naive rm ) The TCR is randomly mutated by one amino acid for 8 times progressively. This process is repeated for 10 runs for each TCR and the resulting one with the highest ERGO prediction score is reported.

Greedy mutation ( greedy ) For each TCR, 10 randomly mutated sequences are generated, each with one amino acid difference from the original sequence. Among the 10 mutated sequences, we select the one that gives the highest binding prediction with the given peptide as the template for the next run. This process is repeated 8 times.

Genetic algorithm ( genetic ) Let M be the sample size. For each TCR, 10 randomly mutated sequences are generated, each with one amino acid difference from the original sequence. All mutated sequences along with the original TCRs are then pooled together, and the top M sequences that give the highest binding prediction are used as the input for the next run. This process is repeated 8 times.

C.2 Generation-based baselines

Monte Carlo tree search ( MCTS ) TCRs are generated by adding amino acids iteratively, resulting in a search tree. When the TCR length reaches 10, the binding score is estimated by ERGO. For each iteration, a random node is selected for expansion and then evaluated by ERGO, and the scores of all its parent nodes are updated accordingly. The tree expansion ends when the length reaches 20. For every generation process, the highest leaf node is added to the output TCR set.

C.3 TCR-dVAE

TCR-dVAE uses a similar objective as IDEL [16], which is a VAE with a mutual information constraint on the latent space. For training, the loss comprises of the following components:

The reconstruction loss: L recon ( x , x ′ )

The KL divergence term for VAE: L KL = D KL ( q θ ( z s , z f | x ) ∥ p ( z s , z f )) .

The reconstruction loss given z s : L s ( x , Φ( z s )) where Φ is an auxiliary decoder.

The classification loss given z f : L f _ cls (ˆ y, y )

The sample-based MI upper bound between the embeddings: L MI ( z f , z s ) . This requires an approximation of the conditional distribution p ( z f | z s ) , which is achieved by a separate neural network.

Here we use our own notations, not the ones used in the original paper, for better comparison.

For the hyperparameters, we use 10 . 0 weight for L recon , L s and L f _ cls , 1 . 0 for the other terms, and a learning rate of 1 e -4 . In practice, we perform annealing [18] on L KL and L MI where their weights gradually increase from 0 to 1 during the first 100 epochs and remain for the rest of the epochs, to make sure the embeddings are as informative as possible. We observe that other hyperparameter settings would easily lead to severe posterior collapse where the embedding is ignored for the reconstruction.

D Evaluation of the optimized sequences

D.1 Training of the autoencoder

We train an LSTM-based autoencoder, which we denote as TCR-AE0, on the 277 million TCR sequences from TCRdb [40]. TCR-AE0 has a latent space of dimension 16 and is trained for 50,000 steps with a batch size of 256.

D.2 Validity score

The validity score combines two scores calculated from TCR-AE0:

The reconstruction-based score is calculated as

where lev ( x ′ , TCR-AE0 ( x ′ )) is the Levenstein distance between the original sequence and the reconstructed sequence, and l ( x ′ ) is the length of the reconstructed sequence. Higher r r means x ′ is better reconstructed from TCR-AE0 and is thus more likely to be a valid TCR sequence.

The density-based score calculates whether the embedding of x ′ follows the same distribution as known TCRs. We learn a Gaussian mixture model from the latent embeddings of known TCRs from TCR-AE0. The likelihood of the embedding e ′ of x ′ from TCR-AE0 falling in the same Gaussian mixture distribution is denoted as P ( e ′ ) . The density-based score is calculated as

where τ = 10 . Higher r d means the latent embedding of x ′ from TCR-AE0 is more likely to follow the same distribution as other valid TCR sequences.

We then define the validity score as r v = r r + r d .

D.3 Validation of the metrics

We compare the TCR-AE-derived evaluation metric scores of three different sources:

(1) all unique CDR3 β sequences from VDJDB.

(2) random segments of length 8 -18 (which is the most frequent lengths of CDR3 β sequences) from random uniport [63] protein sequences of the same size as (1). The conservative 'C' at the beginning and 'F' at the end are added to the segments.

(3) random shuffling of the sequences from (1), where the first 'C' and the last 'F' are kept

We show in Fig. 4 that for both two scores r d and r r , as well as their sum, CDR3 β sequences score much higher than random proteins or shuffled sequences. This shows these scores could be effectively used for the estimation of TCR sequence validity. We choose r v > 1 . 25 as the criteria for valid sequences as it rejects most negatives.

E Extended Results

E.1 Comparison of TCR Engineering Performance

We find consistently improved performance of our method over the baselines in both VDJDB (Table 1) and McPAS-TCR (Table 9). Also, the majority of generated sequences are unique (Table 10) and all are novel (not observed in the original dataset; statistics not shown).

The relative performance between methods in comparison still holds when all peptides, including those that are poorly classified, are taken into account (Table 7). However, when the results are separated by individual peptides (Table 8), we observe that for the poorly classified peptides (AUC<0.8; see Methods), using a positive z f ( TCR-dWAE row) has little difference from using a random z f ( TCR-dWAE (null) row). This further highlights the importance of well-defined and well-classified functional labels. Otherwise, z f does not encode the proper functional information, so it becomes uninformative and does not alter the sequence function better than a random embedding.

E.2 Analysis of the Model

We show extended z f and z s T-SNE patterns in Fig. 5, colored by the ground truth label as well as the predicted label. For the well-classified peptides, there is a clear separation of positives and negatives in the z f space but not z s . There are cases where the true positives are not separable from the true negatives using z f , but the predicted positives and the predicted negatives (by the function classifier Ψ ) are still separated. We consider the latter as a problem with data quality and classification accuracy, not embedding. Meanwhile, the classifier shows consistent performance over the peptides across random seeds with (Fig. 6, left) and without (Fig. 6, right) the Wasserstein loss.

As a result of the Wasserstein loss, the distribution of the embedding space is closer to a multivariate Gaussian (Fig. 7A. It becomes less regularized without the Wasserstein loss (Fig. 7B). Contrary to z f (Fig. 3B in the main text), the T-SNE of z s and first-layer embedding of the encoder for the positive samples cannot distinguish the binding targets from each other (Fig. 8A).

E.3 Analysis of the Generated Sequences

In addition to the results presented in the main text, we also selected 500 random positive and negative sequences from the training set and replaced their z f with the most positive/negative one in the subset. The generated sequences using their original z s and the new z f have binding scores mostly related to the z f , regardless of whether the z s source is positive or negative. This shows z f can be used to encode and transfer binding information, which lays the foundation for the following TCR engineering experiments (Fig. 8B).

The generated sequences have a similar length distribution as their templates (Fig. 8C), meaning no drastic changes are made. We further find that the z s of the modified TCRs show high cosine similarity with those of their templates, while the z f are more similar to the z f used for their generation (Fig. 8D), but not with that of the template. These show that the modified TCRs preserve the 'structural' information from z s and incorporate the new 'functional' information from the modified z f .

As in Fig. 8E, there is no significant correlation between the sequence validity and the binding scores. Also, we find that most binding scores predicted by ERGO are rather extreme (either 1 or 0).

E.4 Interpolation

For each peptide, we perform interpolation of z f between 100 random pairs for 10 steps with z ( r ) f = r z (1) f + (1 -r ) z (2) f , r ∈ [0 , 1] , while z s remains the same. Fig. 8F shows that interpolating between positive z f pairs preserves the positive binding score compared to negative pairs. Here we show one example of interpolation, where the first column is the generated sequence and the second column is the predicted binding score:

Seq 1: CASTESDRRSQNTQYF Seq 2: CASSLSTFTANTAQLFF Peptide: CTPYDINQM

These results further prove that z s could be used to transfer binding information independent of the z s , and that the latent space is well-regularized. Also, throughout the interpolation, the hypervariable region only undergoes minor changes. This shows z s could preserve the structural backbone information within both the conservative and hypervariable regions.

E.5 Extended comparison between sampling methods and model architecture

We experiment with several different generation modes and model architecture. The results are shown in Table 11.

E.5.1 Source of positive embeddings

In addition to using a z f from a random positive sample (labeled as random pos ), we also try two other means of obtaining positive z f 's: (1) best : we use the z f of the sample with the highest classification score from the auxiliary classifier Ψ (2) avg : we use the average z f for all positive samples. In both cases, the same z f is used for all generations. We perform the same experiments with TCR-dVAE , by using the mean of the z f of the positive sample. We observe that all three methods have comparable performance for TCR-dWAE , indicating a well-regularized latent space.

E.5.2 Stochastic generation

As TCR-dVAE is a probabilistic model, we also compared with a random sampling-based method. We use both the predicted mean and variance of the positive source to generate 5 random z f 's and take the average results of all the generated samples (labeled TCR-dVAE-rand ). In this scenario, only the random pos and best modes are applicable. We observe a notable decrease in the performance as the quality of generation varies with the random sampling, probably due to a less regularized latent space.

E.5.3 Transformer decoder

We further test a TCR-dWAE model with a transformer decoder ( TCR-dWAE-attn ) instead of a LSTM decoder. The performance is comparable with those with LSTM (Table 11) while the autoregressive generation is notably slower (Table 10)

E.5.4 Modeling the hypervariable region only

We train a model only using the hypervariable region and run the sequence modification pipeline as in full CDR3beta ( TCR-dWAE (trimmed) ). After the hypervariable region is modified, we add back the conservative region of the input template for each generated sequence.

As shown in Table 11, the performance is worse than modeling the full sequence. To our surprise, it seems that the lower performance is caused by a massive drop in the number of positive sequences, while the sequence validity is not affected much by the trimming, as can be seen from the number of valid sequences. This may indicate some interactions between the two regions that are essential for positive binding. Thus, the modeling and generation should be performed on the entire CDR3 β region, instead of on manually defined segments.

E.6 Comparison between evaluation metrics

E.6.1 Binding prediction

Besides ERGO [23], NetTCR[26] is another sequence-based TCR-peptide binding model with competitive performance. However, NetTCR is only trained on three peptides and shows the best performance when both the α and β chain are provided, which is mostly not available in our dataset. Thus, ERGO provides a more efficient evaluation method for single-chain data and a greater variety of peptides. For the comparison between model architectures, we re-train a single-chain NetTCR on our own training set with more peptides. The evaluation result on the generated sequences in Table 1 in the manuscript (reporting the %positive valid only) with the two classifiers is as follows. We show in Table 12 that the conclusions are similar between the two models.

E.6.2 Uniqueness and novelty score

We calculate the novelty score following [64] for each individual sequence from the positive valid subset in Table 1:

where D 0 is the set of templates and the distance measure d is the Levenshtein distance relative to the length of the longer sequence. This score can then be used to select the 'novel' sequences. We also include a held-out validation set as a reference. We find that around 1 / 3 ∼ 1 / 2 of the sequences remain at cutoff 0.2 and all are excluded at cutoff 0.4 (Table 14). The same pattern is observed in the validation set, which indicates a similar distribution of 'novelty' between the generated sequences and the raw data.

We also report the same novelty and diversity metrics for the positive valid subset as in [64]. We would like to point out that:

(1) These metrics evaluate the set of generated sequences as a whole, while we directly select the sequences that meet the validity and binding score criterion, so they are not our primary objective.

(2) The comparisons are only meaningful within different generative models. For the random mutation-based methods, one can always achieve high novelty and diversity by introducing a sufficient number of mutations.

We show that the results are not very different between TCR-dWAE and TCR-dVAE (Table 13). In addition, we include the case of a VAE that suffers from posterior collapse (where the generation is always limited to the same few patterns). The scores for the first two are at a similar level as the validation set, while the diversity score for the collapsed V AE drops a lot. These suggest that the positive valid sequences presented in the paper are sufficiently diverse and novel.

MCPAS

Zf(true)

0

.0

zs(pred)

0

1

0

-10

Zs

（pred）

1

zs (true)

0

-10

z(pred)

0

10

-10

0

10

Z(pred)

5

0

z(tru

5

0

5

0

5

5

0

5

5

0

5

-5

0

5

Zf(true)

10

0

10

0

5

z（true)

z(pred)

zs (pred)

10

0

=10

Zs(true)

z(true)

zs(true)

(pred)

10

0

10

0

10

0

10

zf(tru

10

0

Zs(true)

zf(prec

10

0

zs(pred)

D

10

0

-5

0

10

0

0

10

0

5

z（pre

5

0

10

0

zs(true)

10

0

广

zs (pred)

10

0

10

0

0

-10

0

10

0

10

0

=5

10

0

zr (true)

10

0

-10

0

=

5

-10

0

Z(pred)

10

0

10

zs(true)

0

0

0

-10

0

10

Zs(pred

0

-5

5

0

5

-10

0

10

0

10

=10

0

10

-5

0

zf(true)

10

0

12

10

0

-10

5

0

5

Zf (pred)

10

0

10

-10

0

11

0

0

10

=10

0

10

10

0

10

-10

0

10

(A)

(B)

(A)

(B)

(D)

(C)

(E)

(F)

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes to disentangle the structural from functional representations when performing machine learning (ML) assisted protein engineering. The authors propose a Wasserstein autoencoder with auxiliary classifier to learn disentangled embeddings for the structural and functional components. Once disentangled embeddings are learned, they generate new sequences with several strategies: randomly selecting positive sequences, averaging positive sequences, and selecting the sequence with the highest classifier prediction. The method is empirically validated on TCR binder generation of modifying a TCR sequence to stay as a TCR while binding to a peptide of interest. The method is compared against a few baselines which their method is able to beat on their metrics.\n\nTheir contribution is learning disentangled structure and function representations to be utilized in TCR engineering.\n\nSTRENGTHS:\n- To the best of knowledge, this is the first exploration of wasserstein autoencoders for the task of protein engineering. The desire of separating functional and structural knowledge and solely optimizing functional properties of a protein is a fundamental challenge.\n\n- The disentanglement measurement is well-formulated and intuitively makes sense. The training objective is justified in this regard of maximizing disentanglement measurement.\n\nWEAKNESSES:\n- Implicit vs. explicit disentanglement. The authors should be careful in stating their method is the \"first approach that utilizes disentangled representations for protein engineering\" (line 15). Their method is a form of implicit disentanglement where the learned embeddings are separated to have a structural and functional component. Prior works such as SMCDiff [1], RFdiffusion [2], constrained hallucination [3] and Cao et al [4] perform explicit disentanglement where one part of the protein is kept fixed and the remaining sequence is optimized. Here the structure is assumed to be known a priori therefore the methods can directly keep fixed the part of the protein not necessary for function. A discussion around the comparison between these two views of the problem is necessary. Additionally they should have benchmarked against the motif-scaffolding methosd if the authors pick TCR for which the structure is very well known. i.e. [5, 6] could ( and should) have been benchmarks.\n\n- Only benchmarking on TCR. I do not find the experiments compelling to support the claim of separating the functional and structural components. TCR is well-studied and lots of data exists. We know very well there is a highly conserved region and only the CDR regions vary depending on the binding target. Therefore TCR seems like a easy task to separate the functional and structural components from the evolutionary data. Perhaps a toy example where the method can detect non-trivial disentanglement would be helpful or another protein system? I can also think of another baseline of only limiting the mutations to the CDR regions. How would this fare against the proposed method?\n\n- Claim of preserving structural backbone. The authors seem to conflate sequence vs. structural backbone conservation. They utlize a sequence classifier to classifier if a sequence is TCR or not. That is fine but this is not the same as keeping the structural backbone the same. What does it mean to stay as a TCR? This definition is not provided. How do we know the structural backbone is conserved based on the classifier predictions? The sequences are not ran through AlphaFold2 which could have supported their claim but we also know AF2 is not sensitive to point mutations. The notion of \"structural embedding\" (line 29) used throughout is misleading because it has no relationship to structure based on what is presented.\n\n- Lack of related work. I find it concering there is little mention of NetTCR throughout the work. As I understand it, NetTCR is a critical comparion  the function classifiers and oracles used in this work. In addition there is no discussion of machine learning methods applied to antibody engineering. The lack of this discussion makes the method appear more novel than it is.\n\n- Missing method and metric clarity. There are missing definitions of symbols (see questions) and some decisions are made arbitrarily. For instance, binding score is positive if $\\psi(z_f`, u) > 0.5$ where 0.5 is the operating point but why 0.5? There needs to be justification for this operating point. Otherwise the operating point could've been cherry picked to make the method look better. \n\n\n\n\n\n[1] https://arxiv.org/abs/2206.04119\n[2] https://www.biorxiv.org/content/10.1101/2022.12.09.519842v1 \n[3] https://www.science.org/doi/10.1126/science.abn2100 \n[4] https://www.nature.com/articles/s41586-022-04654-9 \n[5] https://arxiv.org/abs/2207.06616 \n[6] https://www.biorxiv.org/content/10.1101/2022.07.10.499510v1 \n[7] https://www.nature.com/articles/s42003-021-02610-3\n\nQUESTIONS:\n- As stated in weaknesses, why is there no direct comparison to NetTCR and other antibody engineering methods? I imagine the antibody engineering methods could readily be applied here?\n\n- In line 121 what is beta_1 and beta_2? What are they set to?\n\n- What is $\\Sigma$ on line 117? \n\n- Line 110-111. Why does making the embedding space matc N(0,1) lead to independence of z_f and z_s? For someone not familiar with Wasserstein autoencoders, this needs to be explained better. Additionally what is the meaning of \"gaussian shape of the latent space facilitates generation\" on the same line. I thought the method is deterministic so why is this mentioned? Is there sampling happening?\n\n- Line 160-151. What is meant by \"well-classified by ERGO\"? \n\n- Again I am confused why the authors went out to develop their own classifiers if ERGO and NetTCR are state-of-the-art. Did these methods have a limitation that is not mentioned?\n\n- Line 164-165. \"We generate an additional dataset with 5000 TCR sequences that do not overlap with the training set.\" How are the sequences generated and what is the definition of \"does not overlap\"?\n\n- Line 176-177. For generating positive z_f`, one of the ways is to use the best z_f that produces the highest classifier prediction. How is this generation done? Why wouldn't this lead to adversarial samples? For \"random\" I am unclear of the motivation of this technique. Surprisingly it does very well compared to the other methods. Why does that happen? \n\n- Line 185-186. Why is a autoencoder used as a oracle? It seems very odd to use a continuous latent space for a discrete generation task. Why is a transformer or LSTM not used?\n\n- Comparison to IDEL. Since IDEL is a VAE, shouldn't the authors generate multiple samples or do some sort of greedy search in its latent space for a fair comparison?\n\n- Line 207. It sates the majority of generated valid sequences are unique and novel. What is meant by unique and novel? Is there a sequence identity cutoff? I see no discussion of this in the appendix though the appendix is referenced.\n\nLIMITATIONS:\nThe authors only explain one limitation which is their experiments are performed on a smaller subset selected by ERGO. I don't believe there is sufficient analysis of the failure modes of their method. Why and when does the method fail? Are there examples theirs does good on where baselines fail? This would shed light on the limitations but it is missing.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 1, "DATE": "2023-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn many protein engineering problems, we want to maintain the overall fold of the protein while changing its functional properties. This is analogous to style transfer in computer vision. The paper uses a disentangled Wasserstein autoencoder to learn separate functional and structural embeddings for T-cell receptors (TCR). First, the authors show theoretically that their model should isolate peptide-binding information from information about the TCR scaffold structure. Next, they show that feeding their model new target peptide embeddings results in new TCR sequences that are predicted to be valid and to bind to the target peptide. Finally, they show empirically that the representations are indeed disentangled and demonstrate via ablation studies that all aspects of their model are required for strong performance.\n\nSTRENGTHS:\nThe paper describes an original, performant solution to a significant problem in protein engineering. I found it difficult to believe that similar methods have not been previously applied to protein engineering, but after spending some time searching, that does appear to be the case. In general, the exposition is clear. The experiments are thorough and provide strong evidence that their solution is better than baseline methods and that all components of their model are required for performance.\n\nWEAKNESSES:\n## Major\n\nThe title and introduction make general claims about \"protein engineering,\" but the remainder of the paper focuses on TCRs, and specifically on a short segment called the CDR3Beta region. The authors should either limit the claims to match their experiments, or add experiments to show that this framework will work for protein engineering in general. \n\nIn TCRs, binding is mostly mediated by the variable region within the CDR, and indeed the model learns to mostly change that region, as shown in Figure 2a and b. The results would be a lot more general if the model were applied to a case where there is not such a clean split between the functional region and the scaffold. This also leads to the question of whether the model needs to learn how to disentangle function and scaffold from data. Two possible experiments to validate this: \n1. Only model the hypervariable region and use a constant consensus sequence for the first 4 and last 3 residues\n2. Separately embed the hypervariable and conserved regions directly, instead of having the model learn them from the data. \n\nIn addition, there are other works that explicitly hold the functional part of a protein sequence or structure constant and generate a scaffold around it. See, for example, https://www.biorxiv.org/content/10.1101/2022.12.09.519842v1 and https://www.nature.com/articles/s43588-023-00440-3. This is the reverse of what this paper tries to do, so I don't think it's necessary to do direct comparisons, but it should be mentioned in the related work. \n\n## Minor\n\nTable 1 is difficult to interpret. As their primary empirical result is that no baseline matches their method on both validity and binding, it could be clearer to plot generated sequences by validity and binding colored by method to show that the sequences from their method lie at the pareto front. \n\nIt's unclear to me how much the theorem actually strengthens the paper, or how it's used in the architecture or results. As the authors state right after the theorem, the independent Gaussian prior and classification loss already promote disentanglement. \n\nIt's a little confusing that the \"structural embedding\" is never used to predict a structure, only the CDR sequence. \n\nIn Figure 3c, it would be nice to label the embedding components directly on the figure in addition to describing them in the legend.\n\nQUESTIONS:\nHow well would this method work on more general, larger, proteins where there's a less obvious decomposition between scaffold and function in sequence space? \n\nAlternately, for TCRs, how much better is this method than explicitly decomposing the scaffold and function in sequence space?\n\nWhy use an LSTM decoder? There's already a transformer encoder -- wouldn't it be easier to also use a transformer decoder? \n\nIn general, very little information is given about model and training hyperparameters, including model size, hardware used, and length of training.\n\nLIMITATIONS:\nAs is, while promising, the paper is limited by the fact that empirical results are only presented on a short segment of TCR hypervariable regions. This is not a typical protein engineering case: it only addresses one type of function (binding to peptides), in a non-typical protein, and glosses over the fact that they do not even engineer the entire protein. The method is still interesting, and I am hoping to increase my score (a little) if the authors clearly limit their claims and state limitations, or (a lot) if they demonstrate that this works more generally.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this work, the authors develop an approach to T-cell receptor (TCR) generation using disentangled representation learning (DRL). They introduce a disentangled Wasserstein autoencoder with an auxiliary classifier, and show through a series of experiments that they can separate the “structural” and “functional” representations of TCR sequences. Sample validity and quality are evaluated using pre-trained autoencoders and binding classifiers. Overall, the paper addresses an interesting and challenging problem area (TCR engineering) and DRL is a promising avenue for protein engineering more broadly. However, some questions about the evaluation procedure lead to difficulties in interpreting the results.\n\nSTRENGTHS:\nThe paper considers an important application area, T-cell receptor engineering, which is understudied in the ML community, and is an original application for discrete sequence modeling. The overall methodology and presentation is clear.\n\nWEAKNESSES:\nBecause the authors introduce TCR validity as one of their main evaluation criteria, it requires substantial evidence that it is a meaningful optimization target. Some of the details related to evaluation and limitations are not clear.\n\nQUESTIONS:\nAre there examples in the protein engineering literature for using a pre-trained autoencoder to calculate validity? Is the VDJDB used for training TCR-dWAE a subset of TCRdb? Have the authors considered a simpler metric like “naturalness” (pseudo log-likelihood) from a masked language model or log-likelihoods from an autoregressive language model trained on TCRdb? In the discussion of TCR validity, what does “has a similar embedding pattern as the known TCRs” mean?\n\nThe authors claim that the structural embedding constrains the sequence backbone to ensure validity, but reported %validity is < 60% for TCR-dWAE. The % positive valid is also very low, despite the low number of mutations. As the authors mention, a low number of mutations is desirable for real protein engineering tasks to improve function, but it may be unfeasible to introduce function (peptide-specific binding) with a single mutation.\n\nIn Table 2, is it the case that the average number of mutations is about 0.5 for each method? How is this reconciled with the reported uniqueness and validity? \n\nThe authors limit modeling to the CDR3beta region, which is reasonable for modeling function, but weakens any demonstrations of “disentangled” representations. There are highly conserved motifs at the beginning and end of CDR3beta that the generated samples correctly recover. It is difficult to evaluate claims of structure vs function disentanglement in the high-entropy region of CDR3beta. This argument may be better supported by showing guided sampling that interpolates between two known modes of TCR binders. \n\nHow are Miyazawa-Jernigan energies calculated?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes to use disentangled representations for protein engineering which enables preserving the generic structure while making functional modifications to the protein of interest. To disentangle functional patterns from the rest, they propose to use disentangled Wasserstein auto encoders with auxiliary classifier on the embeddings that encode function. The application of their method has been studied for T-Cell receptor engineering. It is demonstrated that given negative pairs of peptide (TCR target) and TCR, their method can alter TCRs into new sequences with higher validity and binding score relative to other approaches.\n\nSTRENGTHS:\nThe problem is well motivated, and the paper does a good job in adapting the approaches developed in other domains and applying it to protein engineering.\n\nWEAKNESSES:\nThe proposed method is not novel but its application to a new domain is of interest.\n\nQUESTIONS:\nQ1. The results are averaged over four selected peptides (AUROC > 0.8) in Table 1. In the experiments section, it is stated that 10 peptides were used. \n\n1.1.\tCan you provide the results of Table 1 for the rest of the peptides? \n\n1.2.\tAlso, as only 10 peptides were used, can you provide the statistics in Table 1, e.g., r_v\nand r_b, for each peptide?\n\n1.3.\tGiven that you engineered 5000 TCRs per peptide, I would like to see the correlation between the estimated validity and the binding score for each peptide. \n\nQ2. Theoretically you have shown that the two latent embeddings of structure and function are encouraged to be disentangled. However, in the experiments do you have a way of measuring structural similarity between the altered TCR and the input TCR that does not bind to the peptide? I understand that the structural embedding is intact in the generation, but can you quantify it in the sequence space? I do not think that the validity score tells us whether the output folds like the input. Please clarify this. \n\nTo explain my point better, are all TCRs roughly similar in structure? Or their structural properties, e.g., folding, varies a lot. If the former is true, then maybe validity score is a good criterion but if the latter holds then you should quantify the similarity of the generated sequence with improved binding to the input sequence with poor binding.\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}