Contrastive learning of T cell receptor representations (2406.06397v2)

Abstract: Computational prediction of the interaction of T cell receptors (TCRs) and their ligands is a grand challenge in immunology. Despite advances in high-throughput assays, specificity-labelled TCR data remains sparse. In other domains, the pre-training of LLMs on unlabelled data has been successfully used to address data bottlenecks. However, it is unclear how to best pre-train protein LLMs for TCR specificity prediction. Here we introduce a TCR LLM called SCEPTR (Simple Contrastive Embedding of the Primary sequence of T cell Receptors), capable of data-efficient transfer learning. Through our model, we introduce a novel pre-training strategy combining autocontrastive learning and masked-LLMling, which enables SCEPTR to achieve its state-of-the-art performance. In contrast, existing protein LLMs and a variant of SCEPTR pre-trained without autocontrastive learning are outperformed by sequence alignment-based methods. We anticipate that contrastive learning will be a useful paradigm to decode the rules of TCR specificity.

• The paper introduces SCEPTR, a novel TCR language model that integrates autocontrastive learning with masked language modeling for enhanced specificity prediction.
• It demonstrates that SCEPTR outperforms traditional protein language models and sequence alignment methods in few-shot TCR specificity benchmarks.
• Ablation studies underscore how autocontrastive learning and tailored position embeddings mitigate VDJ biases and improve overall model performance.

Contrastive Learning of T Cell Receptor Representations

The paper under review presents an innovative approach to computational prediction of T cell receptor (TCR) specificity through a newly developed TCR LLM named SCEPTR (S). The model addresses the challenge of predicting TCR interactions with peptides presented by major histocompatibility complexes (pMHCs), a significant task in immunology. Despite advancements, the specificity-labelled TCR data is sparse, necessitating more efficient models that can leverage abundant unlabelled data for enhanced predictive accuracy.

The authors propose a novel pre-training strategy for SCEPTR, combining autocontrastive learning and masked-LLMling (MLM). This strategy allows the model to achieve state-of-the-art performance in data-efficient transfer learning. This paper's findings contrast with existing protein LLMs, which the authors demonstrate are outperformed by sequence alignment-based methods in few-shot settings.

Key Findings and Contributions

1. Benchmarking Results:

The authors establish a robust benchmarking framework focused on evaluating model performance in few-shot TCR specificity prediction. The results reveal that:

• Existing protein LLMs (PLMs) such as ProtBert, ESM2, and TCR-BERT perform worse than sequence alignment methods like TCRdist in few-shot settings.
• The newly introduced SCEPTR model achieves superior or at least comparable performance to TCRdist across various pMHCs, outperforming other PLMs significantly.

2. SCEPTR Architecture and Pre-Training:

SCEPTR employs a simplified transformer architecture with a unique token embedding method and a contextualised embedding of a <cls> token for TCR representation. The key innovation lies in the joint optimization of MLM and autocontrastive learning:

• MLM teaches the model to predict masked tokens within a sequence, while autocontrastive learning encourages the model to discriminate between TCR identities by minimizing distances between positive pairs and maximizing distances with a background distribution.
• The result is a model that effectively utilizes the diverse sequence space, mitigating biases introduced by VDJ recombination, and aligning better with TCR-specificity tasks.

3. Ablation Studies:

The extensive ablation studies highlight the importance of various components of the SCEPTR model:

• Autocontrastive learning is pivotal for the model’s performance, significantly enhancing specificity prediction.
• The simplified position embedding strategy of SCEPTR leads to a better-calibrated representation, as opposed to traditional left-aligned methods employed by other models like TCR-BERT.
• The alignment and uniformity terms in contrastive learning further eliminate biases from VDJ recombination, effectively focusing on specificity.

4. Fine-Tuning with Supervised Contrastive Learning:

To further refine TCR specificity predictions, the authors explore supervised contrastive learning:

• Fine-tuned SCEPTR excels in discriminating between specific pMHCs it was trained on, although this specialization comes at the cost of generalizability to unseen pMHC specificities.
• This suggests potential avenues for improving generalization through larger labels and combined unsupervised-supervised learning merges.

Implications and Future Developments

The introduction of SCEPTR demonstrates that autocontrastive learning can significantly enhance the predictive accuracy of TCR specificity, outperforming current PLMs. The implications of this work are both practical and theoretical:

• Practical Applications: SCEPTR can be integrated into TCR analysis pipelines, aiding in the discovery and characterization of TCRs in relation to specific antigens. This can streamline the development of T-cell-based immunotherapies and vaccines.
• Theoretical Advancements: The combination of MLM and contrastive learning sets a precedent for future PLMs addressing non-structure related predictions. It shows potential, especially in domains with high sequence variability and limited labelled data.

Looking ahead, leveraging larger datasets for supervised fine-tuning and integrating multi-modal data, such as phenotypic annotations from single-cell sequencing, may further enhance the model's robustness and generalizability. Additionally, future work could explore how larger model architectures scale performance under the new pre-training paradigm presented herein.

In conclusion, this paper presents a significant advancement in TCR LLMling through the novel implementation of autocontrastive learning, setting a new benchmark in few-shot TCR specificity prediction. The thorough analyses and ablation studies offer valuable insights into improving model performance and present promising avenues for future research in immunological bioinformatics.