dkarthikeyan1/tcrt5_ft_tcrdb

TCRT5 model (finetuned)

Model description

TCRT5 is a seq2seq model designed to for the conditional generation of T-cell receptor (TCR) sequences given a target peptide-MHC (pMHC). It is a transformers model that is built on the T5 architecture operationalized by the associated HuggingFace abstraction. It is released along with this paper.

Intended uses & limitations

This model is designed for auto-regressively generating CDR3 $\beta \backslash beta$ sequences against a pMHC of interest. This means that the model assumes a plausible pMHC is provided as input. We have not tested the model on peptides and MHC sequences where the binding affinity between petpide-MHC is low and do not expect the model will adjust its predictions around this. This model is intended for academic purposes and should not be used in a clinical setting.

How to use

You can use this model directly for conditional CDR3 $\beta \backslash beta$ generation:

import re
from transformers import T5Tokenizer, T5ForConditionalGeneration
tokenizer = T5Tokenizer.from_pretrained('dkarthikeyan1/tcrt5_ft_tcrdb')
tcrt5 = T5ForConditionalGeneration.from_pretrained("dkarthikeyan1/tcrt5_ft_tcrdb")
pmhc = "[PMHC]KLGGALQAK[SEP]YFAMYQENVAQTDVDTLYIIYRDYTWAELAYTWY[EOS]"
encoded_pmhc = tokenizer(pmhc, return_tensors='pt')

# Define the number of TCRs you would like to generate ()
num_tcrs = 10
# Define the number of beams to explore (recommended: 3x the number of TCRs)
num_beams = 30

outputs = tcrt5.generate(**encoded_pmhc, max_new_tokens=25, num_return_sequences=num_tcrs, num_beams=num_beams)

# Use regex to get out the [TCR] tag
cdr3b_sequences = [re.sub(r'\[.*\]', '', x) for x in tokenizer.batch_decode(outputs, skip_special_tokens=True)]

>>> cdr3b_sequences

['CASSLGTGGTDTQYF',
 'CASSPGTGGTDTQYF',
 'CASSLGQGGTEAFF',
 'CASSVGTGGTDTQYF',
 'CASSLGTGGSYEQYF',
 'CASSPGQGGTEAFF',
 'CASSSGTGGTDTQYF',
 'CASSLGGGGTDTQYF',
 'CASSLGGGSYEQYF',
 'CASSLGTGGNQPQHF']

This model can also be used for unconditional generation of CDR3 $\beta \backslash beta$ sequences:

import re
from transformers import T5Tokenizer, T5ForConditionalGeneration
tokenizer = T5Tokenizer.from_pretrained('dkarthikeyan1/tcrt5_ft_tcrdb')
tcrt5 = T5ForConditionalGeneration.from_pretrained("dkarthikeyan1/tcrt5_ft_tcrdb")


# Define the number of TCRs you would like to generate ()
num_tcrs = 10
# Define the number of beams to explore (recommended: 3x the number of TCRs)
num_beams = 30

unconditional_outputs = tcrt5.generate(max_new_tokens=25, num_return_sequences=num_tcrs, num_beams=num_beams)

# Use regex to get out the [TCR] tag
uncond_cdr3b_sequences = [re.sub(r'\[.*\]', '', x) for x in tokenizer.batch_decode(unconditional_outputs, skip_special_tokens=True)]

>>> uncond_cdr3b_sequences

['CASSLGGETQYF',
 'CASSLGQGNTEAFF',
 'CASSLGQGNTGELFF',
 'CASSLGTSGTDTQYF',
 'CASSLGLAGSYNEQFF',
 'CASSLGLAGTDTQYF',
 'CASSLGQGYEQYF',
 'CASSLGLAGGNTGELFF',
 'CASSLGGTGELFF',
 'CASSLGQGAYEQYF']

Note: For conditional generation, we found that the model performance was greatest using beam search decoding. However, we also report a reduction in sequence diversity using this particular decoding method. If you would like to generate more diverse sequence, TCRT5 supports a range of alternative decoding strategies which can be found here and here.

Limitations and bias

One of the known biases of TCRT5's predictions is its preference for sampling high V(D)J recombination probability sequences as computed by OLGA. This can be attenuated with the use of alternative decoding methods such as ancestral sampling.

Training data

TCRT5 was pre-trained on masked span reconstruction of ~14M TCR sequences from TCRdb as well as ~780k peptide-pseudosequence pairs taken from IEDB. Finetuning was done using a parallel corpus of ~330k TCR:peptide-pseudosequence pairs taken from VDJdb, IEDB, McPAS, and semi-synthetic examples from MIRA.

Training procedure

Preprocessing

All amino acid sequences, and V/J gene names were standardized using the tidytcells package. See here. MHC allele information was standardized using mhcgnomes, available here before mapping allele information to the MHC pseudo-sequence as defined in NetMHCpan.

Pre-training

TCRT5 was pretrained with Masked language modeling (MLM): Span reconstruction similar to the original training loss of the T5 paper. For a given sequence, the model masks 15% of the sequence using contiguous spans of random length from length 1-3. This is done via the sentinel tokens introduced in the T5 paper. Then the entire masked sequence is passed into the model and the model is trained to reconstruct a concatenated sequence comprised of the sentinel tokens followed by the masked tokens. This forces the model to learn richer k-mer dependencies of the masked sequences.

Masks 'mlm_probability' tokens grouped into spans of size 'max_span_length' according to the following algorithm:
        * Radnomly generate span lengths that add up to round(mlm_probability*seq_len) (ignoring pad token) for each sequence.
        * Ensure that the spans are not directly adjacent to ensure max_span_length is observed
        * Once the span masks are generated according to T5 standards mask the inputs and generate the targets 
    
    
    Example Input:
    
    CASSLGQGYEQYF
    
    Masked Input:
    
    CASSLG[X]GY[Y]F
    
    Target:
    
    [X]Q[Y]EQY[Z].

Finetuning

TCRT5 was finetuned on peptide-pseudo sequence -> CDR3 $\beta \backslash beta$ source:target pairs using the canonical cross entropy loss.

    Example Input:
    
    [PMHC]KLGGALQAK[SEP]YFAMYQENVAQTDVDTLYIIYRDYTWAELAYTWY[EOS]
    
    
    Target:
 
    [TCR]CASSLGYNEQFF[EOS].

Results

This fine-tuned model achieves the following results on conditional CDR3 $\beta \backslash beta$ generation on our validation set of the top-20 peptide-MHCs with the most abundant known TCRs (in alphabetical order):

1. AVFDRKSDAK_A*11:01
2. CRVRLCCYVL_C*07:02
3. EAAGIGILTV_A*02:01
4. ELAGIGILTV_A*02:01
5. GILGFVFTL_A*02:01
6. GLCTLVAML_A*02:01
7. IVTDFSVIK_A*11:01
8. KLGGALQAK_A*03:01
9. LLLDRLNQL_A*02:01
10. LLWNGPMAV_A*02:01
11. LPRRSGAAGA_B*07:02
12. LVVDFSQFSR_A*11:01
13. NLVPMVATV_A*02:01
14. RAKFKQLL_B*08:01
15. SPRWYFYYL_B*07:02
16. STLPETAAVRR_A*11:01
17. TPRVTGGGAM_B*07:02
18. TTDPSFLGRY_A*01:01
19. YLQPRTFLL_A*02:01
20. YVLDHLIVV_A*02:01

Benchmark results:

Metric	Char-BLEU	F@100	SeqRec%	Diversity (num_seq)	Ave. Jaccard Dissimilarity	Perplexity
	96.4	.09	89.2	1300 (2000 max)	94.4/100	2.48

BibTeX entry and citation info

@article{dkarthikeyan2024tcrtranslate,
  title={TCR-TRANSLATE: Conditional Generation of Real Antigen Specific T-cell Receptor Sequences},
  author={Dhuvarakesh Karthikeyan and Colin Raffel and Benjamin Vincent and Alex Rubinsteyn},
  journal={bioArXiv},
  year={2024},
}