††thanks: Joint last authors††thanks: Joint last authors

# Contrastive learning of T cell receptor representations

Yuta Nagano  Division of Medicine, University College London  Division of
Infection and Immunity, University College London  Andrew Pyo  Center for the
Physics of Biological Function, Princeton University  Martina Milighetti
Division of Infection and Immunity, University College London  Cancer
Institute, University College London  James Henderson  Division of Infection
and Immunity, University College London  Institute for the Physics of Living
Systems, University College London  John Shawe-Taylor  Department of Computer
Science, University College London  Benny Chain  Division of Infection and
Immunity, University College London  Department of Computer Science,
University College London  Andreas Tiffeau-Mayer  Division of Infection and
Immunity, University College London  Institute for the Physics of Living
Systems, University College London

###### 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 language models on unlabelled data has been
successfully used to address data bottlenecks. However, it is unclear how to
best pre-train protein language models for TCR specificity prediction. Here we
introduce a TCR language model 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-language modelling, which
enables SCEPTR to achieve its state-of-the-art performance. In contrast,
existing protein language models 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.

Antigen-specific T cells play important protective and pathogenic roles in
human disease [1]. The recognition of peptides presented on major
histocompatibility complexes (pMHCs) by $\alpha\beta$ T cell receptors (TCRs)
determines the specificity of cellular immune responses [2]. Hyperdiverse
$\alpha\beta$TCRs are generated during T cell development in the thymus by
genetic recombination of germline-encoded V, D (for TCR$\beta$) and J gene
segments with additional diversification through insertions and deletions of
non-template nucleotides at gene segment junctions.

A major goal of systems immunology is to uncover the rules governing which
TCRs interact with which pMHCs [3]. Advances in high-throughput functional
assays of TCR specificity [4, 5, 6] have made the use of machine learning a
promising prospect to discover such rules.

The most direct approach for applying machine learning to TCR specificity
prediction has been to train pMHC-specific models that take an arbitrary TCR
and predict binding [7, 8, 9, 10, 11, 12]. More ambitiously, model
architectures have been proposed that can in principle generalise predictions
to arbitrary pMHCs as well [13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24].
Independent benchmarking studies have shown that both approaches are effective
for predicting TCR binders against pMHCs for which many TCRs have been
experimentally determined [25], but generalisation to pMHCs not seen during
training has largely remained elusive [26] and prediction accuracy is limited
for pMHCs with few known binders [27]. This severely limits the utility of
current predictive tools given that only $\sim 10^{3}$ of the $>10^{15}$
possible pMHCs are currently annotated with any TCRs in VDJdb [28], and given
that for $>95\%$ of them less than 100 specific TCRs are known.

Meanwhile, there is abundant unlabelled TCR sequence data that may be
exploited for unsupervised representation learning. A TCR representation model
that compactly captures important features would provide embeddings useful for
data-efficient training of downstream specificity predictors.

In natural language processing (NLP), unsupervised pre-trained transformers
have demonstrated capacity for transfer learning to diverse downstream tasks
[29, 30, 31]. This has spurred substantial work applying transformers to
protein analysis. Protein language models (PLMs) such as those of the ESM [32,
33] and ProtTrans [34] families have been successfully used in structure-
prediction pipelines and for protein property prediction [35, 36, 37]. PLMs
have also been applied to TCR-pMHC interaction prediction [11, 22, 23, 24],
and the related problem of antibody-antigen interaction prediction [38, 39].
However, there has been limited systematic testing of how competitive PLM
embeddings are in the _few-shot_ setting typical for most ligands – that is,
where only few labelled data points are available for transfer learning.

To address this question, we benchmarked existing PLMs on a standardised few-
shot specificity prediction task, and surprisingly found that they are
inferior to state-of-the-art sequence alignment-based methods. This motivated
us to develop SCEPTR (Simple Contrastive Embedding of the Primary sequence of
T cell Receptors), a novel TCR PLM which closes this gap. Our key innovation
is a pre-training strategy involving an autocontrastive learning procedure
adapted for $\alpha\beta$TCRs, which we show is the primary driver behind
SCEPTR’s improved performance.

Figure 1:  Benchmarking TCR language models against sequence alignment-based
approaches on few-shot TCR specificity prediction. a) TCR similarity can be
quantified using sequence-alignment by taking a (weighted) count of how many
sequence edits turn one TCR into another. b) Learned sequence representations
allow alignment-free sequence comparisons based on distances in the embedding
feature space. c) Sketch of our standardized benchmarking approach to allow
side-by-side comparison of sequence-alignment and embedding methods. Using a
reference set of known TCR binders to a pMHC of interest, we propose nearest-
neighbour prediction as a task for unbiased comparison of the quality of
embeddings for specificity prediction. d) Performance of six different models
on TCR specificity prediction as a function of the number of reference TCRs.
Specificity predictions were made by the nearest neighbour method sketched in
c against six different pMHCs and performance is reported as the AUROC
averaged across the pMHCs.

##  I Results

###  I.1 Benchmarking PLM embeddings on TCR specificity prediction

Given the scarcity of specificity-labelled TCR data, it is of practical
importance to evaluate model performance where access to such data is limited.
Therefore, we set up a benchmarking framework focused on few-shot TCR
specificity prediction.

To conduct our benchmark, we curated a set of specificity-labelled
$\alpha\beta$TCR data from VDJdb [28]. We only included human TCRs with full
$\alpha$ and $\beta$ chain information, and excluded data from an early 10x
Genomics whitepaper [40], as there are known issues with data reliability in
this study [41, 42]. This left us with a total of 7168 $\alpha\beta$TCRs
annotated to 864 pMHCs. Of these, we used the six pMHCs with greater than 300
distinct binder TCRs for our benchmarking task.

We created a benchmarking task that allowed us to directly compare sequence
alignment-based distance metrics such as the state-of-the-art TCRdist [4, 43]
(Fig. 1a) to distances in PLM embedding spaces (Fig. 1b). For each pMHC, we
tested models on their ability to distinguish binder TCRs from non-binders
using embedding distances between a query TCR and its closest neighbour within
a reference set (Fig. 1c). We call this _nearest neighbour prediction_. This
framework is simple and attractive for benchmarking models in the few-shot
regime, since it remains well defined for as few as a single reference TCR and
does not require model specific fine-tuning.

We conducted multiple benchmarks for each pMHC, varying the number of its
cognate TCRs used as the reference set. In each case, we combined the
remaining TCRs for the target with the rest of the filtered VDJdb dataset
(including TCRs annotated to pMHCs other than the six target pMHCs) to create
a test set (see methods III.1). By studying how performance depends on the
size of the reference set, we are effectively probing representation alignment
with TCR co-specificity prediction at different scales.

We benchmarked six models: two alignment-based TCR metrics (CDR3 Levenshtein
distance and TCRdist [4]), two general-purpose PLMs (ProtBert [34] and ESM2
[33]), and two TCR domain-specific language models (TCR-BERT [11] and our own
model SCEPTR). We report performance using the area under the receiver
operator characteristic (AUROC) averaged over the tested pMHCs.

To our surprise, we found that TCR-BERT, ESM2, and ProtBert all fail to
outperform the baseline sequence alignment method (CDR3 Levenshtein) and are
significantly inferior to TCRdist (Figs. 1d / S1). A repeat of the
benchmarking with a broader set of epitopes obtained by including post-
processed 10x Genomics whitepaper data [42], recapitulated these results,
demonstrating the robustness of our findings (Fig. S2). In contrast to
existing PLMs, SCEPTR performs on par with or better than TCRdist. For a
reference set of size 200, SCEPTR performs best among all models for five out
of six tested peptides (Table SI).

We additionally compared models using the average distance between a query TCR
and all references, instead of only the nearest neighbour (Fig. S3). In this
case, SCEPTR outperforms other models by an even wider margin. Interestingly,
all models perform worse compared to their nearest neighbour counterpart. This
finding might be explained mechanistically by the multiplicity of viable
binding solutions with distinct sequence-level features, which are thought to
make up pMHC-specific TCR repertoires [44, 45].

Figure 2:  A visual introduction to how SCEPTR works. a) SCEPTR featurises an
input TCR as the amino acid sequences of its six CDR loops. Each amino acid
residue is vectorised to $\mathbb{R}^{64}$ (see panel b) and are passed along
with the special <cls> token vector through a stack of three self-attention
layers. SCEPTR uses the contextualised embedding of the <cls> token as the
overall TCR representation, in contrast to the average-pooling representations
used by other models. b) SCEPTR’s initial token embedding module uses a simple
one-hot system to encode a token’s amino acid identity and CDR loop number,
and allocates one dimension to encode the token’s relative position within its
CDR loop as a single real-valued scalar. c) Contrastive learning allows us to
explicitly optimise SCEPTR’s representation mapping for TCR co-specificity
prediction. At a high level, contrastive learning encourages representation
models to make full use of the available representation space while keeping
representations of similar input samples close together. d) Contrastive
learning generalises to both the supervised and unsupervised settings. In the
supervised setting, positive pairs can be generated by sampling pairs of TCRs
that are known to bind the same pMHC. In the unsupervised setting, positive
pairs can be generated by generating two independent “views” of the same TCR.
We implement this by only showing a random subset of the input data features
for every view – namely, we remove a proportion of input tokens and sometimes
drop the $\alpha$ or $\beta$ chain entirely (see methods III.3).

###  I.2 Autocontrastive learning as a pre-training strategy

We now briefly summarize SCEPTR’s architecture and autocontrastive pre-
training strategy (see Methods for full details). SCEPTR featurises an input
TCR as the amino acid sequences of its six CDR loops. It uses a simple one-hot
encoding system to embed the amino acid tokens, and uses a stack of three
self-attention layers to generate a 64-dimensional representation vector of
the input receptor (Fig. 2a,b). Unlike existing TCR language models, SCEPTR is
jointly pre-trained using autocontrastive and masked-language modelling (MLM)
(Fig. 2c,d).

To motivate the considerations that have led us to adopt this training
paradigm, some background on transformer architectures and their training by
MLM needs to be introduced. The transformer is a neural network developed in
NLP that uses dot-product attention to flexibly learn long-range dependencies
in sequential data [29]. BERT is an encoder-only variation of the transformer
useful for text analysis and processing [30]. BERT’s innovation was its
ability to be pre-trained in an unsupervised manner through MLM, where
snippets of text are fed to the model with a certain proportion of tokens
(e.g. words) masked, and the model must use the surrounding context to
reconstruct the masked tokens. MLM allowed BERT and its derivative models to
exploit large volumes of unlabelled data to learn grammar and syntax and
achieve high performance on downstream textual tasks with comparatively little
supervised fine-tuning.

While MLM-trained PLMs have been successful in some protein prediction tasks
[32, 33, 34], they have been documented to struggle with others [37]. Our
benchmarking results led us to believe that MLM pre-training may not be
optimal for TCR-pMHC specificity prediction. Firstly, the majority of observed
TCR sequence variation is attributable to the stochastic process of VDJ
recombination. As such, MLM may not teach models much transferable knowledge
for specificity prediction. Secondly, since the low volume of specificity-
labelled TCR data provides limited opportunities for fine-tuning complex
models, representation distances should ideally be directly predictive of co-
specificity.

We were inspired to use contrastive learning to overcome these problems by the
success of our previous work using statistical approaches to uncover patterns
of sequence similarity characteristic of ligand-specific TCR repertoires [44,
46, 45]. Contrastive learning minimises distances between model
representations of positive sample pairs while maximising distances between
background pairs (Fig. 2c) through a loss function of the following form [47,
48]:

$\mathcal{L}_{\textrm{contrastive}}(f)\coloneqq\\\
\operatorname*{\mathbb{E}}_{\begin{subarray}{c}(x,x^{+})\sim
p_{\textrm{pos}}\\\
\\{y_{i}\\}_{i=1}^{N}\overset{\mathrm{iid}}{\sim}p_{\textrm{data}}\end{subarray}}\left[-\log\frac{e^{f(x)^{\top}f(x^{+})}}{e^{f(x)^{\top}f(x^{+})}+\sum_{i}{e^{f(x)^{\top}f(y_{i})}}}\right]$
(1)

where $f:\mathcal{X}\to\mathcal{S}^{m-1}$ is a trainable embedding mapping
from sample observation space $\mathcal{X}$ to points on the m-dimensional
unit hypersphere $\mathcal{S}^{m-1}\subset\mathbb{R}^{m}$, $p_{\textrm{pos}}$
is the joint distribution of positive pairs, $p_{\textrm{data}}$ is the
overall data distribution, and $N\in\mathbb{Z}_{+}$ is some fixed number of
background samples.

There are several well-known variants of this learning approach. In supervised
contrastive learning, positive pairs are generated by sampling observations
known to belong to the same class (Fig. 2d top). In the context of TCRs, we
can define positive pairs to be TCRs annotated to interact with the same pMHC,
in which case contrastive learning regresses distances between TCR pairs to
their probabilities of co-specificity. Autocontrastive learning approximates
such positive pairs through data augmentation by generating two independent
views of the same observation by passing it twice through the model (Fig. 2d
bottom).

Given the scarcity of available labelled data, we opted to use the
autocontrastive approach for purely unsupervised PLM pre-training (see Sec.
I.5 for an application of supervised contrastive learning, more similar to
other recent applications of contrastive learning to TCRs [49, 50]). We
generate different “views” of a TCR by dropout noise as is standard in NLP
[51], but additionally adopted a censoring strategy inspired by masked-
language modeling that randomly removes a proportion of residues or even
complete $\alpha$ or $\beta$ chains. In contrast to the only other study known
to us having explored the application of autocontrastive learning to TCRs
[52], we trained SCEPTR on all six hypervariable loops of the full paired
chain $\alpha\beta$ TCR, as all contribute to TCR-pMHC specificity [45]. That
being said, our chain dropping procedure during censoring ensures that single
chain data are also in distribution for the model, giving SCEPTR flexibility
for downstream applications with bulk sequenced TCR repertoires.

We define SCEPTR’s output representation vector to be a contextualised
embedding of a special input token called <cls> (the naming convention for
<cls> comes from the fact that the output of this vector is often used for
downstream _classification_ [30]), which is always appended to the tokenised
representation of an input TCR (Fig. 2a). This allows SCEPTR to fully exploit
the attention mechanism when generating the overall TCR representation. Such
training of a sequence-level representation is uniquely made possible by
having an objective – the contrastive loss (Eq. 1) – that directly acts on the
representation output. In contrast, MLM-trained PLMs such as ProtBert, ESM2
and TCR-BERT generate sequence embeddings by average-pooling the
contextualised embeddings of each input token at some layer: a destructive
operation which risks diluting information [37].

Figure 3:  Autocontrastive pre-training significantly improves SCEPTR’s
downstream performance. The subplots show performance profiles of SCEPTR,
TCRdist, TCR-BERT, and various ablation variants of SCEPTR on binary
specificity prediction. a) Training SCEPTR solely on MLM results in worse
specificity prediction performance. b) The baseline SCEPTR variant which uses
the <cls> pooling method performs marginally better than the variant which
uses the average-pooling method. However, the average-pooling variant still
performs on par with TCRdist. c) Replacing SCEPTR’s pre-training dataset with
1) the same dataset from Tanno et al., but with $\alpha/\beta$ chain pairing
shuffled, and 2) synthetic data generated by OLGA both result in slightly
inferior specificity prediction performance. d) Restricting SCEPTR’s
featurisation of input TCRs to the amino acids of the $\alpha$ and $\beta$
CDR3 loops significantly worsen downstream performance. Additionally
restricting training to only MLM further degrades performance, and produces a
model with a near-equivalent performance profile to TCR-BERT.

###  I.3 Ablation studies

To understand which modelling choices drive the improved performance of
SCEPTR, we trained variants of SCEPTR ablating a single component of either
its architecture or training at a time, and benchmarked them using the
framework described previously.

To establish the contribution of the autocontrastive learning to SCEPTR’s
performance, we trained the same model only using masked-language modelling:

SCEPTR (MLM only)

This variant is trained only on MLM, without jointly optimising for
autocontrastive learning. Following convention in the transformer field [53],
TCR representation vectors are generated by average-pooling the contextualised
vector embeddings of all constituent amino acid tokens produced by the
penultimate self-attention layer, and $\ell$2-normalising the result.

The MLM-only variant underperforms compared to both SCEPTR and TCRdist,
demonstrating that autocontrastive learning is a necessary ingredient for the
increased performance of SCEPTR in few-shot specificity prediction (Figs. 3a /
S4a).

We next sought to determine how much our pooling strategy and training dataset
choice contributed to SCEPTR’s performance gain. First, we asked whether
autocontrastive learning also improves embeddings generated via token average-
pooling:

SCEPTR (average pooling)

This variant receives both autocontrastive learning and MLM, but uses the
average-pooling method to generate TCR representations.

While SCEPTR’s <cls> embeddings achieve the best results, the autocontrastive
average-pooling variant still performs on par with TCRdist (Figs. 3b / S4b).

Second, we determined how the performance of SCEPTR depends on the precise
dataset used for pre-training. To answer this question we trained two variants
of SCEPTR using size-matched datasets:

SCEPTR (synthetic data)

This variant is trained on a size-matched set unlabelled $\alpha\beta$TCRs
generated by OLGA [54], a probabilistic model of VDJ recombination.

SCEPTR (shuffled data)

This variant is trained on the same set of $\alpha\beta$TCRs as the original
model, but the $\alpha$/$\beta$ chain pairing is randomised.

We find that pairing information and the use of real post-selection
repertoires slightly improves SCEPTR’s performance, but variants trained on
alternative datasets still match TCRdist performance (Figs. 3c / S4c). The
synthetic data models the statistics of VDJ recombination and therefore
estimates a pre-thymic TCR distribution. As such, SCEPTR’s slight performance
improvement using peripheral blood TCRs hints at a potential ability of
language modelling to extract information from imprints left by antigen-driven
selection in memory repertoires [44]. Similarly, the slight performance
increase on paired TCR data suggests some ability to extract information from
pairing biases [55].

Taken together, these ablation studies provide evidence that autocontrastive
learning is the main factor enabling SCEPTR to close the gap between PLMs and
alignment-based methods.

Information-theoretic analysis of the sequence determinants of TCR specificity
demonstrate that all CDR loops and their pairing are important for determining
binding specificity [45]. To understand how much SCEPTR’s improved performance
with respect to TCR-BERT is due to the restriction of the latter model’s input
to the CDR3 alone, we trained variants of SCEPTR restricted to this
hypervariable loop:

SCEPTR (CDR3 only)

This variant only accepts the $\alpha$ and $\beta$ chain CDR3 sequences as
input (without knowledge of the V genes/first two CDR loops of each chain). It
is jointly optimised for MLM and autocontrastive learning.

SCEPTR (CDR3 only, MLM only)

This otherwise equivalent variant is only trained using the MLM objective, and
thus uses the average-pooling representation method.

The results demonstrate that taking into account all CDR loops leads to a
performance gain as expected (Figs. 3d / S4d). We also see that
autocontrastive learning even when restricted to CDR3s leads to a substantial
performance gain, helping the autocontrastive CDR3 variant achieve similar
performance to the full-input MLM-only variant (Fig. 3a,d).

###  I.4 Learning features within embedding spaces

Figure 4:  SCEPTR with nearest-neighbour prediction outperforms other PLMs
with a supervised support vector machine on top. We trained a linear SVC on
top of TCR-BERT’s featurisation of the TCR to predict specificity, and
compared its performance to the nearest neighbour predictions of both SCEPTR
and TCR-BERT using the benchmarking framework from section I.1. We see that
while training an SVC on top of TCR-BERT improves its downstream performance
for few-shot TCR specificity prediction, it still does not outperform nearest
neighbour predictions as made by SCEPTR.

So far we have focused on nearest-neighbor prediction using PLM embeddings as
the most direct test of data-efficient transfer learning that works with as
little as a single reference sequence. If slightly more data is available,
another approach is to train simple supervised predictors atop PLM embeddings.
To test how much such training can improve prediction performance, we trained
linear support vector classifiers (SVC) on the PLM embeddings provided by
different models. In each instance, we trained the classifier to distinguish
reference TCRs from 1000 randomly sampled background TCRs (see methods III.4).

We find that the SVC predictors for ProtBert, ESM2 and TCR-BERT all perform
better than their nearest-neighbour counterparts, but still worse than
SCEPTR’s nearest neighbour predictions (Fig. 4 / S6). We also trained an SVC
atop SCEPTR, which did not lead to further improvement upon the nearest-
neighbor prediction (Fig. S6). These findings highlight how in the low data
regime typical of most pMHCs, misalignment of pre-training to downstream tasks
can only be partially remediated by training on reference TCRs.

###  I.5 Supervised contrastive learning as a fine-tuning strategy

Figure 5:  Supervised contrastive learning improves discrimination between
seen pMHCs. Prediction performance as measured by AUROC on binary one-versus-
rest classification for each of six pMHCs for different models. The fine-tuned
model improves performance by exploiting the discriminative nature of the
overall classification task.

Supervised contrastive learning provides an avenue to further optimise pre-
trained embeddings for TCR specificity prediction. As a proof-of-concept, we
fine-tuned SCEPTR to better discriminate between the six pMHC specificities
used as the benchmarking targets in section I.1.

For this task we took all the TCRs annotated against the target pMHCs from our
labelled TCR dataset, and split them into a training, a validation, and a
testing set. We ensured that no study used for training or validation
contributed any data to the test set, so that the fine-tuned model would not
be able to achieve good performance simply by exploiting inter-dataset biases.
The training set included 200 binders against each target pMHC, totalling to
1200 TCRs. The rest of the TCRs from the same studies were used to construct
the validation set. TCRs from all remaining studies were used for the testing
set, which comprised of 5670 TCRs. SCEPTR was fine-tuned on the training set
with supervised contrastive learning, using the validation loss for early
stopping (appendix III.5).

We used the framework from section I.1 to benchmark the performance of fine-
tuned SCEPTR, using the training set as the references. The results show that
fine-tuning can greatly improve the ability of the model to discriminate
between specificities (Fig. 5). Improvements are most noticeable for the pMHCs
against which other methods achieve relatively low performance. However, model
performance degrades with respect to unseen pMHCs (Fig. S5), perhaps
unsurprisingly given the very limited number of pMHCs represented in the
training data. Interestingly, unlike other models (Fig. S3), fine-tuned SCEPTR
makes better inferences by measuring the average distance between a query TCR
and all reference TCRs instead only the nearest TCR (Fig. S7). This suggests
an ability of supervised contrastive fine-tuning to help the model discover
the commonalities between the multiple different binding solutions thought to
exist for each pMHC.

##  II Discussion

In this study, we have introduced SCEPTR, a pre-trained TCR PLM that achieves
state-of-the-art few-shot TCR-pMHC specificity prediction accuracy. Through
SCEPTR we demonstrate that joint autocontrastive and masked-language pre-
training is a paradigm for learning PLMs better aligned with TCR specificity
prediction tasks. Our model can be readily used for alignment-free TCR
analysis in downstream applications (see code availability) including the
unsupervised discovery of antigen-specific T cell groups (metaclonotypes) by
sequence-based clustering [4, 43].

A limitation of our study is that we did not undertake a complete exploration
of training and architectural hyperparameters. We envisage multiple avenues
that may improve SCEPTR further. Firstly, training could be made more
efficient by optimising the distribution of masked/dropped tokens during pre-
training, taking into account the variable relevance of different parts of the
sequence in determining specificity [45]. Secondly, as certain sequence motifs
appear recurrently (e.g. CDR3 loops often begin with CAS), a more intelligent
tokenisation scheme could offload learning of these primary sequence
statistics into the tokenisation process.

Pre-trained PLMs have achieved high performance on protein stability and
structural predictions [33, 34]. However, we find that existing PLMs fail to
confer similar benefits to predicting TCR-pMHC interactions. This finding adds
to recent work showing that current PLM pre-training is not well-aligned with
certain downstream tasks [37]. Importantly, we show that autocontrastive pre-
training can overcome misalignment, and thus provide a constructive path out
of this impasse which could also be applied outside of the TCR domain.

What determines whether a certain downstream task is aligned with MLM pre-
training? MLM teaches PLMs to predict the conditional distribution of tokens
given sequence context. Thus it stands to reason that amenable downstream
tasks involve predictions of properties that determine the distribution of
observed proteins on sequence space. Observed proteins tend to concentrate in
areas of sequence space with higher protein stability since evolution on
average selects for this property [56]. For datasets containing protein
families whose members have a conserved structure despite primary sequence
variation, co-evolutionary couplings driven by structural constraints
influence allowed sequence variability [57, 58, 33]. These data distributional
properties might explain how MLM can teach PLMs features related to both
stability and structure.

In contrast, the distribution of TCRs over sequence space is primarily shaped
by the biases of VDJ recombination with antigen-specific selection playing an
important, but likely second-order effect [44]. While long-term evolutionary
pressures may act to align recombination statistics with TCR function [59,
60], empirical evidence so far suggests recombination biases primarily
anticipate thymic selection for stability and folding [61]. In contrast,
studies to date have found no clear relationship between probabilities of
recombination and the likelihood of receptors engaging specific pMHCs [62].
Given these considerations, we expect MLM pre-training to align better to
tasks concerning VDJ recombination than to TCR specificity prediction. Indeed,
previous work training PLMs on adaptive immune receptors has demonstrated that
embeddings strongly depend on V/J gene usage and can be used to predict
primarily generation-related properties such as receptor publicity [24, 38].

Why does autocontrastive learning help to generate embeddings better suited
for specificity prediction? An interesting insight comes from an asymptotic
decomposition of the contrastive loss function into the _uniformity_ and
_alignment_ terms [47]:

$\displaystyle\textrm{Unif.}(f)$
$\displaystyle=\log{\mathop{\mathbb{E}}_{x,y\overset{\mathrm{iid}}{\sim}p_{\textrm{data}}}\left[e^{-\lVert
f(x)-f(y)\rVert^{2}}\right]}$ (2) $\displaystyle\textrm{Align.}(f)$
$\displaystyle\coloneqq\mathop{\mathbb{E}}_{(x,x^{+})\sim
p_{\textrm{pos}}}\left[\lVert f(x)-f(x^{+})\rVert\right]$ (3)

Uniformity incentivises the model to make use of the full representation
space, while alignment minimises the expected distance between positive pairs
(e.g. co-specific TCRs) [47]. From this view, contrastive learning on adaptive
immune receptor data encourages PLMs to undo the large-scale distributional
biases created by VDJ recombination through the uniformity term, while helping
to identify features relating to TCR (co-)specificity via the alignment term.
While _auto_ contrastive learning approximates the alignment term through the
generation of pairs of views, it still provides a direct empirical estimate
for the uniformity term. Thus, a key benefit of autocontrastive learning may
be that it reduces the confounding effects of VDJ recombination in embedding
space.

Figure 6:  Model complexity does not correlate with downstream performance.
Model performance as measured by mean 200-shot AUROC (section I.1) does not
scale with model complexity as measured by either a) parameter count or b)
representation dimensionality. Despite being the smallest PLM by a wide
margin, SCEPTR performs better than alternative models.

Comparing SCEPTR to other PLMs suggests that model complexity as measured by
either parameter count or representation dimensionality is not currently the
limiting factor for TCR-pMHC prediction performance (Fig. 6). This is directly
supported by how the CDR3-only, MLM-only variant of SCEPTR performs almost
equivalently to the much larger but similarly pre-trained TCR-BERT (Fig. 3d).
This finding stands in contrast to observations of performance scaling with
model size in general PLMs [33] and antibody language modelling [39]. While
the focus of the current study was on training a simple model, it would be
interesting in future work to investigate performance scaling with model
complexity and training dataset size with our novel training procedure.

Looking forward, there are many exciting avenues to further develop
contrastive learning as a paradigm to crack the TCR code. For example, there
may be ways to exploit the uniformity (Eq. 2) and alignment (Eq. 3)
decomposition to simultaneously train on unlabelled and specificity-labelled
data. A practical benefit of our contrastive learning formulation is that it
does not require any optimisation with respect to the true negative
distribution (i.e. TCRs that are explicitly not co-specific) – a non-trivial
distribution to estimate for TCRs [63].

Another interesting avenue is the use of labels other than pMHC specificity –
such as phenotypic annotations from single-cell data – as additional
supervised contrastive training signals. Finally, while supervised contrastive
learning does not currently lead to generalisable learning beyond training
pMHCs, we expect a transition towards generalisation as larger volumes of
specificity-labelled TCR data become available, as has been the case with
supervised contrastive learning in other fields [64, 65, 66, 67, 51]. Finally,
while the focus of this work given current data limitations has been on
learning TCR embeddings, contrastive learning may also help us learn effective
joint TCR-pMHC embeddings in the future when the joint space (and particularly
the pMHC space) is better sampled, and thus ultimately enable the zero-shot
prediction of TCR-pMHC specificity.

##  III Methods

###  III.1 Model benchmarking

For each pMHC, we varied the number $k$ of reference TCRs where
$k\in\\{1,2,5,10,20,50,100,200\\}$. Within each model-pMHC-$k$-shot
combination, we benchmarked multiple reference-test splits of the data to
ensure robustness. For $k=1$, we benchmarked every possible split. For
$k\in[2,200]$, we benchmarked 100 random splits, where we ensured that the
same splits were used across all models to reduce extraneous variance.
Variations in model performance across different splits were used to generate
confidence intervals for differences in performance, for example in figure S1.

The CDR3 Levenshtein model computes the distance between two TCRs as the sum
of the Levenshtein distances between the receptors’ $\alpha$ and $\beta$
CDR3s.

Note that while SCEPTR’s architecture and training allows it to directly
generate representation vectors for complete $\alpha\beta$ TCR sequences (Fig.
2a, methods III.2), this is not the case for the other PLMs. For TCR-BERT,
ESM2 and ProtBert representations for the $\alpha$ and $\beta$ chains were
independently generated, then concatenated together, and finally average-
pooled to produce an embedding of the heterodimeric receptor (see appendix A).

###  III.2 SCEPTR architecture

SCEPTR (Simple Contrastive Embedding of the Primary sequence of T cell
Receptors) is a BERT-like transformer encoder that maps TCR sequences to
vector embeddings. Like BERT, it is comprised of a tokeniser module, an
embedder module and a self-attention stack (Fig. 2a).

The tokeniser module represents each input TCR as the amino acid sequences of
the first, second and third complementarity-determining regions (CDRs) of each
chain, where each amino acid is a token. A special <cls> token is appended to
each input TCR, as its contextualised embedding will eventually become
SCEPTR’s output representation vector (Fig. 2a,b).

SCEPTR uses a simple, non-trainable embedder module, where a one-hot vector is
used to encode token identity (22 dimensions for 20 amino acids plus special
tokens <cls> and <mask>), and token positions are specified by first one-hot
encoding the containing CDR loop number (6 dimensions), then encoding the
token’s relative position within the loop as a single scalar variable (Fig.
2b). This results in initial token embeddings in $\mathbb{R}^{29}$, which are
passed through a trainable linear projection onto $\mathbb{R}^{64}$. SCEPTR’s
self-attention stack then operates at this fixed dimensionality (Fig. S11).
SCEPTR’s self-attention stack comprises three layers, each with eight
attention heads and a feed-forward dimensionality of 256, and is thus
substantially simpler than existing models. Our tests suggest that relative
position embedding helps SCEPTR learn better calibrated TCR co-specificity
rules (see appendix B).

###  III.3 SCEPTR Pre-training

####  III.3.1 Data

The unlabelled paired-chain $\alpha\beta$TCR sequences used to pre-train SCEPT
were taken from a study by Tanno et al. [68], which provides 965,523 unique
clonotypes sampled from the blood of 15 healthy human subjects. As opposed to
traditional single-cell sequencing, Tanno et al. used a ligation-based
sequencing methods to resolve which $\alpha$ chains paired with which $\beta$
chains. To mitigate potential noise from incorrect chain pairing, we applied
an extra processing step to remove clonotypes that shared the same nucleotide
sequence for either the $\alpha$ or the $\beta$ chain, as previously described
[44]. After filtering for functional TCRs using tidytcells, a TCR gene symbol
standardiser [69], we retained 842,683 distinct clonotypes.

A random sub-sample of 10% of this data was reserved for use as an unseen test
set, containing 84,268 unique clonotypes distributed across 83,979 unique
TCRs. Of the remaining 90% of the data, we filtered out any clonotypes with
amino acid sequences that also appeared in the test set, resulting in a
training set of 753,838 unique clonotypes across 733,070 unique TCRs.

####  III.3.2 Procedure

SCEPTR was jointly optimised for MLM and autocontrastive learning, where the
total loss of a training step was calculated as the sum of the MLM and
autocontrastive (Eq. 4) losses.

We implemented MLM following established procedures [30]. Namely, 15% of input
tokens were masked, and masked tokens had an 80% probability of being replaced
with the <mask> token, a 10% probability of being replaced by a randomly
chosen amino acid distinct from the original, and a 10% probability of
remaining unchanged. The MLM loss was computed as the cross-entropy between
SCEPTR’s predicted token probability distribution and the ground truth.

Our choice of autocontrastive loss function is inspired by related work in NLP
[51] and computer vision [48], but adapted to the TCR setting. Let
$B=\\{\sigma_{i}\\}_{i=1}^{N}$ be a minibatch of $N$ TCRs. We generate two
independent “views” of each TCR $\sigma_{i}$ by passing two censored-variants
of the same receptor through the model. Our censoring procedure removes a
random subset of a fixed proportion (20%) of the residues from the tokenised
representation of the CDR loops and with a 50% chance drops either the full
$\alpha$ or $\beta$ chain. To ensure that censoring does not fundamentally
alter the underlying TCR sequence, the positional encoding for each token
remains fixed relative to the original TCR. In addition to the random
censoring, views also differ due to dropout noise during independent model
passes. Taken together, this procedure maps the minibatch $B$ to the set of
$2N$ TCR views $V=\\{v_{j}\\}_{j=1}^{2N}$, where $v_{2i}$ and $v_{2i-1}$ are
two independent views of the same TCR $\sigma_{i}$ ($i\in\\{1...N\\}$). Where
$k\in I=\\{1...2N\\}$ is an arbitrary index of a view $v_{k}\in V$, let $p(k)$
be the index of the other view generated from the same TCR, and $N(k)=\\{l\in
I:l\not=k\\}$ be the set of all indices apart from $k$ itself. Let
$\mathbf{r}_{k}$ denote SCEPTR’s vector representation of TCR view $v_{k}$.
Then the autocontrastive loss for minibatch $B$ is computed as follows:

$\mathcal{L}_{\textrm{AC}}(B)=\frac{1}{2N}\sum_{k\in
I}-\log\frac{e^{\mathbf{r}_{k}^{\top}\mathbf{r}_{p(k)}/\tau}}{\sum_{n\in
N(k)}e^{\mathbf{r}_{k}^{\top}\mathbf{r}_{n}/\tau}}$ (4)

Here, $\tau$ is a temperature hyper-parameter which we set to 0.05 during
training, following previous literature [51].

We used ADAM (adaptive moment estimation) [70] to perform stochastic gradient
descent. We chose a minibatch size of 1024 samples and trained for 200 epochs,
which equated to 143,200 training steps. The internal dropout noise of
SCEPTR’s self-attention stack was set to 0.1.

Our methodology of randomly censoring residues and even entire chains stands
in contrast to previous work in NLP by Gao  _et al._ [51], who found that
relying only on the internal random drop-out noise of the language model was
sufficient for effective autoconstrastive learning. However, our experiments
suggest that in the TCR domain, residue and chain censoring leads to
embeddings with better downstream TCR specificity prediction performance (Fig.
S8).

###  III.4 Using SVCs to learn PLM features

To train the linear SVCs on top of PLM features, sampled 1000 random
background TCRs from the training partition of the unlabelled Tanno et al.
dataset. For each PLM-pMHC-split combination, we trained a linear SVC using
the PLM embeddings of the $k$ reference TCRs as the positives and those of the
1000 background TCRs as the negatives. The same 1000 background TCRs were used
across model-pMHC-split combinations to ensure consistency. We accounted for
the imbalance between the number of positive and negative samples used during
SVC fitting by weighting the penalty contributions accordingly. Finally, we
tested SVCs using the same benchmarking classification task as previously
described.

###  III.5 SCEPTR fine-tuning with supervised contrastive learning

####  III.5.1 Data

For supervised contrastive fine-tuning we took all TCR binders against the six
best-sampled pMHC targets from our labelled TCR dataset, and split them into a
training, a validation, and a test set such that no study used to construct
the training or validation sets contributed any TCRs to the test set (Table
SII.

####  III.5.2 Procedure

The fine-tuning process involved the joint optimisation of SCEPTR on MLM and
supervised contrastive learning. As during pre-training, the overall loss for
each training step was computed as the unweighted sum of the MLM and
supervised contrastive (Eq. 5) losses. The pre-trained state of SCEPTR was
used as the starting point for fine-tuning. With only 200 TCRs for each target
pMHC to train on, we limited the number of learnable parameters by only
allowing the weights of the final self-attention layer to be trainable.
Additionally, we monitored increases in validation loss for early stopping of
fine-tuning, which occurred after 2 epochs, where one epoch is defined as the
model seeing 100,000 binders for each pMHC. Given our a batch size of 1,024
TCRs, this corresponded to a total of 1,172 training steps.

Our implementation of supervised contrastive learning closely follows the
formulation suggested by Khosla et al. [48]. This approach to supervised
contrastive learning combines loss contributions from true positive pairs,
with those from second views of each positive instance (as in autocontrastive
learning) as well as all views of all other sample points with the same pMHC
label. Let $B=\\{\sigma_{i}\\}_{i=1}^{N}$ be a minibatch of $N$ pMHC-annotated
TCRs. We use the same procedure as in our autocontrastive framework (see
methods III.3) to generate two views of each of the TCRs, producing a set of
$2N$ views $V=\\{v_{j}\\}_{j=1}^{2N}$. Let $Y=\\{y_{i}\\}_{i=1}^{N}$ be the
index-matched pMHC labels for TCRs in $B$, and $\bar{y}_{j}$ denote the labels
mapped to the indices of the views in $V$ such that
$\bar{y}_{2i}=\bar{y}_{2i-1}=y_{i}$. Now given arbitrary sample view index
$k$, let $P(k)=\\{l\in A(k):\bar{y}_{l}=\bar{y}_{k}\\}$ be the set of all
indices whose corresponding samples have the same pMHC label as $v_{k}$, with
cardinality $|P(k)|$. The supervised contrastive loss for TCR minibatch $B$
is:

$\mathcal{L}_{\textrm{SC}}(B)=\\\ \frac{1}{2N}\sum_{k\in
I}\frac{1}{|P(k)|}\sum_{p\in
P(k)}-\log\frac{e^{\mathbf{r}_{k}^{\top}\mathbf{r}_{p}/\tau}}{\sum_{n\in
N(k)}e^{\mathbf{r}_{k}^{\top}\mathbf{r}_{n}/\tau}}$ (5)

Each batch during fine-tuning has an equally balanced number of binders to
each of the six pMHCs.

###### Acknowledgements.

The authors thank Ned Wingreen, Chris Watkins, Linda Li, Rudy Yuen, Sankalan
Bhattacharyya, and Matthew Cowley for useful discussions. YN and MM were
supported by Cancer Research UK studentships under grants BCCG1C8R and A29287,
respectively. The work of ATM was supported in parts by funding by the Royal
Free Charity.

## Code Availability

https://github.com/yutanagano/sceptr

This is a readily usable deployment of SCEPTR and all its variants seen in
this publication.

https://github.com/yutanagano/tcrlm

This repository houses code used for designing and training our models.

https://github.com/yutanagano/libtcrlm

This repository houses some library code that powers both the sceptr and tcrlm
repositories above.

## References

  * Chi  _et al._ [2024] H. Chi, M. Pepper, and P. G. Thomas, Cell 187, 2052 (2024).
  * Davis and Bjorkman [1988] M. M. Davis and P. J. Bjorkman, Nature 334, 395 (1988).
  * Hudson  _et al._ [2023] D. Hudson, R. A. Fernandes, M. Basham, G. Ogg, and H. Koohy, Nature Reviews Immunology , 1 (2023).
  * Dash  _et al._ [2017] P. Dash, A. J. Fiore-Gartland, T. Hertz, G. C. Wang, S. Sharma, A. Souquette, J. C. Crawford, E. B. Clemens, T. H. O. Nguyen, K. Kedzierska, N. L. La Gruta, P. Bradley, and P. G. Thomas, Nature 547, 89 (2017).
  * Dobson  _et al._ [2022] C. S. Dobson, A. N. Reich, S. Gaglione, B. E. Smith, E. J. Kim, J. Dong, L. Ronsard, V. Okonkwo, D. Lingwood, M. Dougan, _et al._ , Nature methods 19, 449 (2022).
  * Joglekar  _et al._ [2019] A. V. Joglekar, M. T. Leonard, J. D. Jeppson, M. Swift, G. Li, S. Wong, S. Peng, J. M. Zaretsky, J. R. Heath, A. Ribas, _et al._ , Nature methods 16, 191 (2019).
  * Gielis  _et al._ [2019] S. Gielis, P. Moris, W. Bittremieux, N. De Neuter, B. Ogunjimi, K. Laukens, and P. Meysman, Frontiers in Immunology 10 (2019).
  * Fischer  _et al._ [2020] D. S. Fischer, Y. Wu, B. Schubert, and F. J. Theis, Molecular Systems Biology 16, e9416 (2020).
  * Jokinen  _et al._ [2021] E. Jokinen, J. Huuhtanen, S. Mustjoki, M. Heinonen, and H. Lähdesmäki, PLOS Computational Biology 17, e1008814 (2021).
  * Montemurro  _et al._ [2021] A. Montemurro, V. Schuster, H. R. Povlsen, A. K. Bentzen, V. Jurtz, W. D. Chronister, A. Crinklaw, S. R. Hadrup, O. Winther, B. Peters, L. E. Jessen, and M. Nielsen, Communications Biology 4, 1 (2021).
  * Wu  _et al._ [2021] K. Wu, K. E. Yost, B. Daniel, J. A. Belk, Y. Xia, T. Egawa, A. Satpathy, H. Y. Chang, and J. Zou, bioRxiv 10.1101/2021.11.18.469186 (2021).
  * Croce  _et al._ [2024] G. Croce, S. Bobisse, D. L. Moreno, J. Schmidt, P. Guillame, A. Harari, and D. Gfeller, Nature Communications 15, 3211 (2024).
  * Weber  _et al._ [2021] A. Weber, J. Born, and M. Rodriguez Martínez, Bioinformatics 37, i237 (2021).
  * Jiang  _et al._ [2023] Y. Jiang, M. Huo, and S. Cheng Li, Briefings in Bioinformatics 24, bbad086 (2023).
  * Lu  _et al._ [2021] T. Lu, Z. Zhang, J. Zhu, Y. Wang, P. Jiang, X. Xiao, C. Bernatchez, J. V. Heymach, D. L. Gibbons, J. Wang, L. Xu, A. Reuben, and T. Wang, Nature Machine Intelligence 3, 864 (2021).
  * Lin  _et al._ [2021] X. Lin, J. T. George, N. P. Schafer, K. Ng Chau, M. E. Birnbaum, C. Clementi, J. N. Onuchic, and H. Levine, Nature Computational Science 1, 362 (2021).
  * Springer  _et al._ [2021] I. Springer, N. Tickotsky, and Y. Louzoun, Frontiers in Immunology 12, 664514 (2021).
  * Cai  _et al._ [2022] M. Cai, S. Bang, P. Zhang, and H. Lee, Frontiers in Immunology 13 (2022).
  * Moris  _et al._ [2021] P. Moris, J. De Pauw, A. Postovskaya, S. Gielis, N. De Neuter, W. Bittremieux, B. Ogunjimi, K. Laukens, and P. Meysman, Briefings in Bioinformatics 22, bbaa318 (2021).
  * Pham  _et al._ [2023] M.-D. N. Pham, T.-N. Nguyen, L. S. Tran, Q.-T. B. Nguyen, T.-P. H. Nguyen, T. M. Q. Pham, H.-N. Nguyen, H. Giang, M.-D. Phan, and V. Nguyen, Bioinformatics 39, btad284 (2023).
  * Gao  _et al._ [2023] Y. Gao, Y. Gao, Y. Fan, C. Zhu, Z. Wei, C. Zhou, G. Chuai, Q. Chen, H. Zhang, and Q. Liu, Nature Machine Intelligence 5, 236 (2023).
  * Kwee  _et al._ [2023] B. P. Y. Kwee, M. Messemaker, E. Marcus, G. Oliveira, W. Scheper, C. J. Wu, J. Teuwen, and T. N. Schumacher, bioRxiv preprint 10.1101/2023.04.25.538237 (2023).
  * Meynard-Piganeau  _et al._ [2024] B. Meynard-Piganeau, C. Feinauer, M. Weigt, A. M. Walczak, and T. Mora, Proceedings of the National Academy of Sciences 121, 10.1101/2023.07.19.549669 (2024).
  * Goldner Kabeli  _et al._ [2024] R. Goldner Kabeli, S. Zevin, A. Abargel, A. Zilberberg, and S. Efroni, Science Advances 10, eadk4670 (2024).
  * Meysman  _et al._ [2023] P. Meysman, J. Barton, B. Bravi, L. Cohen-Lavi, V. Karnaukhov, E. Lilleskov, A. Montemurro, M. Nielsen, T. Mora, P. Pereira, A. Postovskaya, M. R. Martínez, J. Fernandez-de Cossio-Diaz, A. Vujkovic, A. M. Walczak, A. Weber, R. Yin, A. Eugster, and V. Sharma, ImmunoInformatics 9, 100024 (2023).
  * Grazioli  _et al._ [2022] F. Grazioli, A. Mösch, P. Machart, K. Li, I. Alqassem, T. J. O’Donnell, and M. R. Min, Frontiers in immunology 13, 1014256 (2022).
  * Deng  _et al._ [2023] L. Deng, C. Ly, S. Abdollahi, Y. Zhao, I. Prinz, and S. Bonn, Frontiers in Immunology 14, 1128326 (2023).
  * Bagaev  _et al._ [2020] D. V. Bagaev, R. M. A. Vroomans, J. Samir, U. Stervbo, C. Rius, G. Dolton, A. Greenshields-Watson, M. Attaf, E. S. Egorov, I. V. Zvyagin, N. Babel, D. K. Cole, A. J. Godkin, A. K. Sewell, C. Kesmir, D. M. Chudakov, F. Luciani, and M. Shugay, Nucleic Acids Research 48, D1057 (2020).
  * Vaswani  _et al._ [2017] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin, arXiv preprint 1706.03762 (2017).
  * Devlin  _et al._ [2019] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova, arXiv preprint (2019).
  * Brown  _et al._ [2020] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, _et al._ , Advances in neural information processing systems 33, 1877 (2020).
  * Rives  _et al._ [2021] A. Rives, J. Meier, T. Sercu, S. Goyal, Z. Lin, J. Liu, D. Guo, M. Ott, C. L. Zitnick, J. Ma, _et al._ , Proceedings of the National Academy of Sciences 118, e2016239118 (2021).
  * Lin  _et al._ [2023] Z. Lin, H. Akin, R. Rao, B. Hie, Z. Zhu, W. Lu, N. Smetanin, R. Verkuil, O. Kabeli, Y. Shmueli, A. dos Santos Costa, M. Fazel-Zarandi, T. Sercu, S. Candido, and A. Rives, Science 379, 1123 (2023).
  * Elnaggar  _et al._ [2022] A. Elnaggar, M. Heinzinger, C. Dallago, G. Rehawi, Y. Wang, L. Jones, T. Gibbs, T. Feher, C. Angerer, M. Steinegger, D. Bhowmik, and B. Rost, IEEE Transactions on Pattern Analysis and Machine Intelligence 44, 7112 (2022).
  * Elnaggar  _et al._ [2023] A. Elnaggar, H. Essam, W. Salah-Eldin, W. Moustafa, M. Elkerdawy, C. Rochereau, and B. Rost, arXiv preprint 2301.06568 (2023).
  * Wu  _et al._ [2024] K. E. Wu, H. Chang, and J. Zou, bioRxiv 10.1101/2024.05.14.594226 (2024).
  * Li  _et al._ [2024] F.-Z. Li, A. P. Amini, Y. Yue, K. K. Yang, and A. X. Lu, bioRxiv preprint 10.1101/2024.02.05.578959 (2024).
  * Wang  _et al._ [2024] M. Wang, J. Patsenker, H. Li, Y. Kluger, and S. H. Kleinstein, Nucleic Acids Research 52, 548 (2024).
  * Barton  _et al._ [2024] J. Barton, A. Gaspariunas, D. A. Yadin, J. Dias, F. L. Nice, D. H. Minns, O. Snudden, C. Povall, S. V. Tomas, H. Dobson, J. H. R. Farmery, J. Leem, and J. D. Galson, bioRxiv preprint 10.1101/2024.05.22.594943 (2024).
  * 10x Genomics [2020] 10x Genomics, A New Way of Exploring Immunity - Linking Highly Multiplexed Antigen Recognition to Immune Repertoire and Phenotype (2020).
  * Zhang  _et al._ [2021] W. Zhang, P. G. Hawkins, J. He, N. T. Gupta, J. Liu, G. Choonoo, S. W. Jeong, C. R. Chen, A. Dhanik, M. Dillon, R. Deering, L. E. Macdonald, G. Thurston, and G. S. Atwal, Science Advances 7, eabf5835 (2021).
  * Montemurro  _et al._ [2023] A. Montemurro, H. R. Povlsen, L. E. Jessen, and M. Nielsen, Scientific Reports 13, 16147 (2023).
  * Mayer-Blackwell  _et al._ [2021] K. Mayer-Blackwell, S. Schattgen, L. Cohen-Lavi, J. C. Crawford, A. Souquette, J. A. Gaevert, T. Hertz, P. G. Thomas, P. Bradley, and A. Fiore-Gartland, Elife 10, e68605 (2021).
  * Mayer and Callan [2023] A. Mayer and C. G. Callan, Proceedings of the National Academy of Sciences 120, e2213264120 (2023).
  * Henderson  _et al._ [2024] J. Henderson, Y. Nagano, M. Milighetti, and A. Tiffeau-Mayer, arXiv preprint 10.48550/arXiv.2404.12565 (2024).
  * Tiffeau-Mayer [2023] A. Tiffeau-Mayer, arXiv preprint 10.48550/arXiv.2310.03439 (2023).
  * Wang and Isola [2022] T. Wang and P. Isola, arXiv 10.48550/arXiv.2005.10242 (2022).
  * Khosla  _et al._ [2021] P. Khosla, P. Teterwak, C. Wang, A. Sarna, Y. Tian, P. Isola, A. Maschinot, C. Liu, and D. Krishnan, arXiv 10.48550/arXiv.2004.11362 (2021).
  * Drost  _et al._ [2022] F. Drost, L. Schiefelbein, and B. Schubert, bioRxiv preprint , 2022.10.24.513533 (2022).
  * Pertseva  _et al._ [2024] M. Pertseva, O. Follonier, D. Scarcella, and S. T. Reddy, bioRxiv preprint , 2024.04.04.587695 (2024).
  * Gao  _et al._ [2022] T. Gao, X. Yao, and D. Chen, arXiv 10.48550/arXiv.2104.08821 (2022).
  * Fang  _et al._ [2022] Y. Fang, X. Liu, and H. Liu, Briefings in Bioinformatics 23, bbac378 (2022).
  * Li  _et al._ [2020] B. Li, H. Zhou, J. He, M. Wang, Y. Yang, and L. Li, in  _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_, edited by B. Webber, T. Cohn, Y. He, and Y. Liu (Association for Computational Linguistics, Online, 2020) pp. 9119–9130.
  * Murugan  _et al._ [2012] A. Murugan, T. Mora, A. M. Walczak, and C. G. Callan, Proceedings of the National Academy of Sciences 109, 16161 (2012).
  * Milighetti  _et al._ [2024] M. Milighetti, Y. Nagano, J. Henderson, U. Hershberg, A. Tiffeau-Mayer, A.-F. Bitbol, and B. Chain, bioRxiv , 2024 (2024).
  * Bloom  _et al._ [2006] J. D. Bloom, S. T. Labthavikul, C. R. Otey, and F. H. Arnold, Proceedings of the National Academy of Sciences 103, 5869 (2006).
  * Weigt  _et al._ [2009] M. Weigt, R. A. White, H. Szurmant, J. A. Hoch, and T. Hwa, Proceedings of the National Academy of Sciences 106, 67 (2009).
  * Jumper  _et al._ [2021] J. Jumper, R. Evans, A. Pritzel, T. Green, M. Figurnov, O. Ronneberger, K. Tunyasuvunakool, R. Bates, A. Žídek, A. Potapenko, _et al._ , Nature 596, 583 (2021).
  * Mayer  _et al._ [2015] A. Mayer, V. Balasubramanian, T. Mora, and A. M. Walczak, Proceedings of the National Academy of Sciences 112, 5950 (2015).
  * Thomas and Crawford [2019] P. G. Thomas and J. C. Crawford, Current Opinion in Systems Biology 18, 36 (2019).
  * Elhanati  _et al._ [2014] Y. Elhanati, A. Murugan, C. G. Callan, T. Mora, and A. M. Walczak, Proceedings of the National Academy of Sciences 111, 9875 (2014).
  * Sethna  _et al._ [2019] Z. Sethna, Y. Elhanati, C. G. Callan Jr, A. M. Walczak, and T. Mora, Bioinformatics 35, 2974 (2019).
  * Dens  _et al._ [2023] C. Dens, K. Laukens, W. Bittremieux, and P. Meysman, Nature Machine Intelligence 5, 1060 (2023).
  * Schroff  _et al._ [2015] F. Schroff, D. Kalenichenko, and J. Philbin, in  _2015 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ (2015) pp. 815–823.
  * Liu  _et al._ [2017] W. Liu, Y. Wen, Z. Yu, M. Li, B. Raj, and L. Song, in  _2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ (2017) pp. 6738–6746.
  * Chen  _et al._ [2018] S. Chen, Y. Liu, X. Gao, and Z. Han, arXiv 10.48550/arXiv.1804.07573 (2018).
  * Deng  _et al._ [2021] J. Deng, J. Guo, J. Yang, N. Xue, I. Kotsia, and S. Zafeiriou, IEEE Transactions on Pattern Analysis and Machine Intelligence , 1 (2021).
  * Tanno  _et al._ [2020] H. Tanno, T. M. Gould, J. R. McDaniel, W. Cao, Y. Tanno, R. E. Durrett, D. Park, S. J. Cate, W. H. Hildebrand, C. L. Dekker, L. Tian, C. M. Weyand, G. Georgiou, and J. J. Goronzy, Proceedings of the National Academy of Sciences 117, 532 (2020).
  * Nagano and Chain [2023] Y. Nagano and B. Chain, Frontiers in Immunology 14, 10.3389/fimmu.2023.1276106 (2023).
  * Kingma and Ba [2017] D. P. Kingma and J. Ba, arXiv preprint 10.48550/arXiv.1412.6980 (2017).
  * Heather  _et al._ [2022] J. M. Heather, M. J. Spindler, M. H. Alonso, Y. I. Shui, D. G. Millar, D. S. Johnson, M. Cobbold, and A. N. Hata, Nucleic Acids Research 50, e68 (2022).

##  Appendix A Generating TCR vector embeddings using existing protein
language models

###  A.1 TCR-BERT

The TCR-BERT model was downloaded through HuggingFace at
https://huggingface.co/wukevin/tcr-bert. Since TCR-BERT is trained to read one
CDR3 sequence at a time, we generated TCR representations by generating two
independent representations of the $\alpha$ and $\beta$ chain, and
concatenating them together. The TCR-BERT representation of a chain was
generated by feeding the model its CDR3 sequence, then taking the average pool
of the amino acid token embeddings in the 8th self-attention layer, as
recommended by the study authors [11].

###  A.2 ESM2

The ESM2 (T6 8M) model was downloaded through HuggingFace at
https://huggingface.co/facebook/esm2_t6_8M_UR50D. ESM2 is trained on full
protein sequences, but not protein multimers. Therefore, we generated ESM2
representations for the $\alpha$ and $\beta$ chains separately, and
concatenated them to produce the overall TCR representation. To generate the
representation of a TCR chain, we first used Stitchr [71] to reconstruct the
full amino acid sequence of a TCR from its CDR3 sequence and V/J gene. Then,
the resulting sequence of each full chain was fed to ESM2. We took the
average-pooled result of the amino acid token embeddings of the final layer to
generate the overall sequence representation, as recommended [33].

###  A.3 ProtBert

The ProtBert model was downloaded through HuggingFace at
https://huggingface.co/Rostlab/prot_bert. Similarly to ESM2, ProtBert is
trained on full protein sequences. Therefore, we again used Stitchr to
generate full TCR chain amino acid sequences, and fed them to ProtBert to
generate independent $\alpha$ and $\beta$ chain representations. We again as
recommended average-pooled the amino acid token embeddings of the final layer
[34].

##  Appendix B Effects of different position embedding methods

To better understand TCR similarity rules as learned by PLMs, we measured the
average distance penalty incurred within a model’s representation space as a
result of a single amino acid edit at various points along the length of the
$\alpha$/$\beta$ CDR3 loops. To do this, we randomly sampled real TCRs from
the testing partition of the Tanno et al. dataset [68] and synthetically
introduced single residue edits in one of their CDR3 loops. Then, we measured
the distance between the original TCR and the single edit variant according to
a PLM. For each model, we sampled TCRs until we had observed at least 100
cases of: 1) each type of edit (insertions, deletions, substitutions) at each
position, and 2) substitutions from each amino acid to every other. Since CDR3
sequences vary in length, we categorised the edit locations into one of five
bins: C-TERM for edits within the first one-fifth of the CDR3 sequence
counting from the C-terminus, then M1, M2, M3, and N-TERM, in that order. For
this analysis, we investigated SCEPTR and TCR-BERT, since they are the two
best performers out of the PLMs tested (Fig. 1d).

Both SCEPTR and TCR-BERT generally associate insertions and deletions (indels)
with a higher distance penalty compared to substitutions (Fig. S9a). While
SCEPTR uniformly penalises indels across the length of the CDR3, TCR-BERT
assigns higher penalties to those closer to the C-terminus. We hypothesised
that the variation in TCR-BERT’s indel penalties is a side-effect of its
position embedding system. TCR-BERT, like many other transformers, encodes a
token’s position into its initial embedding in a left-aligned manner using a
stack of sinusoidal functions with varying periods [11, 29, 30] This results
in embeddings that are more sensitive to indels near the C-terminus, which
cause a frame-shift in a larger portion of the CDR3 loop and thus lead to a
larger change in the model’s underlying TCR representation. To test this
hypothesis, we trained and evaluated a new SCEPTR variant:

SCEPTR (left-aligned)

This variant uses a traditional transformer embedding system with trainable
token representations and left aligned, stacked sinusoidal position
embeddings.

While we detect no significant difference in downstream performance between
SCEPTR and its left-aligned variant (Fig. S10), this may be because cases
where the differences in their learned rule sets affects performance are
rarely seen in our benchmarking data. The edit penalty profile of the left-
aligned variant shows a similar falloff of indel penalties than TCR-BERT with
higher penalties at the C than N-terminals (Fig. S9b). As their is no clear
biological rationale for this observation, these results suggest that SCEPTR’s
relative position encoding might result in a better-calibrated co-specificity
ruleset. These preliminary findings add to the ongoing discussion around how
to best encode residue position information in the protein language modelling
domain [35].

Interestingly, the penalty falloff seen with SCEPTR (left-aligned) is sharper
than that of TCR-BERT, whose indel penalties plateau past M1. As TCR-BERT is a
substantially deeper model (12 self-attention layers, 12 heads each, embedding
dimensionality 768), it might be partially able to internally un-learn the
left-aligned-ness of the position information. If this is true, then position
embedding choices are particularly important for training smaller, more
efficient models.

##  Appendix C Supplementary Figures/Tables

Figure S1:  Statistical analysis of benchmark performance differences with respect to SCEPTR. This is an accompanying plot to the benchmarking results shown in figure 1. The number of reference sequences varies along the x axis. The y axis shows the difference between the mean AUROC ($\Delta$AUROC) of SCEPTR and other models. The error bars represent standard deviations, which were calculated across pMHCs and data splits.  Figure S2:  Benchmarking PLM embeddings on TCR specificity prediction with Montemurro et al.’s post-processed 10xGenomics dataset included. This is a repeat of the benchmarking study from section I.1 with a larger dataset of labelled TCRs. This includes six more sufficiently sampled pMHC specificities (with epitopes ELAGIGILTV, GLCTLVAML, AVFDRKSDAK, IVTDFSVIK, RAKFKQLL, KLGGALQAK). On both subplots, the number of reference sequences varies along the x axis. a) The y axis shows the models’ AUROCs averaged across pMHCs. b) The y axis shows the distribution of $\Delta$AUROC between SCEPTR and the other models (see Fig. S1). The trends seen in section I.1 are recapitulated.  Figure S3:  Nearest neighbour prediction is more performant than using the average distance to all references. We repeated the benchmarking procedure used to produce figure 1d, but instead of making inferences based on the distance between a query TCR and its closest reference neighbour, we averaged its distance to all references. The results are shown in panel b, with panel a showing the original nearest neighbour prediction-based benchmarking results for comparison. The x axes show the number of reference sequences given to the model, while the y axes show average AUROC. SCEPTR shows state-of-the-art performance in both cases. All models perform better when applied through nearest neighbour prediction.  Figure S4:  Statistical analysis of benchmark performance differences from the ablation studies. This is an accompanying figure to the ablation study results shown in figure 3. The number of reference sequences varies along the x axis. The y axis shows the distributions of $\Delta$AUROC to SCEPTR (see Fig. S1).  Figure S5:  Benchmarking fine-tuned SCEPTR on TCR specificity prediction for unseen pMHCs. This figure shows the result of benchmarking fine-tuned SCEPTR (see section I.5) against TCRdist, TCR-BERT and the baseline SCEPTR model on pMHC targets unseen during SCEPTR’s fine-tuning. Here, we use the benchmarking framework from section I.1, with the prediction targets being all pMHCs other than the six seen during fine-tuning that have at least 120 known binders. The number of reference sequences $k$ varies along the x axis. Here, we benchmark $k\in[1,20]$ as we want to always have at least 100 positive test cases for each data split. a) The y axis shows the models’ AUROCs averaged across pMHCs. b) The y axis shows the distribution of $\Delta$AUROC (see Fig. S1) between the baseline SCEPTR model and the others. Fine-tuned SCEPTR performs significantly worse compared to the baseline model.  Figure S6:  Benchmarking linear support vector classifiers trained on PLM features on TCR specificity prediction. These plots compare the performances of different PLMs applied to few-shot TCR specificity prediction, either through nearest neighbour prediction (models marked as “NN” in the legend, see section I.1) or using a linear support vector classifier trained atop their TCR featurisations (models marked as “SVC” in the legend, see section I.4, methods III.4). Benchmarking was done using the framework outlined in section I.1. In both plots, the number of reference sequences varies along the x axis. a) The y axis shows the models’ AUROCs averaged across pMHCs. b) The y axis shows the distribution of $\Delta$AUROCs (see Fig. S4) between SCEPTR (NN) and the other models. For all PLMs except SCEPTR, training a linear SVC atop the model’s features improves performance. SCEPTR (NN) outperforms all methods, even the SVC trained atop its own features.  Figure S7:  Comparing nearest neighbour and average distance implementations of the baseline and fine-tuned SCEPTR models. This is an accompanying plot to figure 5. Here, we investigate the TCR specificity prediction performance of the baseline and fine-tuned versions of SCEPTR, implemented via the nearest neighbour (NN) or average distance (Avg dist) prediction frameworks (see section I.1). For the baseline model, we see that the nearest neighbour implementation performs better on average, consistent with the results seen in figure S3. In contrast, the average distance implementation of the fine-tuned model greatly outperforms its nearest neighbour counterpart. Our primary hypothesis as to why most models including the baseline SCEPTR model perform better through nearest neighbour prediction (Fig. S3) is that each pMHC has multiple viable binding solutions comprised of TCRs with different primary sequence features, which means that averaging distance to all reference TCRs across binding solutions dilutes signal. The fact that the fine-tuned model no longer shows this property may hint at its ability to better resolve these distinct binding solutions into a single convex cluster.  Figure S8:  Residue and chain dropping provide an improved noise model during autocontrastive learning. This plot summarises benchmarking results comparing SCEPTR with a variant (SCEPTR (dropout noise only)) which does not employ any extra noising operations when producing the two views of the same TCR during autocontrastive learning, solely relying on SCEPTR’s internal dropout noise (see section III.3). The benchmarking framework as outlined in section I.1 is used. The number of reference sequences varies along the x axis. a) The y axis shows the models’ AUROCs averaged across pMHCs. b) The y axis shows the distribution of $\Delta$AUROCs (see Fig. S1) from the dropout noise only variant to the baseline model. While overall performance is similar, the addition of out residue- and chain- dropping noise model improves downstream performance.  Figure S9:  Investigating TCR co-specificity rules as learned by different PLMs. Here we investigate TCR co-specificity rules as learned by various representation models by measuring the expected distance penalties inucrred by single residue edits in different regions of the $\alpha$ and $\beta$ CDR3 loops. We investigate three models: SCEPTR, TCR-BERT, and a SCEPTR variant which replaces its simplified initial embedding module with one that emulates the traditional transformer architecture, including a left-aligned position embedding system (see appendix B). The x axis shows different regions of the CDR3 divided into five bins, where C-TERM represents the first fifth of the loop counting from the C-terminal, N-TERM represents the last fifth of the loop on the N-terminal end, and the middle regions numbered from the C-terminal as shown. The y axis hows the expected distance penalty incurred by different types of single edits. The different lines show the expected penalty curves with respect to insertions (purple), deletions (orange) and substitutions (green). The error bars show the standard deviations. According to all models, substitutions on average incur a smaller distance penalty compared to indels. While SCEPTR uniformly penalises indels, both TCR-BERT and the left-aligned SCEPTR variant assign higher distance penalties to indels closer to the C-terminal.  Figure S10:  SCEPTR with its simplified embedder module performs similarly to a variant with an embedder module emulating the traditional transformer architecture. Here we show the results of benchmarking SCEPTR againt a variant which replaces SCEPTR’s simplified embedder module (see methods III.2) with an implementation emulating the traditional transformer architecture (“left-aligned” variant in plot, see appendix B). The benchmarking framework as outlined in section I.1 is used. The number of reference sequences varies along the x axis. a) The y axis shows the models’ AUROCs averaged across pMHCs. b) The y axis shows the distribution of $\Delta$AUROCs (see Fig. S1) from the left-aligned variant to the baseline model. We detect no significant difference in performance between the two models.  Figure S11:  A simplified schematic depicting the internals of the transformer self-attention stack. The schematic is accurate for the case of a single attention head per layer. In the more general case of $H$ attention heads, each MHA block will have $H$ parallel q, k and v linear projections, each from dimensionality $d$ to dimensionality $d/H$. Each parallel set of q, k and v vectors/matrices undergo the series of operations shown in the schematic. Finally, the final output vector (of shape $d/H\times 1$) from each parallel branch are concatenated together to produce the output of the MHA block (of shape $d\times 1$).  Table SI:  Per-epitope summary of the different models’ perfomances from the nearest-neighbour prediction benchmarking (see section I.1) with the number of reference TCRs k=200. The best AUROC per epitope is shown in bold.  | SCEPTR | TCRdist | CDR3 Levenshtein | TCR-BERT | ESM2 (T6 8M) | ProtBert
---|---|---|---|---|---|---
epitope |  |  |  |  |  |
GILGFVFTL | 0.911 | 0.904 | 0.872 | 0.876 | 0.831 | 0.845
NLVPMVATV | 0.691 | 0.648 | 0.632 | 0.655 | 0.652 | 0.629
SPRWYFYYL | 0.728 | 0.695 | 0.610 | 0.637 | 0.604 | 0.575
TFEYVSQPFLMDLE | 0.976 | 0.970 | 0.964 | 0.966 | 0.937 | 0.950
TTDPSFLGRY | 0.708 | 0.720 | 0.600 | 0.576 | 0.579 | 0.564
YLQPRTFLL | 0.775 | 0.762 | 0.743 | 0.698 | 0.697 | 0.669
Table SII: The different studies that contributed TCR data to the training/validation and test splits for the supervised contrastive learning fine-tuning task. For datasets without a PubMed ID the table indicates the VDJdb github issue number corresponding to the dataset inclusion. epitope |  training/validation |  test
---|---|---
GILGFVFTL |  PMID:28636592 |  PMID:12796775, PMID:18275829, PMID:28250417, PMID:28931605, PMID:7807026, PMID:28423320, PMID:28636589, PMID:27645996, PMID:29483513, PMID:29997621, PMID:34793243, VDJdbID:215
NLVPMVATV |  PMID:28636592, VDJdbID:332 |  PMID:19542454, PMID:26429912, PMID:19864595, PMID:28423320, PMID:16237109, PMID:28636589, PMID:36711524, PMID:28623251, PMID:9971792, PMID:17709536, PMID:28934479, PMID:34793243, VDJdbID:252
SPRWYFYYL |  PMID:33951417, PMID:35750048 |  PMID:33945786, PMID:34793243
TFEYVSQPFLMDLE |  PMID:35750048 |  PMID:37030296
TTDPSFLGRY |  PMID:35383307 |  PMID:35750048
YLQPRTFLL |  PMID:35383307, PMID:34793243 |  PMID:34685626, PMID:37030296, PMID:33664060, PMID:33951417, PMID:35750048, VDJdbID:215