CLEST-TCRdist
Contrastive LEarning with Soft Targets from TCRdist. Checkpoint
SCEPTR6L_AC_soft_800k_20ep_bs1024.
A paired-chain αβ T-cell receptor encoder. It maps a TCR to a 128-dimensional vector in which receptors that recognise the same peptide sit close together, so retrieval and clustering of specificity groups reduce to nearest-neighbour search. The model reads only the receptor, never the epitope.
Install
pip install git+https://github.com/20eddibae/CLEST-TCRDist.git
pip install --no-deps sceptr==1.2.0
sceptr goes in with --no-deps because its metadata asks for numpy 2 while tcrdist3 needs
1.26. Everything it actually imports is already installed by the line above.
Use it
import pandas as pd
from clest_tcrdist import load_pretrained
model = load_pretrained() # pulls ~5 MB here, then caches
tcrs = pd.DataFrame({
"TRAV": ["TRAV35*01"], "CDR3A": ["CAGQLGGGSQGNLIF"], "TRAJ": ["TRAJ42*01"],
"TRBV": ["TRBV19*01"], "CDR3B": ["CASSIRSSYEQYF"], "TRBJ": ["TRBJ2-7*01"],
})
vectors = model.calc_vector_representations(tcrs) # (N, 128)
distances = model.calc_cdist_matrix(tcrs, tcrs) # (N, N)
The returned object is a sceptr.model.Sceptr, so anything written against the SCEPTR API
works unchanged. Columns are IMGT symbols and amino-acid junctions. A missing chain or gene
can be left empty, and single-chain input is accepted.
Few-shot retrieval
Given a few known binders for one epitope, rank a repertoire:
from clest_tcrdist import read_tcrs, retrieve
repertoire = read_tcrs("repertoire.csv") # any common column naming
seeds = read_tcrs("known_binders.csv")
ranked = retrieve(repertoire, seeds, model=model) # sorted best first
or from the shell:
python -m clest_tcrdist.retrieve repertoire.csv --seeds known_binders.csv -o ranked.csv
Scoring is distance to the mean seed embedding by default, or to the nearest single seed
with --scoring nearest. Averaging over the seeds is worth about 0.05 AUROC to this model
and is what the headline number below uses.
Clustering
python -m clest_tcrdist.cluster repertoire.csv -o clustered.csv
Pairs within a radius are linked and connected components become clusters, the same
construction GIANA and TCRdist use, with the radius measured in this embedding. The radius
is calibrated from the data unless you set --threshold, and it controls a purity against
retention tradeoff: on 120 TCRs across four epitopes, purity is 1.00 out to 20% retention,
then 0.88 at 30% and 0.74 at 40%. --sweep prints that curve for a labelled file. Compare
clustering methods at matched retention, never at matched radius.
Single-chain (beta-only) mode
Most repertoires are bulk TCRb, and GIANA and TCRdist are usually run that way, so every
entry point takes chain="beta" / --chain beta:
from clest_tcrdist.cluster import embed
vectors = embed(model, repertoire, chain="beta") # TRBV, CDR3B, TRBJ only
python -m clest_tcrdist.retrieve repertoire.csv --seeds binders.csv --chain beta
python -m clest_tcrdist.cluster repertoire.csv --chain beta
Same weights, not a second checkpoint: pretraining dropped whole chains at random, so a lone beta chain is an input the model saw in training. A file with no alpha columns already encodes this way; the flag lets you withhold an alpha chain you do have, which is what makes a comparison against a beta-only method fair.
On the six-epitope VDJdb backtest this reads AUROC@50 0.843 paired, 0.706 beta-only, against GIANA 0.679, iSMART 0.687 and TCRdist restricted to beta 0.668. One caveat worth knowing before you pick a checkpoint: the plain autocontrastive control scores 0.732 in beta-only mode, so the soft-label objective's paired-chain advantage does not survive dropping the alpha chain, most likely because its TCRdist teacher is computed over all six CDR loops. For beta-only work treat the two as comparable.
How it scores
Few-shot retrieval on the six-epitope VDJdb backtest of the SCEPTR paper, 100 resamplings of the reference set, mean AUROC over epitopes, nearest-neighbour scoring:
| Reference TCRs | CLEST-TCRdist | TCRdist | SCEPTR (released) | CDR3 Levenshtein |
|---|---|---|---|---|
| 1 | 0.701 | 0.684 | 0.643 | 0.621 |
| 10 | 0.765 | 0.756 | 0.710 | 0.687 |
| 50 | 0.794 | 0.779 | 0.737 | 0.699 |
| 200 | 0.814 | 0.802 | 0.754 | 0.708 |
Read this before quoting a number
These weights are one training run, not an average. The checkpoint predates seeding in our pipeline. Five seeded replicates of the identical configuration average 0.7750 ± 0.0076 at nearest-neighbour and 0.8104 ± 0.0053 under average-distance scoring, where this run gives 0.7945 and 0.8423. It sits above all five in every cell, 2.5 seed standard deviations out at nearest-neighbour and 6.0 under averaging. Treat the replicate means as what the recipe gives you, and this checkpoint as a lucky draw from it. All five replicates are in the paper's supplement.
What the objective is worth is measured against its own control, not against TCRdist. With the full set of replicates the soft-label target buys +0.039 AUROC@50 at nearest-neighbour over the matched autocontrastive control (paired over six epitopes, p = 0.043, positive on five) and +0.048 under averaging (p = 0.119). Across estimators the range is +0.038 to +0.056. That single-variable comparison is the contribution.
It does not reliably beat TCRdist, and we do not claim it does. The margin in the table above does not survive two controls: reselecting hyperparameters on held-out epitopes costs about 0.013, and giving TCRdist the read-out from its own publication rather than the 1-NN rule inherited here closes the rest. Over seeds the sign of the 1-NN margin changes. What holds up is the gap over the control.
It is not an epitope-binding classifier. There is no peptide or MHC input. It tells you which receptors resemble each other in a space shaped by specificity; it cannot tell you that a receptor binds a peptide you have not given it examples of.
Training
- 6-layer Pre-LN transformer with NormFormer norms, d_model 128, 8 heads, dropout 0.1,
CdrSimpleEmbedder(29-dim token embedding), CLS readout, 1,203,860 parameters. - 827,802 unlabeled paired αβ TCRs from Tanno et al. 2022 (Zenodo 14288119). No epitope labels anywhere in pretraining.
- 20 epochs, batch 1024, Noam schedule with 2000 warmup steps, contrastive temperature 0.05, masked-token auxiliary loss with 0.1 label smoothing.
- Shipped epoch is the one with the lowest validation contrastive loss, epoch 20. The
repository contains
config.jsonandlog.csvfor the whole run.
Also here
- Fine-tuned variant, for when you have epitope labels to train on.
- Benchmark data, the VDJdb snapshot these numbers come from.
- Code, including the training configs and the benchmark that produced the table.
Citation
The manuscript is in preparation; cite the repository until it appears. Please also cite SCEPTR (Nagano et al., Cell Systems 2025), TCRdist (Dash et al., Nature 2017), the Tanno et al. 2022 paired repertoire, and VDJdb (Bagaev et al. 2020).
MIT licensed.
- Downloads last month
- 20