Disentangled Wasserstein Autoencoder for T-Cell Receptor Engineering
Authors: Tianxiao Li
Hongyu Guo
Filippo Grazioli
Mark Gerstein
Martin Renqiang Min
Venue: NeurIPS 2023 poster (2023)
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
Tianxiao Li Hongyu Guo Filippo Grazioli', .Mrk Gerstein 2r, Martin Renqiang Min It+
NEC Laboratories America, Yale University
National Research Council Canada, NEC Laboratories Europe
tianxiao. li@yale edu, hongyu. guo@uottawa ca, flpgrz@outlook. com
mark. sorkkggesseiinyale.. edu, renqiangênec caaeecclabb...oo com
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 critical for protein engineering computationally non- trivial, and requires significant domain knowledge. To automate this process from a data-driven perspective, we propose 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 fructure-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.
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 enzymes [T, and the characteristic immunoglobulin fold versus the antigen-binding aumplementarityddeemmiin
region (CDR) in immunoproteins 2 3]. The latter usually defines the protein't 's key function, but cannot work on its own without the stabilizing effect of the former. This dichotomy is similar to the content- style 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 77] 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 "functional" embedding and "structural" embedding. This approach results in 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 g or static and dynamic parts of videos 10 for tasks such style transfer [11 8 and conditional generation 22 13
"Equal contribution. 'Corresponding author.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Attaining the aforementioned disentangled embeddings in discrete sequences such as protein se- quences, 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] [5] Cheng, et al. [6 improves the disentanglement with a mutual information (MI) upper bound on the embedding space of 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 Wasserstein autoencoder (WAE) [17] framework that achieves disentangled embeddings with theoretical guarantee, using a simpler loss function. Also, WAE could be trained deterministically. avoiding several practical challenges of VAE in general, especially on sequences [T8 [9]. 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 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
[1 (Fig. [T]. Designing TCRs with higher affinity to the target peptide is thus of high interest in immunotherapy [22]. Various driven methods have been proposed to enhance the accuracy of TCR binding prediction 25 12 [3 26 (2). However, there has been limited research on leveraging machine learning for TCR engineering. A related application is the motif scaffolding problem [2 where structural "scaffold" is generated supporting 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 CDR33 region of TCRs where there is sufficient data and a clearly defined functional role (peptide binding). Using large TCR- 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, assume that similarity between the "structure-related" parts of the sequence indicates similarity in the structure. Furthermore, by modifying only the functional embedding, 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".
approach is able to generate
We summarize main contributions follows:
To our knowledge, we are the first to formulate computational protein design as style transfer problem and leverage disentangled embeddings for TCR engineering, resulting in more interpretable and efficient conditional generation and property manipulation.
We propose disentangled Wasserstein autoencoder with auxiliary classifier, which effectively isolates the cenction-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 could not bind to. introduce 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 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

---

(A)
(B)
Viral infection or tumor
Input sequence
82
Reconstructed sequence
MHC
Gaussian
peptide
TCR
1 Interaction
(C)
Template sequence
6,
Modified sequence
peptide
Modified g'y
TCRa
TCRB
Figure 1: (A) Top: The TCR recognizes antigenic peptides provided by the major histocompatibility complex (MHC) with high specificity; bottom: the 3D structure of the TCR-peptide-Me binding interface (PDB: 5HHO); the CDRs are highlighted. (B) The disentangled autoencoder framework, where the input x, i.e., the CDR33, is embedded into functional embedding 2 (orange bar) and structural embedding 2, (green bar). (C) Method for sequence engineering with input x. Zg of the template sequence and a modified z',, which represents the desired peptide binding property, are fed to the decoder to generate engineered TCRs
that define the peptide binding property. TCRs that bind to the same peptide should have similar functional patterns. The latter refers all other patterns that do not relate to the function but could affect the validity. Following 223 we here limit our modeling to the CDR33 region since it the most active region for TCR binding. In the rest the paper, we use "TCR" and "CDR33" interchangeably.
2.2 Disentangled Wasserstein Autoencoder
Our proposed framework, named TCR-dWAE, leverages disentangled Wasserstein autoencoder that learns embeddings corresponding to the functional and structural patterns. In this setting, the input data sample for the model is triplet {x, u. 3}. where x the TCR sequence, the peptide sequence, and y is the binary label indicating the interaction.
In detail, given input triplet (x, u, y), the embedding space of is separated into two parts: = concat(z; ).), where Z is the functional embedding, and z, is the structural embedding. The schematic of the model is shown in Fig. BB.
2.2.1 Encoders and Auxiliary Classifier
We use two separate encoders for the embeddings, respectively:
(x))
where e (s, / correspond to "structure" and "function".
First, the functional embedding y is encoded by the functional encoder (x). In order to make sure z carries information about binding the given peptide u. we introduce auxiliary classifier V (z u) that takes z and the peptide u as input and predicts the probability of positive binding label qo (y Zy, u).
ÿ = (Y = 1 zy, u) = ( u).
We define the binding prediction loss as binary cross entropy:
Cr yog 9 (1 y) log( 1 9)
3

---

Second, the structural embedding is encoded by the structural encoder ④ (x). To enforce to contain all the information other than the peptide binding-reiie related patterns, we leverage a sequence reconstruction loss, as will be discussed in detail in Section 2.2.3
2.2.2 Disentanglement the Embeddings
To attain disentanglement between 2f and 2x. 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 Qz where = concat(zj- 2ss) and an isotropic multivariate Gaussian prior Zo Pz where Pz = (0, Ta):
L Z MMD Pz, Qz).
The MMD is estimated as follows: given the embeddings {Z1, 22, .... Zn of an input batch of size we randomly sample from the Gaussian prior {21. ž2, ..., Zn with the same sample size. We then use the linear time unbiased estimator proposed by B3 to estimate the MMD:
[n/2]
MMD Qz)
Σ
(((22-1, 22i 1), (Z2i, k2;));
11/2
where h((2,, 2,); (z.. 2j)) kz, + (Z, 3,) ((, 2z) -
k(zz 2.) and & is the kernel function. Here we use a radial basis function (RBF) 34 with o = as the kernel.
By minimizing this loss, the joint distribution of the embeddings matches ( Ia), so that 2 and Z, are independent. Also, the Gaussian shape of the latent space facilitates generation [3
2.2.3 Decoder and Overall Training Objective
The decoder takes zy, z, and peptide u as input and reconstructs the original sequence as x'. It also acts as regularizer to enforce the structural embedding 2, to contain 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 the sequence:
x' l'(concat(z 2z. Zy, u))
1 Σ
Lrecon (x, x')
(1) log(x"(0) (1 (1 x loglll 1og(1=x(())
where is the length of the sequence and (t) is the probability distribution over the amino acids at the i-th position.
Combining all these losses with weights 31 and 32, we have the final objective function, which then can be optimized through gradient descent in end-to-end fashion:
= Crecon + B1 B + B2L
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 We introduce measurement of disentanglement as follows:
16
D Z.; X U)- = -VI(Z,; X U) + Vi UJ - VZ UU,
where VI is the variation of information, VI(X; H(X) H(Y) - 2I; which is measure independence between two random variables. For simplicity, we omit the condition U (peptide) in the following parts.
This measurement reaches 0 when 2 and Z, are totally independent, i.e. disentangled. It could further be simplified as:
XX) + + I(X; T(X; 2,)].

---

Note that H (X) is constant. Also, according to data processing inequality, as 2 - - forms Markov chain, we have [ (z:z,) 2 (y: 2y). Combining the results above, we have the upper bound of the disentanglement objective:
[ Z.) (x;; ((; 2,) / ZZ, ((;; Zy)
Next, we show how our framework could minimize each part of the upper bound in 2
Maximizing ((x Similar to 10 we have the following theorem:
Theorem 2.1 Given the encoder Qa 2
X), decoder P,((X 2), prior P(Z), and the data
distribution PD
(Q(Z) II P(Z)) = Epp Drr X) -(X;
where QZZ) is the marginal distribution of the encoder when X PD and Z Qo (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 Qo Z X) and the prior PZZ). Detailed proof of the theorem can be found the Appendix [.I]] This also applies to the two separate parts of Z, 2 and Z.. In practice, because the marginal cannot be measured directly, we minimize the aforementioned kernel
MMD instead.
As a result, there is no need for additional constraints on the information content of Z. because [[X; Zy) automatically maximized by the objective. We also empirically verify in Section 3.2.3 that supervision on Z, does not improve the model performance.
Maximizing Zy) I(Y; Zy) has an lower bound as follows:
T(Y: 2y) 2 H(Y) + ())(), log 90 (Y Zz);
where qu (Y Z is the predicted probability by the auxiliary classifier (note we omitted the condition U here). Thus, maximizing the performance of classifier v would maximize I(Y; Zy).
Minimizing I (Z;;; Z.) Minimization of the Wasserstein loss forces the distribution of the em- bedding space Z to approach an isotropic multivariate Gaussian prior Pz -N Ia), 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 and Z..
As conclusion, our objective can can jointly minimize (Zy; Z.) and maximize 1 [(X; Z.) and [(Y Z₁). The former ensures independence between the embeddings, and the latter enforces them to learn separate information, achieving disentanglement.
Experiments
3.1 Setup
Datasets Interacting TCR-peptide pairs are obtained from VDJDB (36,. merged with experimentally- validated negative pairs from NetTCR [26). We then expended negatives to 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- ECR-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.I). D.. For each peptide, select from VDJDB an additional set of 5000 TCR sequences that do not overlap with the training set. This set used for the TCR engineering experiment.
Implementation Details The TCR-dWAE model uses two transformer encoders [38 for O., θ and long short-term memory (LSTM) recurrent neural network decoder 39 for T. The auxiliary classifier 2- layer perceptron. Hyperparameters are selected through grid search. All results are averaged across five random seeds. See Appendix B.2 for more details.
5

---

%vlid
#mut/len
&positive valid
TCR-dWAE
1.36+0.03
0.45=0.09
0.61
0.49=0.05
0.23+0.02
TCR-dVAE
1.44+0.02
0.24=0.02
0.64+0.05
0.4+0.02
0.16#0.01
greedy
0.34+0.0
0.79=0.0
0.02+0.0
0.34+0.0
0.02+0.0
genetic
0.3940.03
1.0+0.0
0.02+0.0
NA
0.02+0.0
naive rm
0.35+0.0
0.2+0.0
0.03+0.0
0.35=0.01
0.0=0.0
MCTS TCR-dWAE (null)
-0.1+0.0
0.95=0.0
0.0+0.0
NA
0.0=0.0
1.51
0.05=0.03
0.85+0.01
0.45=0.07
0.04+0.03
original
1.59
0.01
0.92
NA
0.01
Table 1: Performance comparison. original is the evaluation metrics the template TCRs without any modification. Only unique sequences are counted. Results are averaged across selected peptides (i.e., SSYRRPVGI, TTPESANL, FROYVVREEKTTRAEOOSSSS CTPYDINQM) and then averaged across five random seeds.
3.2 TCR Engineering
3.2.1
Manipulating TCR Binding via Functional Embeddings
As shown in Fig. [CC given any TCR sequence template x. we combine its original Z, and new functional embedding x', that positive for the target peptide u which then fed to the decoder to generate a new TCR that could potentially bind to u while maintaining the backbone of x. For each generation, obtain the z'y from random positive sample in the dataset.
3.2.2 Metrics and Baselines
We use the following metrics 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, 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 the distribution of the latent space, we consider it as 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 ERGO classifier and binding probability ты = ERGO(x', u) is calculated.
fed into the pre-trained
We compare TCR-dWAE with the following baselines (see Appendix
C for more details):
Random mutation- based For each iteration, new random mutation is added the template.
The best mutated sequence (one with the highest ERGO score) is selected on either a sample-level (greedy), population-les (genetic), or without any selection (naive rm). then goes into the next round.
Generation- This includes Monte Carlo tree search (MCTS), where random residues are itera- tively added until the sequence reaches a maximum length. For each generation process, the leaf with the highest ERGO score added to the results.
TCR-dVAE We also include a VAE baseline inspired by the objectives of IDEL from [16), which minimizes MI upper bound to achieve disentanglement in the embedding space (Appendix
3.2.3 Main Results
We use peptides with AUROC > 0.8 by the classifier V for the engineering task (SSYRRPVGI. CTPYDINQM), The results are listed in Table The cutoff for "valid" is r, 2 1.25, as it can successfully separate known TCRs from other protein sequences. The cutoff for "positive" TCR rs 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%
TTPESANL, TRDYYDREVKKTRAAQQASS

---

(A)
BB)
Average attention (functional encoder)
Generated valid sequence
CIPYDINQM
CASS
QYF
TTPESANL
SSYRRPVGI
0.05
FADIWDDRFYKTLBAEDASS
Generated positive sequence
Average attention (structura encoder)
CASS
QYE
CTPYDINQM
0.10
SSYRRPVGI
Known sequence
FRDIVDBFFKTLBBECASEE
CASS
Oz
Mutation frequency
CTPYDINQM
TTPESANL
SSYRRPVGI
FROYVORFVKTUAAECOSSEE
(C)
Generated
Template
Generated pos
Template
CASSPCEST EEENNDTTY
0.6
0.6
0.4
0.0
0.0
0.2
0.2
-0.4
-0.6
-0.6
SSYRRPVG
SSYRAPVG1
CTPYD INQ
Figure 2: (A) Consensus motif of the first and last 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 the functional encoder for the positive TCRs; (middle) average attention patterns of the structural encoder for the positive (botttoottt))) 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-base methods generate valid and positive sequences compared to other methods. We also include the results on McPAS-TCR in the Appendix Table where the observations are similar.
One advantage of TCR-dWAE is that 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-distribttion 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 enginecring in pass, requiring 10x less time (Appendix Table IO).. Compared to VAB-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 we've found be rather sensitive hyperparameters (see Appendix C.3 and E.3 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 similar conservative region motif as known TCR CDR3 is 41 (Fig. A). The average attention score of the functional encoder concentrates on the hypervariable part (Fig. 2B, top), which has similar distribution as the mutation frequency of the engineered sequences compared to their templates (Fig, 2B, bottom). On conservative regions
the other hand, we can also observe some, though less frequent, changes through modifying the 2y, while some residues in the hypervariable region maintained (Fig. (C)) Furthermore, the attention patterns of the functional and structural encoders, despite having their preferences, overlap at some positions and do not have clear-cut separation. Neither pattern
7

---

&positive valid
z 1
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 029
0.008+0.000
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 (2,, z.) Random z
0.9044 = 0.0020
0.7535 0.0085
Fully random
0.6016 0.0030
Table 3:
Reconstruction accuracy with altered embeddings (VDJDB).
aligns with the hypervariable/cossrvvtvv separation (Fig. 23 top and middle). These results indicate that both 2f and learn to encode patterns from the entire sequence, and do not fit into the manual separation of "functional" hypervariable region and "structural" conservative region.
We also show some examples where residues introduced by the positive y to the engineered sequence have the potential of forming lower-energy (i.e. more stable) interactions with the target peptide (Fig. (C), using the classic Miyazawa-Jernigan contact energy 442 which a knowledge-based pairwise energy matrix for amino acid interactions. These results demonstrate that the positive patterns carried by z have learned some biologically relevant information such as the favored binding energy of the TCR-pMHC complex, and are successfully introduced the engineered TCRs.
3.3 Ablation Study
For quantitative comparison of the disentanglement, we calculate sample-based MMD between the embeddings of positive and negative samples. Ideally, a disentangled latent space would make zy of the positive and negative groups very different, and z, indistinguishable. As shown in Table 2 removing the Wasserstein loss (Wass-) or the auxiliary classifier (w-- leads to both poorer performance and poorer disentanglement.
We also attempt to explicitly control the information content of z, by adding another decoder using only 2, (L. This in practice results in similar, slightly better, disentanglement compared to the original model, but not necessarily better performance. These observations agree with our conclusions in Section 2
3.4 Analysis of the Embedding Space
We select balanced subset from the test set (332 positives and 332 negatives) and obtain their z and n, embeddings. T-SNE 43
visualization Fig. 3 shows that for each peptide, the binding and non-binding TCRs can be well separated by 2, but not z, (more examples in Appendix Fig. 5)
Also, among the positive samples, 2, show strong clustering patterns corresponding to their peptide
targets (Fig. 2B) As result of the Wasserstein loss constraint, there is minimal correlation between z and 25 (Fig. 3C).
Furthermore, as shown in Table the reconstruction accuracy is the highest when both (z,, z.) are provided, and significantly impaired when either embedding missing. These results imply that the embedding space is properly disentangled as intended: zy and 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)
SSYRRPVGI
(B)
2, of positive sequences
(true)
(true)
SSYRRPVSI
10
5
COM
-5
TTPESANL
Zy (true)
(true)
-10
-15
-20
20
CTPYDINQM
TTPESANL
-10
FRDYVOR
FRDYWDRFYKTRRAEQASO
2r (true)
2, (true)
(C)
-5
10
-10
CTPYDINQM
(true)
(true)
0-
-5 11
Figure 3: (A) T-SNE patterns of 2 and z, for selected peptides, colored by ground-trutt labels. (B) t-SNE of the / of the positive TCRs. Colors correspond to their binding peptide. between dimensions of z and 2s, where the orange color corresponds to zy and green to Zg.
(C) Correlation
("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 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. 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 "Cp-peptidee 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 function versus structure, believe it can be further extended broader definition of protein functions or other molecular contexts.
Related Works
TCR-peptide binding prediction Most human TCRs are heterodimers comprised of TCR& chain and a TCR3 chain [21] Each chain contains three complementarity- -eteraiiiinddeeriiinnnngg regions (CDRs), CDR1, 2 and 3, which are highly variable due to gene recombination. A large amount 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 Other methods assess the binding energy for the TCR- -eptide pair using physical structural In recent years, several deep neural network architectures have been applied to
46 24
modeling 47
48
this task 49
26
50
51 27], using both TCR and peptide sequences as input. Most studies 23
24

---

focus on the CDR33 sequence and some also include other information like the CDR3a, depending on the availability of the data.
Computational biological sequence design Traditional methods of computational biological se- quence design uses search algorithms, where random mutations are introduced and then selected by specified criteria [52 Deep generative models like variational autoencoder (VAE) [55] or generative adversarial network (GAN) [6] have also been successfully applied biological sequences such as DNA, Other recent approaches include reinforcement learning (RL) 17]. and iterative refine- ment [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 its nature, such as separating the "style" and the "content" of a piece of text [16] or picture [8], or static (time-invariant) and dynamic (time-variant) elements of video [ 10] The disentangled representations can be used for style transfer [11] 8 [6] and conditional generation
Several regularization methods have been used achieve disentanglement with without a priori knowledge, including controls on the latent space capacity [57 58) adversarial loss [12), cyclic reconstruction [[]]. mutual information
12
constraints (59 and self-supervising auxiliary tasks
References
[1] Peter K Robinson. Enzymes: principles and biotechnological applications. Essays 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, Mingli Song. Neural style transfer: A review. IEEE transactions on visualization and computer graphics, 26(1 1):3365 3385, 2019.
[5] Martina Toshevska 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, (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.
Yizhe Zhu, Martin Renqiang Min, Asim Kadav, Hans Peter Graf. S3vae: Self sepe-suue 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 Rengiang Min, Ligong Han, Li Erran Li, and Xuan Zhang. Disentangled recurrent wasserstein autocncoder. In International Conference on Learning Representations, 2021.
[11] Hsin- 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 I. Denton et al. Unsupervised learning of disentangled representations from video. Advances in neural information processing xystems, 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 Lan Boureau. attribute text rewriting. In International Conference on Learning Representations. 2019.
10

---

[15] Vineet John, Lili Mou, Hareesh Bahuleyan, and Olga Vechtomova. Disentangled representation learning for non-parallel text style transfer. arXiv preprint 80X1v:1808.04339, 2018.
[16] Pengyu Cheng, Martin Renqiang Min, Dinghan Shen, Christopher Malon, Yizhe Zhang, Yitong Li, and Lawrence Carin. Improving disentangled text representation leanenn with information-theorriicc guidance. arXiv preprint arXiv:2006,00993 2020.
[17] Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein auto-encoders. arXiv preprint arXiv:/71..01558, 2017.
[18] Samuel R Bowman, Luke Vilnis, Oriol Vinyals, Andrew M Dai, Rafal Jozefowicz, and Samy Bengio. Generating sentences from a continuous space. arXiv preprint 1rXiv:1511.06449 2015.
[19] Partha Ghosh, Mehdi SM Sajjadi, Antonio Vergari, Michael Black, and Bernhard Schölkopf. From variational to deterministic autoencoders. arXiv preprint 19t.:::990.124330, 2019.
[20] Balachandra K Gorentla and Xiao-Ping Zhong. T cell receptor signal transduction in lymphocytes. Journal of clinical & cellular immunoloyy, 2012(Suppl 12):005, 2012.
[21] Kinjal Shah, Amr Al-Haidari, Jianmin Sun, and Julhash U Kazi. T cell receptor (tcr) signaling in health and discase. Signal transduction and torgeted therapy, 6(1):1-26, 2021.
[22] Nicholas P Restifo, Mark E Dudley, and Steven A Rosenberg. Adoptive immunotherapy for cancer: harnessing the t cell response. Nature Reviews Immumology, 12(4):269-2) 2816999-2111110 2012.
[23] Ido Springer, Hanan Besser, Nili Tickotsky- Moskovitz, Shirit Dvorkin, and Yoram Louzoun. Prediction of specific ter- peptide binding from large dictionaries of ter- tcr-peptide pairs. Frontiers in immunology, page 1803, 2020.
[24] Anna Weber. Jannis Born, and María Rodriguez Martínez. Titan: T-cell receptor specificity prediction with bimodal attention networks. Bioinformatics, 37(Supplement__))2233
2021.
[25] 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,22 2023.
[26] Alessandro Montemurro, Viktoria Schuster, Helle Rus Povlsen, Amalie Kai Bentzen, Vanessa Jurtz. William D Chronister, Austin Crinklaw, Sine R Hadrup, Ole Winther, Bjoem Peters, et al. Netter-2.0 enables accurate prediction of tcr-peptide binding by using paired tera and 3 sequence data. Communications biology, 101):1-13, 2021.
[27] Filippo Grazioli, Pierre Machart, Anja Mösch, Kai Li, Leonardo V Castorina, Nico Pfeifer, and Mar- tin Renqiang Min. Attentive variational information bottleneck for ter cer-peptide interaction prediction. Bioinformatics. 39baac2,,, 2023.
[28] Erik Procko, Geoffrey Y Berguig, Betty W Shen, Yifan Song. Shani Frayo, Anthony J Convertine, Daciana Margineantu, Garrett Booth, Bruno Correia, Yuanhua Cheng, al. A computationally designed inhibitor of an epstein-barr viral bel-2 protein induces apoptosis in infected cells. Cell, 157(7):: 16641166666 2014.
[29] Jue Wang, Sidney Lisanza, David Juergens, Doug Tischer, Joseph I. Watson, Karla M Castro, Robert Ragotte, Amijai Saragovi, Lukas F Milles, Minkyung Back, et al. Scaffolding protein functional sites using deep learning. Science, 37766644)387-399, 2022.
[30] 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-seaffolding problem. arXiv preprint arXiv:2206.04119 2022.
2022.
[31] Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian I. Trippe, Jason Yim, Helen E Fisenach, Woody Ahem, 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, 12, 2022.
[32] 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.
[33] Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf. and Alexander Smola. A kernel sample test. The Journal of Machine Learning Research, 13(1):723-75 773, 2012.
11

---

[34] 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 nearal information processing systems, 25, 2012.
[35] Diederik Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint 1r1ii:13122.111
2013.
[36] Dmitry Bagaev, Renske MA Vroomans, Jerome Samir. Ulrik Stervbo, Cristina Rius, Garry Dolton, Alexander Greenshields- Watson, Meriem Attaf, Evgeny Egorov, Ivan Zvyagin, et al. Vdjdb in 2019; database extension, new analysis infrastructure and a t-cell receptor motif compendium. Nucleie Acids Research, (3(DI))D1057-D1060 2020.
[37] Nili Tickotsky, Tal Sugiv, Jaime Prilusky, Fric Shifrut, and Nir Friedman. Mepas-ter; a manually curated catalogue of aasooicatd cell receptor sequences. Bioinformatics, 33(18):2924-25 2929, 2017.
[38] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances neural information processing systems, 30, 2017.
[39] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 983::1355 1997.
[40] Si-Yi Chen, Tao Yue, Qian Lei, and An- Guo. Terdb: a comprehensive database for t-cell receptor sequences with powerful search function. Nucleic Acids Research, 49(D1)))4688-447)) 2021.
[41] Douglas Freeman, René L Warren, John R Webb, Brad H Nelson, and Robert A Holt, Profiling the I-cell receptor beta-chain repertoire by massively paraliel sequencing. Genome research, 19(10):1117-1824 2009.
[42] Sanzo Miyazawa and Robert I. Jernigan. Estimation of effective interresidue contact energies from protein crystal structures: quasi-chemical approximation. Macromolecules, 18(3):334-552, 1985.
[43] Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9)1 2008
[44] Pradyot Dash, Andrew GarreaGa Tomer Hertz, George C. Wang, Shalini Sharma, Aisha Souquette, Jeremy Chase Crawford, Bridic Clemens, Thi HO Nguyen, Katherine Kedzierska, et al. Quantifiable predictive features define epitope-specific call receptor repertoires. Nature, 547(7661))88) 9989,,3, 2017.
[45] Yao Tong. Jiayin Wang. Tian Zheng, Xuanping Zhang. Xiao Xiao, Xiaoyan Zhu, Xin Lai, and Xiang Liu. Sete: Sequence- based ensemble learning approach for ter epitope binding prediction. Computational Biology and Chemistry, 87:107281, 2020.
[46] Nicolas De Neuter, Wout Bittremieux, Charlic Beirnaert, Bart Cuypers, Aida Mrzic, Pieter Moris, Arvid Suls, Viggo Van Tendeloo, Benson Ogunjimi, Kris Laukens, et al. On the feasibility of mining ed8+ cell receptor patterns underlying immunogenic peptide recognition. Immunogeneticx, 1063):15--16,, 2018.
[47] Kamilla Kjærgaard Jensen, Vasileios Rantos, Emma Christine Jappe. Tobias Hegelund Olsen, Mar- tin Closter Jespersen, Vanessa Jurtz, Leon Eyrich Jessen, Esteban Lanzarotti, Swapnil Mahajan, Bjoemn Peters, et al. Terpmhemodels: Structural modelling of ter-pmhe class complexes. Scientific reports, 9(1): 1-12, 2019.
[48] 1-Hsin Liu, Yu-Shu Lo, and Jinn-Moon Yang. Genome-wide structural modelling of ter-pmhe interactions. BMC genomics, 14(5):1-13 13, 2013.
[49] David S Fischer, Yihan Wa, Benjamin Schubert, and Fabian J Theis. Predicting antigen specificity of single cells based on edr 3 regions. Molecular systems biology, 16(8):94416 2020.
[50] Pieter Moris, Joey De Pauw, Anna Postovskaya, Sofie Gielis, Nicolas De Neuter, Wout Bittremiexx Benson Ogunjimi, Kris Laukens, and Pieter Meysman. Current challenges for unseen-epitope ter interaction prediction and a new perspective derived from image classification. Briefings in Bioinformatics, 22(4): b8aa318, 2021.
[51] 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.
[52] Frances H Amold. Design by directed evolution. Accounts of chemical research, 31(3): 1111 1998.
12

---

[53] Alexander Hartemink, David K Gifford, and Julia Khodor. Automated constraint-based nucleotide sequence selection for dna computation. Biosystems, 5221-33:2 1999.
[54] Yon Shin, In-Hee Lee, Dongmin Kim, and Byoung-Tai Taoung-Takk Zhang. Multiohjective evolutionary opti- mization of dna sequences for reliable dna computing. IEEE transactions evolutionary computation, (2):143-155) 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 1rXiv:1712.06145: 2017.
[56] Wengong Jin, Jeremy Wohlwend. Regina Barzilay, and Tommi Jaakkola. Iterative refinement graph neural network for antibody sequence-structure co-design. arXiv preprint 10.1v:2110.06624, 2021.
[57] Trina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts constrained variational framework. In International conference on learning representations, 2017.
[58] Christopher Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner. Understanding disentangling in 9-vae. arXiv preprint (rXXiv:/0440335000 2018.
[59] Chen, 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-100 1992.
[61] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXivv: 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))D400-D48, 2021.
[64] Moksh Jain, Emmanuel Bengio, Alex Bermandez-Garcit 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 Muchine Learning, pages 9786-9801. PMLR, 2022.
13

---

APPENDIX
A Data preparation
A.1 Combination of data sources
TCR-peptide interaction data are obtained from VDIDB and MCPAS ki 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 removed.
Only peptides with > 10
A.2 Filtering by ERGO performance
trains two separate models for VDJDB and MCPAS, the following filtering process is also performed separately the two datasets. For this and all subsequent ERGO-based predictions, pre-trained weights from https: bestpe://github.. httpchuuu.oommoooooo wwom/ooooonnalll wwwooounnllbbaaeee ab/ERGOO
Since ERGO (3
the
Additional negative samples are generated as follows: 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 positive sample counts (Table 5)
To the specificity of TCR recognition the following study, we did second round of filtering both the TCRs and the peptides. We pair all TCRs with at least one positive binding event and peptides in the filtered dataset. Any unobserved pair is treated as negative. This dataset is then provided to ERGO. Performance is shown in Table We discard peptides with AUPR < 0.7 and TCRs that have more than one positive prediction have least wrong prediction.
After that, downsample all peptides most 400 positive TCRs. This number is chosen that the resulting dataset more balanced across peptides. The final number samples for each peptide can be found in Table information, for every TCR in the positive set, we add its
Φ To make sure the model captures peptide-specife unobserved pairings with other peptides to the negative set. We then split the TCRs into train/testvallidation sets with a ratio of 8:1:1, and all pairings of each TCR to the respective subset, to ensure all TCRs in the test and validation sets not negative/positive ratio of original dataset.
the training. For the training set, the positive samples up-sampled by the
B
Model details
Notation
Meaning
functional encoder
O
structural encoder decoder
V
auxiliary functional classifier a data point with TCR x. peptide u and binding label
{x, u, 5}
z
functional embedding
Z
structural embedding
concatenation of (z,, 2,)
x'
reconstructed/gennrated sequence from the decoder
(()
the probability distribution over amino acids at the i-th position in x
concat(X1, .... x.,
concatenation of vectors {x1, .... Xn)
Table 4: Notations used for this paper. Sequences are represented as / x [ matrices where / is the
length [V] the number of amino acids.
B.1 Proof of Theorem 1
We use density functions for simplicity. Let (z x) be encoder and P,(x [x be decoder. joint generative distribution:
p(x, (x 12)(22).
14

---

where p(z) is the prior. Also, we have the joint inference distribution:
(xx z) = go xx)
where (x( is the data distribution.
g(x, x) (xxxx log
I(x: 2)
pD(x)q(()
Epp Σ pD px)x00(( x) loy go(z x))p(x)
pr(x)1(()
pp Σ 20 x) log go(z x)
((z))
Epp Σ 90 x) log
20(z x)
Epp Σ 90 x)
log
g(z)
p(z)
p(z)
Epp Σ 20 x) log
20(z x)
Σ ((z) log
g(z)
p(z)
p(z)
=Ep [Tx (Qa X) P(Z))] TKK. (Q(Z) k
P(Z)).
Thus,
|| P())) (Q(Z)
Epp DKr (Q&(Z) X) P(Z)) I(;;))
B.2 Implementation and training details
All imput 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 an embedding layer that transforms the one-hot encoded sequence x into continuous-valus vectors of 128 dimensions:
W x.
Both Z and z, encoders are l-layer transformer encoders with attention heads and an intermediate size of 128. The transformer layer utilizes the multi-head self-attention mechanism. For each attention head
Q: W;9 e, K; = w:" Vi V==W""
(e)a: softmax Q.K. JV.
Vdk
where is the dimension of Q; and K. The outputs the attention heads are then aggregated follows: Multihead(e) = concat(Atte (e), Att.) 2-layer MIP with 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, and z, respectively. The functional classifier is 2-layer MLP with 32-dimension hidden layer. The decoder is 2-layer LSTM with 256 hidden dimensions.
The hyperparameters selected with grid search and models with the best generation results reported. Specifically, weights of all losses are selected from [1.0, 0.1]. The dimension of 2 is fixed and x, to 32. We train each model with 200 epochs and learning rate of le 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 the results significantly, so the value is fixed 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. 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:
[B1 1.0, 82 0.1, epoch 200]
where Is's weights of the losses:
L Lrecon BB _ol B2hw
For the visualization and analysis of the model trained VDJDB, we 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 0.5 probability of using the previous predicted token, instead of the original token, to calculate the hidden state the next position. This 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 batch size of 256 (128 per GPU). Training with 200 epochs typically takes ~ 4 hours.
15

---

C
Baseline methods
We compare our model with types of methods for the generation of the optimized TCR x': (1) mutation- based, which iteratively adds random mutations the template sequence; and (2) generation- which generates novel sequences of the pre-determined length range. For both types of methods, the modified//ennerate sequences are selected by peptide binding scores from the respective pre-trained ERGO. The experiments are performed 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 rums for each TCR and the resulting one with the highest ERGO prediction score 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 the template for the next rum. This process is repeated times.
Genetic algorithm (genetic) Let M be the sample size. For each TCR, 10 randomly mutated sequences are generated, 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- baselines
Monte Carlo tree search (MCTS) TCRs are generated by adding amino acids iteratively. resulting in 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 the output TCR set.
C.3 TCR-dVAE
TCR dVAE uses similar objective as IDEL [16., which is VAE with mutual information constraint on the latent space. For training. the loss comprises of the following components:
The reconstruction loss: Lrecon (x, x')
The KL divergence term for VAE: LKL = DKL(golz, 2y(x) || (p(2,,2)))
The reconstruction loss given L. (x, ((z.)) where & is auxiliary decoder.
The classification loss given zy: Ly_ete (ύ.9)
The basele-based upper bound between the embeddings: Lar(zy, 2.). This requires an approxi- mation of the conditional distribution p(z, (2.). which is achieved by separate neural network.
Here use
notations, the ones used in the original paper, for better comparison.
For the hyperparameters,
10.0 weight for Crecons L. and Cy 1.0 for the other terms, and learning rate le 4. practice, we perform annealing [88] on Lx and LM1 where their weights gradually increase from to 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 L.STM-base autoencoder, which we denote as TCR-AEO, on the 277 million TCR sequences from TCRdb 44). TCR-AB0 has a latent space of dimension 16 and is trained for 50,000 steps with batch size of 256.
16

---

D.2
Validity score
The validity score combines two scores calculated from TCR-AEO:
The reconstruction--aasee score is calculated as
rx((x)) = lev(x', (CC-AE0(x'))/1(x>
where lev(x', TCR-AE0(X')) ( is the Levenstein distance between the original sequence and the recon- structed sequence, and [(x') is the length of the reconstructed sequence. Higher r, means x' is better reconstructed from TCR-AE0 and is thus more likely to be valid TCR sequence.
The density-basee
calculates whether the embedding of x' follows the known TCRs. We learn Gaussian mixture model from the latent embeddings of known TCRs from TCR-AEO. The likelihood of the embedding of x' from TCR-AEO falling in the same Gaussian mixture distribution is denoted as P(e'). The density-based score is calculated as
distribution as
Ta ) = exp(1 + log P(e'),
where 10. Higher ra means the latent embedding of x' from TCR-AEO is more likely to follow the same distribution as other valid TCR sequences.
We then define the validity score as
ra-
D.3 Validation of the metrics
We compare the TCR-AE-derived evaluation metric scores three different sources:
(1) all unique CDR33 sequences from VDIDB.
(2) random segments length uniport 633 protein sequences of the same size as (1). The conservative 'C' at the beginning and 'F' at the end are added to the segments.
(which is most frequent lengths of CDR33 sequences) from random
(3) random shuffling of the sequences from (I), where the first 'C" and the last 'F" are kept
We show in Fig. that for both two scores and Pr, as well their sum, CDR33 sequences score much higher than random proteins shuffled sequences. This shows these scores could be effectively used for the estimation of TCR sequence validity. We choose Ty 1.25 the criteria for valid sequences 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 T. Also, the majority of generated sequences are unique (Table TT) and all are novel (not observed the original dataset; statistics not shown).
The relative performance between methods comparison still holds when all peptides, including those that are poorly classified, are taken into account (Table However, when the results are separated by individual peptides (Table we observe that for the poorly classified peptides (AUC<0)8; see Methods), using positive z row) has little difference from using random z (TCR-dWAE (null) row). This further highlights the importance of well-defined and well-classified functional labels. Otherwise, z, does encode the proper functional information, so becomes uninformative and does not alter the sequence function better than random embedding.
E.2 Analysis of the Model
We show extended z and u, T-SNE patterns in Fig. colored the ground truth label as well the predicted label. For the well-classified peptides, there is clear separation of positives and negatives the z space but not z,, There are cases where the true positives are not separable from the true negatives using 2,, but the predicted positives and the predicted negatives (by the function classifier V)) are still separated. We consider the latter as 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. right) the Wasserstein loss.
As a result of Wasserstein loss, the distribution of embedding space is closer to a multivariate Gaussian (Fig. 7A. becomes less regularized without the Wasserstein loss (Fig. BB. Contrary to 2 (Fig. [3 in the main text), the T-SNE of z, and lirst-layer embedding of the encoder for the positive samples cannot distinguish the binding targets from each other (Fig. N)
17

---

E.3 Analysis of the Generated Sequences
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 2 with the most positive/negative one in the subset. The generated sequences using their original and the is positive negative. This shows zf
2 have binding scores mostly related to the z,, regardless
of whether the z, information, which lays foundation for the following TCR engineering experiments (Fig. (B).
used to encode and transfer binding
The generated sequences have a similar length distribution as their templates (Fig. SC), meaning no drastic changes are made. We further find that the 2, of the modified TCRs show high cosine similarity with those of their templates, while the Zy are more similar to the used for their generation (Fig. (D)) but not with that of the template. These show that the modified TCRs preserve the "structural" information from 2, and incorporate the new "functional" information from the modified z,-
As in Fig. E there significant correlation between the sequence validity and the binding scores. Also, find that most binding scores predicted by ERGO are rather extreme (either or O).
E.4 Interpolation
(1) For each peptide. we perform interpolation of 2 between 100 random pairs for 10 steps with 5 r2/ +
(1-77 f)) (2) r (0, 1]. while z, remains the same. Fig. 8F shows that interpolating between positive z 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
Interpolated zy, with Zg of Seq 1
Sequence
CASTESDRRSQNTQYP 1.0 (original)
CASTESDRRSSNVTQYF
1.0
CASTESDRRSQNTQYY
1.0
CASTESDRNSMQPOQF
1.0
CASTESDRNSSOOPOFF
1.0
CASTESTKNNSNQQQQF
1.0
CASTESTKNSSOPPQYF
1.0
CASSESTMNSNQQQQYF
1.0
CASSESTTISSOPPQYF
1.0
CASSESTTNSMQPQYF
1.0
CASSESTTANTAQLFF 1.0
Interpolated zy, with Za Seq 2
Sequence
Tb
CASSLSIFFTANTAQQFF
1.0 (original)
CASSLSTFFTATTAQLFFFF
1.0
CASSLSIFFTAVTAQLFF
1.0
CASSLSTFFTANTAQLFF
1.0
CASSLSIFFTANTAQQLF
1.0
CASSISTFFTANTAQLFF
1.0
CASSLSTRRTANTAQLLF
1.0
CASSLSTRTTANNTQLFF
0.16
CASSLSTTRTAKNTOLFF
0.97
CASSLSTRITSQNQOYF
1.0
CASSLSTRDSSTNTQFF 1.0
These results further prove that z, could used to transfer binding information independent of z., and that the latent space is well-regularized. Also, throughout the interpolation, the hypervariable region only undergoes minor changes. This shows z. could preserve the structural backbone information within both the conservative and hypervariable regions.
18

---

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
E.5.1 Source of positive embeddings
In addition using 2 from random positive sample (labeled random pos), also try other means of obtaining positive z zs's s: (1) best: we use the 2 the sample with the highest classification score from the auxiliary classifier V// (2) avg: we use the average ž, for all positive samples. In both cases, the same 2 is used for all generations. We perform the same experiments with TCR-dVAE, by using the mean of the 2 of the positive sample. We observe that all three methods have comparable performance for TCR-, indicating well-regularized latent space.
E.5.2 Stochastic generation
As TCR-dVAE is probabilistic model, also compared with a random sampling-based method. We both the predicted mean and variance of the positive source to generate random x, and take the average results of all the generated samples (labeled (CR-dVAE-rand) In this scenario, only the random pos and best modes are applicable. We observe notable decrease in the performance as the quality of generation varies with the random sampling, probably due to less regularized latent space.
E.5.3 Transformer decoder
We further test a TCR-dWAE model with transformer decoder (TCR-dWAE-ard) instead of a LSTM decoder. The performance comparable with those with LSTM (Table while the autoregressive gencration is notably slower (Table[0)) 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 I1, 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, indicate some interactions between the two regions that are essential for positive binding. Thus, the modeling and generation should be performed on the entire CDR33 region. instead of on manually defined segments.
be from the number of valid sequences. This may
E.6
Comparison between evaluation metries
E.6.1 Binding prediction
Besides ERGO 23 Ne(TCR[6) (6) is another sequence- TCR-peptide binding model with competitive performance, However, NetTCR is only trained on three peptides and shows the best performance when both the and f chain are provided, which is mostly not available in our dataset. Thus, ERGO provides a more efficient evaluation method for single-chain data and greater variety of peptides. For the comparison between model architectures, we re-train single-chain NeITCR on result on the generated sequences Table in manuscript (reporting the %positive valid only) with two classifiers is as follows. We show in Table [2 that the conclusions are similar between the two models.
training set with more peptides. The evaluation
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:
Nov(æ) min d(x, so
where Do is the set of templates and the distance measure d is the Levenshtein distance relative to the length of the longer sequence. This validation set as a reference. We that around 1/3 ~ 1/2 the sequences remain at cutoff 0.2 and are excluded at cutoff 0.4 (Table [1].. The same pattern is observed in the validation set, which indicates a similar distribution of "novelty" between the generated sequences and the raw data.
then used to select the "novel" sequences. We also include held-out
We also report the same novelty and diversity metrics for the positive valid subset as in point out that:
64 We would like to
19

---

(I) These metrics evaluate the set of generated sequences whole, while we directly select the sequences that meet the validity and binding score criterion, they not our primary objective.
(2) The comparisons 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). addition, we include the case 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 VAE drops lot. These suggest that the positive valid sequences presented in the paper sufficiently diverse and novel.
20

---

source
#pos
auroc
aupr
AVEDRKSDAK
vdjdh
1641
0.94
0.71
CTPYDINQM
vdjdb
500
0.99
0.81
ELAGIGILTV
vdjdb
1410
0.95
0.79
BRDTVDRRYKTLRAEDAAQEE
vdjdb
367
0.98
0.85
GII.GFFFLL
vdjdh
3408
0.95
0.89
GLCTLLVAMI
vdjdh
962
0.92
0.73
NTDESVIK
vdjdh
548
0.94
0.62
KRWIILGLNK
vdjdb
319
0.95
0.54
NLVPMVATV
vdjdb
4421
0.94
0.85
RAKFKOLL
vdjdb
830
0.94
0.75
SSLENFRAYV
vdjdb
322
0,99
0.57
SSYRRPVGI
vdjdb
337
0.99
0.81
STPESANL
vdjdb
234
0.99
0.35
TTPESANL
vdjdb
511
0.99
0.75
ASNENMETM
mepas
26.5
0.98
0.63
CRVLCCYVI.
mepas
435
0.95
EAAGIGILTV
mepas
272
0.97
0.55
FRCPRRFCF
mepas
266
0.96
0.58
GILGFVFTL
mepas
1142
0.96
0.9
GLCTLVVMM
mcpas
828
0.95
0.85
LPRRSGAAGA
mepas
2142
0.96
0.88
NLVPMVATV
mepas
543
0.93
0.78
RFYKTLRAEQASQ
mepas
304
0.99
0.91
SSLENFRAYV
mepas
416
0.99
0.78
SSYRRPYGI
mepas
337
0.99
0.83
TPRVTGGGAM
mepas
274
0.95
0.52
VTEHDTLLY
mepas
273
0.95
0.45
WEDLFCCDESLSSPEPPSSEE
mepas
364
0.98
0.93
Table Statistics and ERGO prediction performance for the selected peptides from the first round.
VDJDB
MCPAS
#pos
#all
#pos
fall
NLYPMVATV
2880
5478
NLVPMVATV
1792
3810
GLCTLVAML
2880
5478
RFYKTLR TEQASS
1528
3579
RAKFKOLL
2880
5478
WEDLFCCDISLSSPEPPSSE
1928
3929
AVFDRKSDAK
5478
GILGEVFTL
2560
4482
SSYRRPVGI
2268
4934
SSYRRPVGI
1504
3558
GILGEVETL
2880
5478
SSLENFRAYV
1824
3838
TTPESANL
2286
4950
CRVLCCYVL
1680
3712
TRDYVDREYKTLRREOAASSR
2034
4726
LPRRSGAAGA
2560
4482
ELAGIGILTV
5478
GLCTIVAML
2560
4482
CTPYDINQM
2394
5046
Table 6: Statistics of the training data by peptide.
21

---

VDIDB
Tv
%valid
#mut/len
%positive valid
TCR-dWAR
1.47=0.02
0.31+0.04
0.76=0.03
0.47+0.06
0.18=0.01
TCR-dVAE
1.51=0.01
0.2+0.01
0.73=0.02
0.38+0.02
0.15=0.01
greedy
0.31=0.0
0.88+0.0
0.02=0.0
0.32+0.0
0.02=0.0
genetic
0.36=0.02
1.040.0
0.02+0.0
1.0+0.0
0.02+0.0
naive rm
0,3320.0
0.39+0.0
0.02+0.0
0.35+0.0
0.01+0.0
MCTS
-0.1=0.0
0.95+0.0
0.0+0.0
0.0+0.0
0.0=0.0
TCR-dWAE (null)
1.51=0.01
0.08+0.02
0.86#0.01
0.45+00.06
0.07=0.02
original
1.59
0.03
0.92
0.0
0.03
Table 7: Performance comparison for VDJDB, averaged across all peptides (ELAGIGILTV, GLCTL- VAML, AVFDRKSDAK, SSYRRPVGI, RAKFKQLL, CTPYDINQM, TTPESANL, NLVPMVATV, FRDDVDRRFYKILRAEEAAQQEE GILGFVFTL)
22

---

vvllid
#mut/len
Spositive valid
SSYRRPVGI
TCR-dWAE
1.11=0.14
0.48+0.18
0.37=0.16
0.48+0.06
0.06#0.01
TCR-dVAE
1.35=0.06
0.25+0.07
0.68=0.05
0.38+0.02
0.06=0.002
TCR-dWAE (null)
1.5+0.01
0.03+0.05
0.85=0.01
0.43+0.05
0.02=0.04
original
1.59
0.02
0.92
0.0
0.01
TTPESANL
TCR-dWAE
1.41=0.06
0.6440.11
0.56=0.06
0.5240.04
0.34+0.06
TCR-dVAE
1.42=0.02
0.4+0.07
0.74+0.02
0.46+0.01
0.1000.01
TCR-dWAE (null)
1.5=0.01
0.05+0.04
0.85=0.01
0.46+0.08
0.04=0.03
original
1.59
0.01
0.92
0.0
0.01
BRDYYVDRFYKKTRRFFAASSO
TCR-dWAR
1.59=0.03
0.28+0.06
0.86=0.03
0.47=0.07
0.22=0.04
TCR-dVAE
1.58=0.02
0.35+0.06
0.91#0.01
0.3940.02
0.22+0.01
TCR-dWAE (null)
1.51=0.01
0.05+0.04
0.86#0.01
0,45+0.07
0.04+0.03
original
1.59
0.02
0.92
0.0
0.01
CTPYDINQM
TCR-dWAE
1.36=0.03
0.51+0.07
0.66=0.04
0.49+0.05
0.29=0.04
TCR-dVAE
1.41=0.02
0.3620.06
0.76#0.03
0.43+0.02
0.18
TCR-dWAE (null)
1.5+0.01
0.08+0.02
0.85=0.01
0.4540.07
0.06#0.02
original
1.59
0.02
0.92
0.0
0.01
ELAGIGILTV*
TCR-dWAE
1.5620.03
0.140.01
0,9+0.02
0.44+0.06
0.09#0.01
TCR-dVAE
1.58=0.02
0.13+0.02
0.92=0.02
0.3240.02
0.1+0.01
TCR-dWAE (null)
1.51=0.02
0.08+0.0
0.86=0.02
0,45+0.05
0.07=0.0
original
1.59
0.05
0.92
0.0
0.04
GI.CTIVAMI.*
TCR-dWAE
1.5620.02
0.11+0.01
0.89=0.02
0.47+0.06
0.1+0.01
TCR-dVAE
1.5820.02
0.14+0.03
0.91=0.02
0.3540.03
0.110..02
TCR-dWAE (null)
1.51=0.01
0.09+0.0
0.85=0.01
0.46+0.06
0.07=0.01
original
1.59
0.04
0.92
0.0 0.04
AVFDRKSDAK*
TCR-dWAE
1.58=0.03
0.15+0.01
0.9a0.02
0.45+0.06
0.13=0.01
TCR-dVAE
1.59=0.02
0.19:0.02
0.91#0.01
0.33+0.02
0.1420.01
TCR-dWAE (null)
1.51=0.01
0.15+0.01
0.85+0.01
0.46+0.06
0.1240.01
original
1.59
0.15
0.92
0.0
0.14
RAKFKOLL*
TCR-dWAE
1.58=0.03
0.08+0.0
0.91=0.02
0.43+0.07
0.08=0.0
TCR-dVAE
1.58=0.02
0.13+0.01
0.92=0.01
0.32+0.02
0.0.00
TCR-dWAE (null)
1.51=0.01
0.07+0.01
0.86#0.01
0,44+0.05 0.0 0.04
0.06#0.01
original
1.59
0.04
0.92
NLVPMVATV
TCR-dWAE
1.57=0.02
0.2+0.0
0.9=0.02
0.45+0.07
0.1740.01
TCR-dVAF
1.59=0.02
0.21+0.0
0.92=0.02
0.33+0.02
0.17=0.0
TCR-dWAR (null)
1.51=0.01
0.18+0.01
0.86+0.01
0.46+0.06
0.15
original
1.59
0.08
0.92
0,0
0.07
GILGEVFTL*
TCR-dWAE
1,56=0.02
0.19+0.05
0.89=0.02
0.46+0.06
0.1740.04
TCR-dVAE
1.57=0.02
0,23+0.04
0.91=0.02
0.34+0.02
0.18:0.02
TCR-dWAE (null)
1.51+0.01
0.140.0
0.86#0.01
0.44+0.06
0.09=0.0
original
1.59
0.05
0.92
0.0
0.05
Table 8:
Performance comparison for VDJDB, separated by peptides. Peptides with * have
classification AUC < 0.8.
23

---

MCPAS
Fv
%valid
#mut/len
%positive valid
TCR-dWAE
1.39=0.07
0.29+0.03
0.69=0.06
0.47+0.03
0.17=0.01
TCR-dVAE
1.45=0.01
0.31
0.47=0.04
0.4+0.02
0.1+0.01
greedy
0,33+0.0
0.92+0.0
0.02=0.0
0.33+0.0
0.02+0.0
genetic
0.34+0.03
1.0#0.0
0.02=0.0
NA
0.02=0.0
naive rm
0.31
0.43+0.0
0.02=0.0
0.35+0.01
0.01=0.0
MCTS
-0.11
0.94#0.0
0.0=0.0
NA
0.0+0.0
TCR-dWAE (null)
1.45#0.06
0.08+0.0
0.79#0.05
0.41+0.03
0.06#0.01
Table 9: Performance comparison for MCPAS, averaged across selected peptides (SSYRRPVGI, WEDLFCDESSLSSEPPSSSE, SSLENFRAYV, RFYKTLRAEQASO, GLCTLVAML, CRVLC- CYVL)
VDJDB
MCPAS
valid:all
unique:valid
running time
valid:all
unique:valid
TCR-dWAE
0.79=0.02
0.95+0.01
56
0.72+0.07
0.95#0.01
TCR-dWAF-att
0.77=0.06
0.97=0.01
5553
TCR-dWAF (null)
0.86+0.01
1.0=0.0
56
0.8+0.05
0,99=0.0
TCR-dVAE
0.84=0.01
0.65+0.02
58
0.8+0.01
0081=0.06
greedy
0.02+0.0
1,0+0.0
832
0.02+0.0
1.0=0.0
genetic
0.02+0.0
0.73+0.04
1389
0.0340.0
0.71
naive rm
0.02+0.0
1.00..0
106
0.02+0.0
1.0+0.0
MCTS
0.0+0.0
0.0=0.0
4559
0.0+0.0
0.04#0.08
Table 10: Additional performance comparison. This table shows the ratio of valid sequences and unique valid sequences, as well as the running time per 5000 samples for CTPYDINQM.
Fv
F6
%valid
#mut/len
%positive valid f
TCR-dWAE (best)
1.33+0.07
0.53=0.18
0,45+0.14
0.5+0.05
0.18+0.05
TCR-dWAE (avg)
1.37+0.13
0.51=0.21
0.53+0.18
0.46=0.09
0.2+0.06
TCR-dWAE (random pos)
1.36+0.03
0.48=0009
0.61#0.06
0.49=0.05
0.23+0.02
TCR-dVAE-rane (best)
0.31+0.06
0.01=0.01
0.0=0.0
0.21=0.3
0.0=0.0
TCR-dVAE-rand (random pos)
1.44+0.02
0.35#0.05
0.45+0.04
0,4+0.01
0.1
TCR-dWAE-attn (random pos)
1.46+0.06
0.32+0.08
0.75+0.06
0.54+0.05
0.19+0.04
TCR-dWAE (trimmed)
1.35+0.08
0.19=0.07
0.68+0.08
0.78=0.0
0.11
TCR-dWAE (null)
1.51+0.01
0.05=0.03
0.85+0.01
0.45=0.07
0.04+0.03
original
1.59
0.01
0.92
0.0
0.01
Table 11: Comparison between model designs and modes of obtaining positive zy's.
24

---

ERGO
NetTCR
TCR-dWAE
0.24+0.01
0.2+0.01
TCR-dVAB
0.16+0.01
0.15+0.01
greedy
0.02+0.0
0.01=0.0
genetic
0.02+0.0
0.0+0.0
naive rm
0.0=0.0
0.0+0.0
MCTS
0.0=0.0
0.0+0.0
Table 12: Performance using different evaluation metrics (ERGO and NeTTCR)
diversity
novelty
TCR-dWAE
0.20+0.0
0.20=0.0
TCR-dVAE
0.23+0.0
0.19=0,0
VAE (collapse)
0.14+0.01
0.17=0.01
Validation
0.26
0.18
Table 13: Diversity and novelty scores for the generated sequences.
0.1
0.2
0.4
TCR-dWAE 1174.6=38.7 TCR-dVAR
1130.4+34.0
1+2.00 290.9+25.6
0.1+0.1 0.0+0.0
809.6+49.5
774.6+49.8
Validation
621
551
191
Table 14: Number of "novel" sequences with different novelty score cutoffs.
25

---

10-
00
ol3
eandoon plotel
shotted
racdon protatn
stortee
copp
randor protein
sholited
Figure 4: Distribution of AE--aEd evaluation metrics on known CDR38's, randomly selected protein segments and randomly shuffled CDR33's.
e-larell
PANTAL
CPACONE
-16
-30
colanesi
(mrall
Itrud
Arignodi
CPYOMN
e-Prael
e-forest
di Itruel
(trocc
By Aareedl
(prad)
A-ppodd
16
-18
-28
16
CAVAN
33
(maa)
gpadd
Itruel
A-1ppedi
grods
Itruel
(predi
Spredt
20 Bruut
2: Sorest
Itrunl
-1
-10
-38
Figure 5: T-SNE of 2 and z, embeddings for all peptides in VDJDB (left) and MCPAS (right). Points are colored by the label. "True" means the ground truth label. "Pred" refers to label predcited by the function classifier V.
26

---

0110.1 seed=42
000000 seed=42
0.6
0.4
0.4
0.2
0.2
0.0
0.0
0.00
0.25
0.75
0.00
0.25
0.50
0.75
1,0.11 456
01=0.0,
446d
1.0
1.0
0.8
0.8
0.6
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0,10.1, seed
0000.0 seed- 789
1.0
0.8
0.8
0.6
0.6
0.4
0.4
0.2
0.2
0.0
0.0
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
01:0.1,
seed=987
0100.0,
seed=987
TTPESANL (auroc 0.98)
SSYRAPVGI (auroc 0.94)
FRDYVDRFYKLRREOASOE (auroc 0.86)
CTPYDINQM (auroc 0.91)
LGFVFTL 0.79)
0.4
GLCTLVAML 0.78)
NLVPMVATV (auroc 0.72)
RAKFKQLL (aurec 0.76)
0.2
0.2
AVEDRKSDAK (auroc D.64)
ELAGIGILTV (auroc 0.66)
0.0
0.0
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Figure 6: ROC of function classifier V by peptide, with different hyperparameter settings and random seeds.
27

---

A
(B)
32
40
Figure 7: Distribution of the latent embeddings with (A) and without (B) Wasserstein loss. Orange lines correspond to dimensions of Zy and green lines 2f. The distribution is estimated using gaussian kde from the scipy package.
28

---

(A)
postove sequances
input embedding positive sequences
21
3.5
2.2
SOYSAPVI
TIPESALL
CTRIDIOON
-30
Soquence length distriutionn
binding
CTPYDINOM
TIPESANL
1.00
1200
0.75
050
20
10
20
SSYRAPVGI
OPPLOWOM
FIPESSR
SCYMPPSID
Generated pontivve
Template
20
10
20
(D)
Generated
Genorated
Generated modrficd
Codbv da
8.5
2.0
00-
800
CIVORTION
PRONEN MAASEi
TTRESAR
PROMEER TETUREIAAA
CTP rasso
PRO PRO Betrutioe
Strmepett
(E)
Correl ation between validity and score
(F)
Interpolation of
PROYVOR/VKTLMMEOSSS
SSYRAPVG
pearson 0.18
pearson -0.31
65-
nig
CTPYDINOM
TIPESANL
FRDIVORPPRTLRREEDS
DAX
790
Figure 8: (A) T-SNE of n, (left) and first layer embedding the encoder (right) of positive TCRs, colored by their binding peptides. (B) The average binding score of generated positive and negative TCRs. (C) The length distribution of template and optimized TCRs (CDR33 region) from VDJDB. (D) Cosine similarity between z, of the optimized sequences their templates (left), z of the optimized sequences vs their templates (middle), z of the optimized sequences vs the modified (right). (E) Scatter plot between the validity score and binding score of the engineered sequences (for simplicity, only 500 out of 5000 points are shown). (F) Binding score of interpolated samples between positive and negative pairs.
29


Reviews:
[{"review_id": "UnQKPj7h78", "reviewer": "Reviewer_ULRU", "summary": "This 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. ", "questions": "Q1. 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.", "limitations": "NA", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The 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. ", "weaknesses": "The proposed method is not novel but its application to a new domain is of interest.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your comments. Here we respond to your points and questions in the sequential order:\n\n**1. Additional results for Table 1** \n\nThe other six peptides were discarded because of poor classification (with AUC < 0.8). As we rely on the classifier to control the information content of the functional embedding, poor classification leads to uninformative disentanglement, i.e. the functional embedding does not really encode the correct binding information. \n\nStatistics averaged over all 10 peptides:\n|                   | $\\bar{r_v}$   | $\\bar{r_b}$   |  %valid   | #mut/len   | %positive valid   |\n|:------------------|:------------|:----------|:-----------------|:--------------|:--------------------------|\n| TCR-dWAE (best)   | 1.43±0.02   | 0.37±0.09 | 0.55±0.09        | 0.5±0.05      | 0.15±0.04                 |\n| TCR-dWAE (avg pos)    | 1.48±0.07   | 0.32±0.11 | 0.66±0.11        | 0.43±0.08     | 0.15±0.03                 |\n| TCR-dWAE (random pos) | 1.47±0.02   | 0.31±0.04 | 0.76±0.03        | 0.47±0.06     | 0.18±0.01                 |\n| IDEL  (random pos)  | 1.51±0.01   | 0.26±0.03 | 0.55±0.02        | 0.38±0.02     | 0.1±0.01                  |\n| TCR-dWAE (null)   | 1.51±0.01   | 0.08±0.02 | 0.86±0.01        | 0.45±0.06     | 0.07±0.02                 |\n| original          | 1.59    | 0.03  | 0.92        | 0.0      | 0.03                  |\n\n\n\nStatistics separated by peptides, under \"random pos\" mode (due to space limits, we only show 2 well classified and 2 poorly classified examples, but the pattern is similar for all. The full table will be added to the new version):\n\n*FRDYVDRFYKTLRAEQASQE (AUC = 0.88)*\n|                   | $\\bar{r_v}$   | $\\bar{r_b}$   |  %valid   | #mut/len   | %positive valid   |\n|:----------------|:------------|:----------|:-----------------|:--------------|:--------------------------|\n| TCR-dWAE   (random pos)   | 1.59±0.03   | 0.28±0.06 | 0.86±0.03        | 0.47±0.07     | 0.22±0.04                 |\n| IDEL  (random pos)    | 1.58±0.01   | 0.35±0.05 | 0.5±0.03         | 0.39±0.02     | 0.12±0.02                 |\n| TCR-dWAE (null) | 1.51±0.01   | 0.05±0.04 | 0.86±0.01        | 0.45±0.07     | 0.04±0.03                 |\n\n\n*CTPYDINQM (AUC = 0.91)*\n|                   | $\\bar{r_v}$   | $\\bar{r_b}$   |  %valid   | #mut/len   | %positive valid   |\n|:----------------|:------------|:----------|:-----------------|:--------------|:--------------------------|\n| TCR-dWAE  (random pos)  | 1.36±0.03   | 0.51±0.07 | 0.66±0.04        | 0.49±0.05     | 0.29±0.04                 |\n| IDEL  (random pos)  | 1.4±0.02    | 0.37±0.07 | 0.49±0.05        | 0.44±0.02     | 0.13±0.01                 |\n| TCR-dWAE (null) | 1.5±0.01    | 0.08±0.02 | 0.85±0.01        | 0.45±0.07     | 0.06±0.02                 |\n\n*ELAGIGILTV (AUC=0.66)*\n|                   | $\\bar{r_v}$   | $\\bar{r_b}$   |  %valid   | #mut/len   | %positive valid   |\n|:----------------|:------------|:----------|:-----------------|:--------------|:--------------------------|\n| TCR-dWAE  (random pos)  | 1.57±0.03   | 0.1±0.01  | 0.9±0.02         | 0.44±0.06     | 0.09±0.01                 |\n| IDEL  (random pos) | 1.58±0.02   | 0.13±0.01 | 0.64±0.02        | 0.32±0.02     | 0.07±0.01                 |\n| TCR-dWAE (null) | 1.51±0.02   | 0.08±0.0  | 0.86±0.02        | 0.45±0.05     | 0.07±0.0                  |\n\n*AVFDRKSDAK (AUC = 0.64)*\n|                   | $\\bar{r_v}$   | $\\bar{r_b}$   |  %valid   | #mut/len   | %positive valid   |\n|:----------------|:------------|:----------|:-----------------|:--------------|:--------------------------|\n| TCR-dWAE   (random pos)    | 1.58±0.03   | 0.15±0.01 | 0.9±0.02         | 0.45±0.06     | 0.13±0.01                 |\n| IDEL   (random pos)    | 1.59±0.02   | 0.19±0.02 | 0.64±0.03        | 0.34±0.02     | 0.1±0.01                  |\n| TCR-dWAE (null) | 1.51±0.01   | 0.15±0.01 | 0.85±0.01        | 0.46±0.06     | 0.12±0.01                 |\n\nWe also exclude some baselines as they are uniformly poor.\n\nFor the poorly classified peptides (AUC < 0.8),  using a positive $z_f$’s is not much different from using a random $z_f$ (the \"null\" row), meaning in this case the $z_f$ is uninformative.\n\nThe correlation between $r_v$ and $r_b$ by peptide is shown in Attached Fig. C. As expected, there is no significant correlation.\n\n**2. Structural similarity**\n\nWe are focusing on the CDR3beta region, a loop whose structure is difficult to predict. We do acknowledge that it’s hard to directly quantify structural similarities with our current method, due to lack of 3D structural data and reliable prediction methods. Thus, our “structural pattern” can only be used to indirectly measure structural similarity through sequence similarity. Specifically, the autoencoder for the validation could learn frequent sequential patterns and their importance among the a set of known TCRs, and we assume that this indirectly implies structural similarities, based on the assumption that the structure is defined by the sequence. We will state these limitations in the next version.\n"}, {"review_id": "c36EDTCVR5", "reviewer": "Reviewer_mx2A", "summary": "In 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.", "questions": "Are 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", "limitations": "Yes", "rating": 7, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The 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. ", "weaknesses": "Because 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your comments. Here we respond to your points and questions in the sequential order:\n\n**1. Validity metrics**\n\nWe follow the same evaluation metrics as in [1]. Specifically, with a trained autoencoder, we use (1) reconstruction accuracy (2) the likelihood in the distribution of the embeddings of known TCR sequences (estimated by a Gaussian mixture model) as a metric of the “likelihood” of a new TCR. This is similar to autoencoder-based out-of-distribution (OOD) detection methods (refs below).\n\nTCRdb is a much larger database with 277 million TCRs. Most unique CDR3beta sequences in VDJDB should appear in it (we didn’t do a comprehensive search, but the majority of a subset we have tried can be found in TCRdb). However, as it is thousands of times larger than the entire VDJDB, we believe the autoencoder trained on it should learn more generic patterns that can be used for the validation.\n\n“Has a similar embedding pattern as the known TCRs” means when the generated sequences are fed to the autoencoder trained on TCRdb, their embedding patterns have a high likelihood in the distribution of the embeddings of known TCRs. We also show in Appendix Fig. S1 that our metric could separate real TCRs from other protein segments or random sequences very well. Thus, we believe it is sufficient for the evaluation. \n\nRef:\n[1] Chen, Ziqi, et al. \"T-Cell Receptor Optimization with Reinforcement Learning and Mutation Polices for Precision Immunotherapy.\" International Conference on Research in Computational Molecular Biology. Cham: Springer Nature Switzerland, 2023.\n[2] Pimentel, Marco AF, et al. \"A review of novelty detection.\" Signal processing 99 (2014): 215-249.\n[3] Zong, Bo, et al. \"Deep autoencoding gaussian mixture model for unsupervised anomaly detection.\" International conference on learning representations. 2018.\n\n**2. Number of mutations**\n\nSorry for the confusion. In fact, the average number of mutations reported in Table 1 is divided by the sequence length, not the raw number. This means we have on average ~50% of the original sequence modified. We will clarify this in the future. Thanks for pointing this out.\n\n**3. Structural vs function within the hypervariable region; interpolation**\n\nWe have done interpolation on 100 positive and negative pairs (Attached Fig E). Overall, interpolation between positive $z_f$’s generates far more positive samples than negative ones. We also observe in an example (Attached Fig B) that throughout the interpolation, changing $z_f$ only affects some parts of the hypervariable region, and sometimes it also modifies the conservative region, while some others from both regions are preserved.\n\nFrom these observations, we can have two conclusions:\n(1) It’s not that $z_s$ is solely responsible for the reconstruction of the conservative region. \n(2) There is still a functional/structural separation within the hypervariable region, where some positions can be modified to alter functions while some cannot.\n\nThis is further supported by Attached Fig A that the average attention patterns of the functional and structural encoders, despite having their own positional preferences, overlap at some positions in the hypervariable region. Neither patterns align with the boundary of hypervariable/conservative regions. \n\nThus, the functional and structural embeddings may have learned higher-level interactions between sequential patterns all across the sequence, including but not limited to the hypervariable region. \n\n**4. Miyazawa-Jernigan energy**\n\nMiyazawa-Jerningan energy is based on a heuristic score matrix of pairwise interaction energy between amino acids. We sum up the score of each pair of “interacting” amino acids (whose distance falls below a certain threshold; in this case, we consider all amino acids in the CDR3beta and the peptide interact with each other).\n"}, {"review_id": "dN7CUHO4lb", "reviewer": "Reviewer_TMVT", "summary": "In 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. ", "questions": "How 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. ", "limitations": "As 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. ", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The 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. ", "weaknesses": "## 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your comments. Here we respond to your points and questions in the sequential order:\n\n**1. Limited scope**\n\nWe focus on TCR here because it is easier to define a quantifiable “function” (binding vs non-binding). For many other scenarios, such a definition might be ambiguous. Also, there is quite a large amount of labelled TCR-peptide pairs that are sufficient for machine learning, while for other applications it may require careful data collection and curation. We would like to leave them for future studies.\n\nWe do agree that the claim on “protein engineering” as a whole is too broad for the paper. Thus, we will limit our claim to the engineering of TCR CDR3beta and modify the title, introduction and conclusion accordingly.\n\n**2. Conservative and hypervariable regions**\n\nWe observe that the structural and functional patterns learned by the model could span the entire sequence, as modifying $z_f$ could potentially lead to modifications in both the conservative region and the hypervariable region (Fig 2B and Attached Fig B), and that the functional and structural encoders attend to some overlapping positions, which does not correspond to the boundary of hypervariable/conservative regions (Attached Fig A). Thus, a manual cut between a “functional” hypervariable region and a “structural” conservative region may not be a good idea. \n\n**3. Motif scaffolding**\n\nWe very much appreciate the suggestion. We will add a paragraph to the introduction and related work on motif scaffolding. \n\n**4. Table 1; Pareto efficiency**\n\nThanks for the suggestion. We find that the pareto efficiency analysis is less informative in our case because the predicted binding scores are rather extreme (either 0 or 1; Attached Fig. C), so the pareto front always overlaps with $r_b=1$ (Attached Fig. D) . Thus, we would still like to keep Table 1 as we feel that the count of valid binders from a hard cutoff on both scores, rather than the scores themselves, is a better evaluation criterion. We will also improve its readability later.\n\n**5. Theorem 2.1**\n\nFrom a theoretical perspective, the disentanglement we want to achieve has two parts:\n\n1. Independence, or minimal mutual information, between the dimensions. This is indeed enforced by the Gaussian prior.\n\n2. Besides independence, we also would want different parts of the embedding (i.e. $z_s$ and $z_f$) to correctly encode the desired information. Otherwise, the disentanglement would be uninformative. $z_s$ should have information about the majority of the sequence. $z_f$ should have information on some key parts of the sequence and the functional label. Thus, both embeddings need maximal mutual information with the raw sequence $x$. The theorem proves that this objective can be achieved.\n\n**6. The use of “structural embedding”**\n\nThe “structure” here means “sequential patterns that define the structural backbone”. For simplicity, we use the word “structural”. However, we are also open to suggestions for more accurate wording.\n\nIt is hard to predict the “structure” in the sense of protein 3D structure because the CDR3 region is a very flexible loop, and high-quality structure and reliable prediction methods are scarce. Thus, we can only rely on the available sequential information of known CDR3beta's to determine whether generated ones are valid, or whether they preserve the structural backbone, based on the intuition that the structure is defined by the sequence. If certain pivotal residues and motifs (which we try to make the “structural embedding” learn) are similar between two sequences (e.g. generated and known), they should have similar structures. We will state its limitations in the discussion section of the paper.\n\n**7. Fig. 3C**\n\nThanks for the suggestion. We will modify the figure accordingly.\n\n**8. Generalization to proteins**\n\nThat should depend on the availability and nature of data. If there is a weaker decomposition between scaffold and function (or any feature in general), or if there is not a clear definition of the function, the model may not be able to learn a meaningful disentanglement between such features (e.g. the poorly classified peptides not included in Table 1). In fact, we devise our model based on the assumption that there is some natural yet unknown decomposition within the data. We will further discuss these limitations in the discussion section of the paper.\n\n\n**9. Explicit decomposition**\n\nIf you mean hand-picking the structural and functional parts of the sequence, it would be impractical on a larger scale as we can only rely on some rough heuristic rules that heavily depend on the context.  As we've shown in question 2, in our problem, the functional and structural patterns cannot be manually separated and need to be learned from the data.\n\n**10. LSTM and transformer decoder**\n\nAs our dataset is not very complex, we believe an LSTM is sufficient and is easier to implement. Also, in our application, there is only one embedding for the whole sequence, while a typical transformer decoder would accept token-wise embeddings. The latter may not be easily manipulated for sequence engineering. \n\nNevertheless, we have experimented with a slightly modified transformer decoder that accepts one single embedding. The performance as shown below is similar to LSTM, while its autoregressive sampling is much slower. \n\n|                   | $\\bar{r_v} $  | $\\bar{r_b}$   | %valid   | #mut/len   | %positive valid   |\n|:------------------|:------------|:----------|:-----------------|:--------------|:--------------------------|\n| TCR-dWAE (random pos) | 1.36±0.03   | 0.48±0.09 | 0.61±0.06        | 0.49±0.05     | 0.23±0.02                 |\n| TCR-dWAE-transformer  (random pos)  | 1.46±0.06   | 0.32±0.08 | 0.75±0.06        | 0.54±0.05     | 0.19±0.04                 |\n\n**11. Hyperparameters & hardware**\n\nHyperparameters and hardware settings are specified in Appendix 2.2. \n"}, {"review_id": "t9U0HcBoPS", "reviewer": "Reviewer_4twu", "summary": "The 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.", "questions": "- 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.", "limitations": "The 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.", "rating": 5, "confidence": 4, "soundness": 1, "presentation": 2, "contribution": 2, "strengths": "- 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.", "weaknesses": "- 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your comments. Here we respond to your points and questions in the sequential order:\n\n **1. Explicit and implicit**\n\nBy explicit and implicit, here we are talking about the nature of the latent space. For some generative models, such as diffusion models,  there is no “explicit” latent space. Methods like VAE and WAE “explicitly” model the distribution of the latent space, so we could directly enforce disentanglement on that latent space.\n\nOn the other hand, the “explicit” you mean with the examples is manually defining functional and structural parts of the sequence. There is no such prior knowledge for the CDR3beta region. Thus, we need to learn self-supervised disentanglement on an **explicit latent space** because we cannot do such separation manually and “explicitly”.\n\nSpecifically, motif scaffolding is a reverse problem of ours. It extends a functional “motif” by generating a “scaffold”, while we seek to modify the function of an existing TCR sequence. It also requires prior knowledge of the functional part, while in our case it needs to be learned from the data. Thus, it is not directly comparable here. \n\n\n**2. Usage of TCR**\n\nIn fact, we actually are focusing on the CDR3beta region, as mentioned in 2.1. We also find that the function/structure separation is non-trivial within CDR3beta as the model learns patterns all across the region.\n\n**3. Structural backbone**\n\nWe acknowledge that our study is limited by the lack of high-quality TCR structures or reliable prediction methods. Thus, we can only estimate structural similarity indirectly through the embedding patterns of the sequence, based on the intuition that the sequence defines the structure.\n\n**4. Related work; NetTCR**\n\nNetTCR is a classifier. It cannot be directly used for sequence engineering. \n\nMost computational antibody engineering methods either use physical modeling, which is slow on a large scale, or rely on 3D structural data, which is scarce for TCR. Thus, they are not comparable to ours.\n\n**5. 0.5 as the cutoff**\n\nTo select positive $z_f$’s in \"random\" and \"avg\" modes, we actually do not use the classifier $\\Psi$ (which is used only in \"best\" mode). Instead, we directly use the $z_f$’s  of the sequences with positive labels from the dataset. The information in 3.2.1 is inaccurate. Thanks for pointing out. We will correct it in the next version.\n\n**6. Symbols**\n\nBeta_1 and beta_2 are the weights of the losses. The values are 1.0 and 0.1, as mentioned in Appendix 2.2. Σ means sum. We will further clarify it in the manuscript.\n\n**7. WAE and the latent space**\n\nWAE models the distribution of the latent space by enforcing a shape on it. The training and decoding are deterministic, but it still has a latent space distribution that could be sampled from.\n\nFor our experiments, however, we manipulate existing embeddings and deterministically decode them to generate new samples, **with no random sampling involved**, which is different from VAEs. \n\nThe Gaussian distribution facilitates such generation because it has a regular shape. If one manipulates  (Section 3.1) a latent embedding, the resulting “new” embedding is highly likely to remain in the distribution and thus could be decoded properly. With an irregular shape, it may fall out of the distribution leading to poor-quality generations.\n\n\n**8. Well-classified by ERGO**\n\nWe have explained the selection process in Appendix 1.2. We will also clarify it in the main text.\n\n**9.“Develop our own classifiers”**\n\nERGO and NetTCR all take the **raw sequence** as the input. In order to manipulate the latent space, however, we’ll need a classifier on the **latent embeddings**. It is used as an auxiliary to control the content of the embeddings.\n\n**10. Additional dataset**\n\nWe select 5000 additional sequences from VDJDB which are not seen in (\"do not overlap with\") the training set.\n\n**11.\"Best\" and “random” mode**\n\nFor the “best” method we obtain the $z_f$’s from the sequence in the dataset with the highest binding score from $\\Psi$. Also, we’re not sure why the reviewer believes this would lead to adversarial samples. \n\nAs stated in 3.2.1, the “random” method uses the $z_f$ from a random **positive** sample. With a well-regularized latent space, any positive $z_f$ should carry patterns related to positive binding. Thus the high performance of “random” is expected. For better clarity, we will rename it to \"random pos\".\n\n**12. Using the autoencoder, not LSTM or transformer**\n\nThe goal here is to quantify how likely the generated sequence comes from the same distribution of known TCRs (i.e. valid), similar to autoencoder-based out-of-distribution detection (OOD) methods (refs below). For such a task, continuous-valued embeddings from an autoencoder are more computationally efficient. We show in Appendix Fig. S1 that these metrics can separate real TCRs from random protein segments.\n\nAlso, transformers and LSTMs can be used in many model architectures, including autoencoders. In fact, we are using **LSTM-based autoencoder** here as is mentioned in Appendix 4.1. \n\n**13. Comparison to IDEL**\n\nThanks for bringing this up.  We have re-done the experiment with 5 random samples for each embedding, and choose the one with the highest binding score. \n\n**14. Meaning of unique and novel**\n\nSorry for the confusion. We’re looking for literally “unique” and “novel” sequences through direct sequence comparison, and there is no sequence similarity involved.\n\n**15. Failure cases**\n\nOur method would fail when:\n\n(1) There is no sufficient data. \n\n(2) The functional label is not well-defined or well-classified (e.g. the poorly classified peptides we do not include in Table 1),  or when there is no inherent (yet unknown) separation between the features. \n\n"}]

Author rebuttal:
We really appreciate the reviewers’ valuable comments. Here we would like to address some common concerns that several reviewers have brought up.

We have also attached a pdf with the figures for the rebuttal, and refer to them as “Attached Fig XX”.

**1. The usage of “structure embedding”**

We agree that the wording is a little confusing. The “structure” here means “sequential patterns that relate to the 3D structural backbone”.  For simplicity, we use the word “structural” in the rest of the paper. However, we are also open to suggestions of more accurate wording.

In practice, it is difficult to predict the real “structure” in the sense of protein 3D structure for CDR3beta region because it is a very flexible loop, and high-quality structure is scarce. Additionally, it cannot be predicted from the CDR3beta sequence alone, ignoring the rest of the TCR. Thus, we can only rely on the available sequential information of known CDR3beta's to determine whether generated ones are valid, or whether they preserve the structural backbone, based on the intuition that the structure is defined by the sequence. Thus,  if certain pivotal residues and motifs (which we try to make the “structural embedding” learn) are similar between two sequences (e.g. generated and known), they should have similar structures. 

**2. Separating structures and functions**

There are also concerns that the separation of “structural” and “functional” embedding in the latent space is unnecessary, as CDR3beta already has a conservative region (first four and last three residues) and a hypervariable region, which can be manually defined as the “structural” and “functional” parts of the sequence. Moreover, it might be possible to remove the conservative region entirely as the structural embedding might only learn to memorize those residues. 

We would like to address that from the following perspectives:

(1) There are actually quite a few variations within the conservative regions in real TCRs. It is uncertain whether it interacts with the hypervariable region to maintain validity. For example, some combinations might be considered invalid. Thus, we think it is necessary to still include it into the modeling.

(2) From the engineering experiments, we discover that by only modifying $z_f$, there are changes in both the conservative region and the hypervariable region, though the latter is more frequent. On the other hand, the variable region is not fully replaced and some residues are maintained. This can be seen in the mutation hotspots in Fig. 2B and the interpolation example in Attached Fig. B. For the interpolation experiment, we obtain the embedding of two sequences (CASTESDRRSQNTQYF and  CASSLSTFTANTAQLFF) that binds to peptide CTPYDINQM. We interpolate their  $z_f$ by  $z^{(r)}_f = rz^{(1)}_f + (1-r)z^{(2)}_f, r\in[0,1]$, and combine it with the $z_s$ of each sample to do the generation. We observe that throughout the process, some amino acids in both regions are gradually modified from one binding mode to another while the rest remains.

(3) We further show in the Attached Fig. A that the attention patterns of the functional and structural encoders, despite having their own positional preferences, also overlap at some positions. In addition, neither agrees with the boundary of the conservative region and the hypervariable region.

These observations indicate that both the structural and functional patterns the model learns are higher-level patterns and could span both the conservative and hypervariable regions, and cannot be fit into manually defined windows. Thus, we believe a manual separation or exclusion of certain residues is not possible, and such patterns need to be learned from the data.


**3. Generalization to other proteins**

Deep learning-based protein engineering is limited by the availability of data. In our specific scenario, we need both a well-defined functional label, and samples that are large enough for the machine learning. To our knowledge, there have not been a lot of problems that meet both requirements. Because of these concerns, we limit our experiments to the CDR3beta region of TCR.

Nevertheless, we do agree that “protein engineering” is too broad a claim for the experiments in the paper. Thus, we will change the title to TCR engineering, and clearly state the limitations of our study in the next version.

**4. Updated Table 1**

In response to some reviewers’ suggestions, we have updated the results in Table 1 in two ways: (1) we have included the results of 1 additional random seed, resulting in 5 in total; (2) as IDEL is a VAE, we do random generation 5 times given the mean and variance of $z_f$ and $z_s$ and choose the one with the highest binding score, instead of directly using the mean as in the previous Table 1. “Best” and “avg” modes do not apply to this generation procedure so we exclude them for IDEL. Some baselines are not included here because their performance remains low after the change.

|                   | $\bar{r_v}$   | $\bar{r_b}$   |  %valid   | #mut/len   | %positive valid   |
|:------------------|:------------|:----------|:-----------------|:--------------|:--------------------------|
| TCR-dWAE (best)   | 1.33±0.07   | 0.53±0.18 | 0.45±0.14        | 0.5±0.05      | 0.18±0.05                 |
| TCR-dWAE (avg pos)    | 1.37±0.13   | 0.51±0.21 | 0.53±0.18        | 0.46±0.09     | 0.2±0.06                  |
| TCR-dWAE (random pos) | 1.36±0.03   | 0.48±0.09 | 0.61±0.06        | 0.49±0.05     | 0.23±0.02                 |
| IDEL  (random pos)  | 1.44±0.02   | 0.35±0.05 | 0.45±0.04        | 0.42±0.01     | 0.1±0.01                  |
| TCR-dWAE (null)   | 1.51±0.01   | 0.05±0.03 | 0.85±0.01        | 0.45±0.07     | 0.04±0.03                 |
| original          | 1.59    | 0.01  | 0.92         | 0.0       | 0.01                  |




Meta-review: This paper proposes a method for T-cell receptor engineering, which is designed to preserve the 3D structure while making functional changes (i.e. binding) to the protein of interest. The method consists of training a disentangled Wasserstein autoencoder with an auxiliary classifier to learn separate functional and structural embeddings. The reviewers all agreed that the proposed method (disentangled WAE) is not novel but its application to a new and important domain is interesting and well-designed. Reviewers were somewhat skeptical of the evaluation and potential limited applicability of this approach, however, the authors addressed these concerns during the rebuttal period and will incorporate their feedback into the manuscript, including updating the title.
Decision: Accept (poster)