Tranception: protein fitness prediction with autoregressive transformers and inference-time retrieval

In plain words

Biologists want to know which changes to a protein's amino-acid sequence will break it. Testing every change in a lab is slow, so software tries to predict the effect. Older tools look at many related proteins from different species, line them up, and see which positions vary. That works well only when enough related sequences exist and they line up cleanly, which is not true for many proteins.

Tranception is a language model for protein sequences. It learns to guess the next amino acid from the ones before it, trained on a very large set of unaligned sequences. At prediction time it can also look up related sequences and lean on the amino acids seen at each position. Combining the two gives a score for each mutant. The authors also build ProteinGym, a collection of measured mutation experiments used to compare such tools.

TL;DR

• A decoder-only transformer trained on protein sequences, with a convolutional attention variant and a grouped distance bias §3.1 §3.2.
• At inference it fuses its own predictions with a retrieved alignment profile, with weight 0.6 §4.2 app. B.5.
• On the ProteinGym substitution benchmark it is only slightly ahead of EVE overall, but clearly ahead on shallow alignments and multiple mutants §6.2.
• Use it when a protein has few homologs, when you need indels, or when mutants carry several changes. Check the overall margin before assuming it wins.

Problem

Most fitness predictors rely on multiple sequence alignments of homologous proteins, because those alignments carry homology, phylogeny and structure information §4.1. Such alignments need to be deep and diverse. Many protein families cannot be aligned that way, which limits the scope of these methods §4.1 §6.2. Alignment-based models also struggle with multiple mutants and cannot score insertions or deletions, because they depend on the fixed coordinate system of their training alignment app. D.

Protein language models trained on unaligned sequences avoid the alignment requirement, but earlier ones were weaker than alignment methods on shallow families §6.2. The paper also needs a benchmark that covers more protein families, mutation depths and indels than earlier ones §5 app. F.

Method

Architecture. Tranception is a decoder-only transformer. Its only prediction target is the next residue, so it is trained with a left-to-right likelihood over the 20 standard amino acids §3.4 eq. 1 §3.3.

Tranception attention §3.1. Each layer splits its attention heads into four groups. The query, key and value projections of each group pass through a depthwise convolution. The first group has no convolution, and the other three use different kernel sizes. This is meant to make heads specialise on patterns of different length. The paper uses squared ReLU activations, as in Primer, and the grouping idea, as in Inception. The kernel sizes are not given in the text; see Paper vs code.

Grouped ALiBi §3.2. Learned or sinusoidal position embeddings are removed. Instead, each head's query-key score gets a term that grows with distance, as in ALiBi, but the slopes are applied per head group so that each group can attend at a different range.

Training data §3.3 app. B.2. The final model is trained on UniRef100, after mild filtering. The paper reports about 249 million training sequences and a 1% validation split. The context is 1,024 tokens, which covers about 98% of the sequences without truncation. Longer sequences are cut to a random contiguous slice during training. Indeterminate residues are imputed at random during training and inference.

Mirror augmentation §3.3 app. B.4. Each training sequence is reversed with probability 0.5. At inference, the score of each sequence is averaged over its two directions §3.4.

Scoring a mutant §3.4 eq. 2. The fitness of a mutant is the log-likelihood ratio between the mutant and the wild type:

F_x = \log \frac{P(x^{mut})}{P(x^{wt})} eq. 2

Sequences longer than the context are cut to a window. The paper picks the window that gives the most context around the mutation, or around the mean position for multiple mutants app. B.4.

Retrieval inference §4.1 §4.2 app. B.5. For each wild-type protein, a multiple sequence alignment is retrieved once and reused for every mutant of that protein. The retrieved sequences are weighted with the scheme of Hopf et al. §4.2. At each aligned column, the amino-acid distribution is computed from the weighted sequences with a small pseudocount. For substitutions the same alignment is used for wild type and mutant. For indels, the alignment columns are edited to match the mutant. Columns that are inserted in the mutant get no retrieval term.

Fusion §4.2 eq. 3 eq. 4. The autoregressive log-probability is combined with the retrieval log-probability at each covered position:

\log P(x) = \frac{1}{C}\left[(1-\alpha)\log P_A(x) + \alpha \log P_R(x)\right] eq. 3

The constant C cancels in the log-likelihood ratio, so it is not estimated. The per-position form is:

\log P(x) \propto \sum_{i=1}^{l}\left[(1-\alpha)\log P_A(x_i \mid x_{<i}) + \alpha \log P_R(x_i)\right] eq. 4

The retrieval weight alpha is chosen by grid search on a validation subset of the benchmark, and the paper uses 0.6 throughout app. B.5.

# Retrieval profile for one wild-type family  [§4.1] [app. B.5]
# Input:  msa (N, L) aligned residues; weights (N,) sequence weights
# Output: log_profile (L, 20) per-position log-probabilities

def retrieval_profile(msa, weights):
    # Each sequence's one-hot rows get a 1e-5 pseudocount  [app. B.5]
    pseudo_one_hots = one_hot(msa) + 1e-5                 # (N, L, 20)
    # Weight each sequence, then sum over sequences per column  [§4.2]
    weighted = pseudo_one_hots * weights[:, None, None]   # (N, L, 20)
    counts = weighted.sum(axis=0)                         # (L, 20)
    # Normalise each column to a probability distribution  [app. B.5]
    probs = counts / counts.sum(axis=-1, keepdims=True)   # (L, 20)
    return log(probs)

# Fuse the two inference modes at each position  [§4.2] [eq. 4]
# Input:  ar_log_probs (L, 20) from the causal model; profile (L, 20);
#         covered (L,) booleans; alpha is the retrieval weight (0.6)
# Output: fused (L, 20)

def fuse(ar_log_probs, profile, covered, alpha):
    # Weighted arithmetic mean of the two log-probabilities  [eq. 4]
    mixed = (1 - alpha) * ar_log_probs + alpha * profile  # (L, 20)
    # Positions outside the alignment keep the autoregressive term  [app. B.5]
    return where(covered[:, None], mixed, ar_log_probs)   # (L, 20)

# Log-likelihood of one sequence in one direction  [§3.4] [eq. 1] [eq. 4]
# Sequences longer than the context must already be windowed  [app. B.4]

def sequence_log_likelihood(sequence, profile, covered, alpha):
    # Causal pass: predict each residue from the residues before it  [eq. 1]
    ar_log_probs = causal_model(sequence)                 # (L, 20)
    fused = fuse(ar_log_probs, profile, covered, alpha)   # (L, 20)
    # Sum the fused log-probability of the residue present  [eq. 4]
    return sum(fused[i, sequence[i]] for i in range(len(sequence)))

# Log-likelihood ratio in one reading direction  [§3.4] [eq. 2]

def directional_ratio(mutant, wild_type, profile, covered, alpha):
    # Difference of log-likelihoods, mutant minus wild type  [eq. 2]
    return (
        sequence_log_likelihood(mutant, profile, covered, alpha)
        - sequence_log_likelihood(wild_type, profile, covered, alpha)
    )

# Final fitness of a mutant  [§3.4] [app. B.4] [app. B.5]
# Input:  mutant, wild_type as residue strings
# Output: fitness score (float, higher means fitter)

def mutant_fitness(mutant, wild_type, profile, covered, alpha):
    # Left to right, as the model was trained  [§3.4]
    forward = directional_ratio(
        mutant, wild_type, profile, covered, alpha
    )
    # Right to left: reverse the strings and the profile together  [§3.4]
    backward = directional_ratio(
        mutant[::-1],
        wild_type[::-1],
        profile[::-1],
        covered[::-1],
        alpha,
    )
    # Plain average of the two directions  [§3.4]
    return (forward + backward) / 2

Hyperparameters

Results

Author-reported. Setup: ProteinGym substitution benchmark, 87 DMS assays from the paper's curation, Spearman rank correlation unless stated, single-seed scoring for all models §6.1 §6.2 app. F.

Paper vs code

The official repository is OATML-Markslab/Tranception, MIT licensed (copyright OATML and Marks Lab). The paper links it app. D app. F. The clone is at commit 2ddf40e1db9d2d180d1b5fc9d1b39ad5b04fbb6d. The repository contains inference and scoring code only. A training script was not found, so the training settings in Hyperparameters come from the paper alone.

• Kernel sizes. The paper says each head group gets a different convolution kernel §3.1 but gives no sizes. The code uses no convolution for the first group and kernels 3, 5 and 7 for the others model_pytorch.py model_pytorch.py. Inferred from code only.
• Grouped ALiBi slopes. The paper says each group learns attention at different distances §3.2. The code computes slopes for one quarter of the heads and repeats that list four times model_pytorch.py. The head groups therefore share the same slope values, and only the convolution differs between them model_pytorch.py.
• Window for long sequences. The paper says truncation was used "for simplicity" app. B.4, but also describes choosing the optimal window per mutation. The code default is the optimal window score_tranception_proteingym.py scoring_utils.py.
• Mirror scoring. Matches the paper: forward and reversed scores are averaged §3.4 model_pytorch.py. The reversed profile is flipped too model_pytorch.py.
• Retrieval weight and pseudocount. Default weight 0.6 and pseudocount 1e-5 match the paper config.py msa_utils.py.
• Gaps. The paper says gaps are ignored in the pseudocounts app. B.5. The code adds the pseudocount to every one-hot row before weighting msa_utils.py. I did not check how gaps are encoded, so whether gaps are fully excluded is not verified.
• Indel windows. In indel mode, the code keeps the whole mutated sequence as the window scoring_utils.py. Its own comment recommends the sliding window for indels in sequences longer than the context scoring_utils.py. The paper does not describe this.
• Indeterminate residues. The code replaces X, B, J and Z with standard residues when encoding model_pytorch.py (encode_batch). The paper describes random imputation §3.3; I did not check the replacement logic in detail.

Pitfalls

• Small overall margin. Over the whole substitution benchmark, the gap to EVE is 0.003 Spearman and 0.003 AUC. The tables read do not report uncertainty estimates or significance tests §6.2. Read "outperforms all baselines" as "highest mean on the benchmark".
• EVE leads on deeper alignments. Tranception is behind EVE in the medium and high depth groups §6.2. The advantage is in shallow alignments and multiple mutants.
• Validation overlap. Architecture choices and the retrieval weight were tuned on 10 of the 87 substitution assays, which are part of the same benchmark. The "full set" ablation numbers therefore include the tuning assays app. B.1 app. B.5.
• Retrieval weight is a judgment call. On the full set, alpha 0.7 gives slightly higher Spearman (0.446) than 0.6 (0.445) app. B.5. The paper keeps 0.6 because it was chosen on the validation set.
• MSA building is part of the method. Reproducing the retrieval results requires building alignments as described, with jackhmmer, the bit-score search and the evolutionary-coupling selection app. C. The paper points to a reference file in its repository with per-assay alignment settings. I did not inspect that file.
• Sequence weights depend on family. The paper cites Hopf et al. for weights but does not state theta. The code default is 0.2, and its docstring suggests 0.01 for viruses msa_utils.py msa_utils.py.
• Few multiple-mutant assays. Only 11 of the 87 substitution assays contain multiple mutants §7. The large multi-mutant gain rests on those 11 assays.
• Cost. The largest model took about two weeks on 64 A100 GPUs, per the paper app. B.3. Inference also needs a retrieved alignment per protein.
• Scoring is per residue. Each sequence is scored by one forward pass that predicts all positions at once, using masks to keep the model causal §4.2. Scoring a mutant still requires running the wild type too.
• Learning rate schedule is partly unstated. The paper says the rate is annealed over 10k steps to 3e-4, then decreased linearly until the end of training. The final value and the starting value are not stated app. B.3.

Lineage

• Builds on: the squared ReLU and depthwise convolution from Primer, and the split-by-kernel idea from Inception §3.1. The distance bias comes from ALiBi, modified to act per head group §3.2.
• Compared against: the site-independent model, EVmutation, DeepSequence, EVE, Wavenet, ESM-1v and MSA Transformer §6.1. Their scoring details are in app. D.
• Superseded by: not known from content already on this wiki.