Scalable Universal T-Cell Receptor Embeddings from Adaptive Immune Repertoires

Thanks for the insightful comments. On i) clinical relevance and ii) model interpretability, please refer to the overall response above.

The paper lacks rigorous benchmarks against established methods beyond simple logistic regression. Disease and HLA classification tasks do not adequately demonstrate the model’s robustness, especially with limited sensitivity for certain conditions (e.g., HSV) at larger embedding scales.

Response

The benchmarks used for comparison are the most rigorous available. Prior work found little to no improvement over logistic regression for disease prediction tasks (Pradier et al., 2023; Kanduri et al., 2022). Please refer to the shared response above on the motivation for classification tasks and comparisons with prior works for more details.

The limited sensitivity observed in certain tasks reflects the complex and inherently ambiguous biology rather than model limitations. For instance, a challenge with HSV is that there are two closely related viruses, HSV-1 and HSV-2, which share a significant portion of their genome, resulting in a subset of identical antigens presented to the immune system when infected by either virus. Distinguishing whether HSV-specific TCRs recognize unique antigens from HSV-1 or HSV-2, or shared antigens presented by both, is inherently difficult for any model given the sparsity and high dimensionality of repertoire data. Notably, the task-agnostic JL-Glove repertoire embeddings perform as well as the AIRIVA task-specific model trained to disentangle these two viruses (Table 3).

How does the model handle rare but potentially significant TCRs? Given the emphasis on co-occurrence, it is unclear how rare TCRs are represented, as these could provide unique insights in immune responses but may not frequently co-occur with other TCRs.

Response

Due to their enormous range in generation probability through random V(D)J recombination, a significant fraction of TCRs are only observed in one subject and thus are not part of our co-occurrence analysis. Our goal however is not to capture all possible signals in the TCR repertoire but to aggregate sufficient information to associate a large fraction of embedded TCRs with HLAs and diseases. By aggregating these associations, we generate a dense vector representation that encodes both immune genetics and disease exposure history. The method presented is a significant step toward this objective. As datasets grow and more repertoires are sequenced, we will capture additional signals from rarer TCRs (Figure 5).

How does the embedding perform across other immunological datasets?

Response

We focus on learning TCR embeddings from bulk TCR sequencing data, as it provides deep sampling, each repertoire is characterized by TCRs in $O(10^7)$ and a large sample size $N$, which is crucial for learning meaningful TCR embeddings, as illustrated in Figures 5 and 10. Most publicly available labeled single-cell TCR sequencing data have low sample sizes in the order $< 30$ samples, and the number of TCRs per patient is even lower $O(10^4)$, making it difficult to benchmark on repertoire classification tasks. However, the proposed TCR embeddings could be combined with other immunological representations, both at the T cell level (e.g., single-cell RNA) and the patient level (e.g., B cell receptors repertoires and histopathology), to form a comprehensive mulit-modal representation.

Additional clarity on computational scaling challenges would be helpful, especially given the potential high-dimensional space of TCR repertoires

Response

JL-GloVe addresses scalability in training the GloVe objective in two ways: (i) computing the normalized JL embeddings for a large number of TCRs $K$ is extremely fast since we don't need to compute the co-occurrence matrix $C$ (see Equations 2, Proposition 2.2, and Algorithm 2 in the Appendix), and (ii) by initializing with normalized JL embeddings, the GloVe algorithm converges faster in fewer epochs, enabling training with only a fraction of the co-occurrence matrix $C$ while achieving similar performance on downstream tasks (see Figures 3, 14, and 15). Additionally, we compute the co-occurrence matrix on CPUs using distributed learning in Spark according to Algorithm 1 in the Appendix and leverage PyTorch Lightning's distributed data parallel training to train across GPUs (see Appendix A.6). We understand the resource constraints within the community, therefore, all TCR embeddings derived from the public cohorts, along with the PyTorch code, will be made publicly available.

What additional features could improve biological relevance?

Response

The experimental results show that the proposed TCR embeddings encode immune genetics (HLA) and pathogenic exposure history. Additional multi-modal learning with other T cell modalities, such as single-cell RNA, could enhance the encoding of TCR function. We consider this important future work.

We thank the reviewer for engaging during the discussion period. We are deeply grateful that they found our rebuttal addressed most of their concerns and promised to increase their score to recommend acceptance of our paper upon revision.

We would like to inform the reviewer that a revised version of our paper, with changes highlighted in blue, has been uploaded. Additionally, we gently remind the reviewer that they can currently adjust their score during the discussion period.

Thanks for the insightful comments. On i) interpretability and ii) comparisons to previous works, please refer to the overall comments above.

The biological definitions presented in the study are somewhat unclear. For instance, when the authors refer to TCR embedding, it is important to specify whether they mean both the TCR alpha and beta full chains, the CDR3 regions of both chains, or only the CDR3 region of the TCR beta chain. Additionally, do the authors take into account V(D)J gene information when using the CDR3 region?

Response

We thank the reviewer for this comment and have clarified in Section 2 that we embed the bioidentity of the TCR$\beta$ chain, where the bioidentity is defined as the CDR3+Vgene+Jgene. However, the JL-GloVe algorithm is general and could be repurposed to learn embeddings of the TCR$\alpha$ chain or both chains, given large TCR repertoire datasets.

Given that TCRs are highly cross-reactive, the authors need to provide further explanation on why using co-occurrence information alone is effective for TCR embedding.

Response

TCRs do indeed have a high theoretical potential for cross-reactivity. However, evidence suggests that meaningful cross-reactivity in vivo is rare (Ishizuka et al., 2009; Petrova et al., 2012). Previous work using disease-labeled repertoire datasets have also identified extensive sets of TCRs strongly associated with individual diseases (Emerson et al., 2017; Greissl et al., 2021; May et al., 2024) and thus not cross-reactive to any common antigenic exposures. One exception are homologous antigens, like those responding to HSV-1 and HSV-2, or other homologous viruses.

The repertoires of different subjects contain varying numbers of TCRs. How do the authors address this variability when representing them with a matrix of the same TCR dimensionality?

Response

In Section 3, we detail how we derive repertoire embeddings, given a varying number of TCRs per repertoire (subject), via a set mean pooling transformation of TCR embedding dimensions, as specified in Equation 8. Refer to overall comment above on generating TCR repertoire embeddings for more details.

Considering the high cross-reactivity of TCRs, how do the authors define the TCR-level ground truth without relying on wet-lab-based experiments?

Response

Previous work has identified TCRs associated with infections such as CMV (Emerson et al., 2017), SARS-CoV-2 (Snyder et al., 2020), EBV, HSV-1, HSV-2, Parvovirus (May et al., 2024), and Lyme disease (Greissl et al., 2021), achieving high sensitivity and specificity in distinguishing cases from controls. The observed high specificity of these TCRs, which make up a substantial fraction of our embedded TCRs, would not be achievable if substantial cross-reactivity were prevalent in vivo. In the context of SARS-CoV-2 infections, specificity of TCRs (i.e., ground level truth) has been established via wet-lab experiments using functional assays directly demonstrating T cell activation by antigens (Snyder et al., 2020).

These in vitro experiments have been validated statistically in vivo. (Pradier et al., 2023) show that TCRs purported to bind spike and non-spike derived antigens by MIRA do indeed separate subjects who have natural infection from vaccination, further demonstrating that the functional assay used to associate TCRs to antigens is reproducing in vivo biology. In this work, we rely on these wet lab experiments to provide ground truth for our validation analysis demonstrating that TCRs binding antigens derived from the same disease have high cosine similarity (Figure 6a) and to show that subjects who have natural infection have greater overlap with TCRs associated with our covid classifier than subjects who only have vaccination (Figure 6b).

When discussing classification tasks, it would be helpful to clarify whether the focus is on receptor-level classification or repertoire-level classification.

Response

Thank you for the suggestion. We have clarified that all quantitative classification tasks are at the repertoire level (Section 4.2) and that all TCR-based results are qualitative. Refer to the shared response on motivation for classification benchmarking tasks for more details.

Furthermore, given different receptors have clone frequencies within the repertoire, it appears that the authors do not consider clone frequency in their repertoire-level embedding.

Response

We thank the reviewer for this comment and have added text in Section 2 clarifying that we do not consider clone frequencies within the repertoire, as we have found binary indicators of the presence or absence of TCRs to be more robust.

Thanks for the insightful comments. On i) comparisons to previous works, ii) generating TCR repertoire embeddings and iii) clinical relevance, please refer to the overall comments above.

The authors observe that the disease classification performance is sensitive to the embedding dimension (d) and the number of TCRs (K). A more systematic exploration of the impact of these parameters can be done A more detailed analysis of the impact of different embedding dimensions across various dataset sizes would be valuable.This would aid other researchers in configuring JL-GLOVE for datasets of different sizes or resolutions, thereby increasing the framework’s accessibility and practical utility.

Response

We seek to highlight Figure 13, which provides additional experimental results exploring the sensitivity of the dimension $d$ and the number of TCRs $K$, indicating that in general, we need embedding dimensions $d \ge 100$ to obtain reasonable disease prediction performance. Further, theoretical results (Theorem 2.1) show that $d=O(\epsilon^{-2} \log K)$ for distortion $0<\epsilon^{-2}<1$.

Thanks for the insightful comments.

I believe this paper should be in the topic area of applications to physical sciences (biology / immunology) rather than unsupervised, self-supervised, semi-supervised, and supervised representation learning. There are few novelties in terms of methods development, but the application of these methods are extremely impactful and a nice way to show the power of representation learning.

Response

We agree that the application of the JL-GloVe algorithm to TCR embeddings is a novel contribution and fits well within the immunology community. However, this work will also be of interest to the broader machine learning community for two key reasons: i) The proposed JL-GloVe algorithm is general and could be repurposed to learn scalable co-occurrence-based representations with lower memory usage and faster training time, especially in contexts where data structures do not align with current representation learning paradigms used in other domains such as text or images; and ii) Multi-modal representation learning at both the T cell and patient levels will benefit from TCR embeddings that account for immune genetics and exposure history. Given these aspects, we prefer to maintain our current topic area.

You mention other methods for set level representations, why not use those? Average is nice in it's simplicity, but does this simplicity cost performance? Would be nice to see a benchmark against set representation methods. You may also be interested in OTKE method for set level representation.

Response

Thank you for bringing the OTKE method to our attention; we have included it as an alternative set pooling method in the paper (Section 3). Refer to shared response above on generating TCR repertoire embeddings for more details.

Motivation for classification benchmarking tasks (Reviewers Qp7p, DdmS)

We seek to clarify that the JL-GloVe TCR embeddings encode immune genetics (HLA) and pathogenic exposure history. In Section 4.2, we present quantitative repertoire-level classification tasks based on shared task-agnostic repertoire embeddings generated according to Equation 8, primarily to validate this hypothesis. Here, by task-agnostic, we simply mean that we have not derived these embeddings by optimizing for any specific downstream modeling application. We emphasize our goal is not to build the most optimal model for any specific task but rather to demonstrate the generalizability and robustness of the inferred TCR embeddings across 5 disease prediction and 145 HLA inference binary classification tasks. We have clarified this in Section 4.2.

Comparisons to previous works (Reviewers Qp7p, ovF4, DdmS)

Thank you for bringing these works to our attention. We have included them in the introduction and provided comparisons with DeepRC on the Emerson dataset (Tables 2 and 4). Note that DeepRC reports a best AUROC of $0.831 \pm 0.002$ compared to JL-GloVe's $0.99 \pm 0.02$ and ESLG's $0.93 \pm 0.01$. These results align with previous findings showing that TCR protein sequence models struggle to outperform simple regularized logistic regression models, such as ESLG (Kanduri et al., 2022). We wish to emphasize that the proposed JL-GloVe TCR embeddings are task-agnostic and serve as a complement to the TCR protein sequences used as features in supervised learning approaches, including DeepTCR, TCRAI, DeepAIR, and DeepRC. Since JL-GloVe TCR embeddings encode immune genetics (HLA) and pathogenic exposure history, they can be effectively utilized in conjunction with TCR protein sequences for tasks such as repertoire classification or TCRpMHC binding prediction, which are employed in these methods. Furthermore, qualitative results (Figure 7) demonstrate that TCR protein sequence-based embeddings fail to cluster the embedding space by HLA or disease exposure, unlike the proposed JL-GloVe co-occurrence-based TCR embeddings. Additionally, the JL-GloVe algorithm is general and could be repurposed for other immunological datasets, including B-cell receptor repertoires used in DeepID. We have emphasized this aspect in Sect 2.3

Clinical relevance (Reviewers ovF4, DdmS)

We illustrate the power of JL-GloVe embeddings by showing that they can distinguish between individuals vaccinated against COVID-19 and those with natural infections based on distinct co-occurrence patterns (Figure 6b). A similar separation is observed for TCRs associated with HSV-1 and HSV-2 infections (Table 3). Our approach can facilitate similar analyses for other homologous viruses where labels may be unavailable. The most important clinical applications will be in the ability to study the interdependence between diseases. For instance, a strong link between EBV infection and multiple sclerosis has recently been identified by Bjornevik et al. (2022); uncovering this link required disease exposure history for a large cohort of patients. We plan to release embeddings trained on public cohorts, which we hope will benefit research groups with smaller datasets to explore disease interdependency.

Generating TCR repertoire embeddings (Reviewers Qp7p, ovF4, x8Nf):

We opted for parameter-free mean pooling primarily for simplicity and easier interpretability in repertoire classification tasks, as we wanted to highlight that the TCR embeddings capture immune genetics and pathogenic exposure history without learning additional parameters. The proposed simple mean pooling approach will serve as a strong baseline for more sophisticated set pooling approaches, including OTKE and set transformers. We agree that experimental results show that as we increase the number of TCRs K, mean pooling results in noisier repertoire embeddings. This necessitates the sophisticated methodological development that adapts a set-based pooling mechanism to derive task-agnostic (universal) repertoire embeddings, which account for the challenges posed by TCR repertoire data, which are heterogeneous, high-dimensional, sparse, and mostly unlabeled. We consider this to be a key area for future work.

Interpretability of TCR embeddings (Reviewers Qp7p, DdmS):

Note that the JL-GloVe algorithm is derived from TCR co-occurrence to maximize the cosine similarity of co-occurring TCRs, along with an interpretable monotonic function derived from marginal occurrence to down-weight commonly occurring TCRs. Wet-lab results demonstrate that TCRs targeting the same antigen have high cosine similarity (Figures 4a and 6a). Additionally, we assume a simple linear mean pooling function for repertoire embeddings, which enables: i) stratification of repertoires by antigen exposure (Figure 6b), and ii) recovery of disease/HLA-associated TCRs (Figure 12) through simple dot products.