ProGen2: Exploring the Boundaries of Protein Language Models

In plain words

A protein is a long string of letters, one per building block. This work treats those strings the way large language models treat text: train a model to guess the next letter, feed it a very large pile of real protein sequences, and make the model bigger. Once trained, the same model can write new sequences and can score how plausible a given sequence looks.

The authors build a family of such models, from a small one to one with 6.4 billion parameters. The training data mixes sequences from several databases, including antibody sequences from immune repertoires. They ask three questions: how well the models predict held-out natural sequences, whether the sequences they write are new yet look like real proteins, and whether their scores predict how well a protein variant works in a lab measurement.

The most interesting answer concerns the third question. The biggest model is not the best at ranking variants of lab-measured proteins; a mid-sized model does better there. The authors read this as a sign that the data a model learns from matters as much as its size.

TL;DR

ProGen2 is a family of decoder-only transformers trained with next-token prediction on protein sequences, from 151M to 6.4B parameters §3.1 tab. 1. Larger models give lower held-out perplexity §4.1 tab. 2, but zero-shot fitness ranking peaks at 764M parameters and does not improve with size in general §4.3 tab. 3. Generated sequences are often predicted to fold like natural proteins, usually at low sequence identity §4.2 fig. 1. The official repository is BSD-3 licensed and contains inference and likelihood code, but not the training code or the training data README.md. All numbers are author-reported.

Problem

Engineering a protein usually means directed evolution: mutate a starting sequence, measure each variant, and repeat §1. Models that help choose variants have been either family-specific or universal. Family-specific models depend on alignments, which are shallow for proteins with few relatives, hard to align for families such as antibodies, and need new sequences aligned to the training set §2. Transformer models trained on unaligned sequences avoid these problems, and autoregressive ones had already been applied to fitness prediction §2.

The paper asks what happens when model size and data are both scaled up, and it measures three things: how well the model captures the distribution of natural sequences, how well it generates novel sequences, and how well its likelihood ranks experimentally measured variants without fine-tuning on those variants §1 §4.

Method

ProGen2 models are autoregressive transformers trained with a next-token language-modelling loss. The family is trained at 151M, 764M, 2.7B and 6.4B parameters §3.1. The architecture is a standard decoder with left-to-right causal masking, rotary positional encodings, and an attention block and a feed-forward block that run in parallel on the same normalised input §3.1. The paper states the residual update as x(t+1) = x(t) + attn(ln(x(t))) + mlp(ln(x(t))), in place of the sequential form §3.1.

Pretraining data is a mixture of UniRef90 and BFD30 for the standard models; a second variant mixes UniRef90 with representative clusters from UniProtKB, Metaclust, SRC and MERC, and is called BFD90 §3.2. The antibody model is trained on 554M antibody sequences from the Observed Antibody Space, clustered at 85% identity §3.2. Each sequence is wrapped in a one- or two-character terminal token at each end, then provided "as-is and flipped" §3.2. Proteins are packed into a batch until the maximum token length is filled §3.2.

Steps, as the paper describes them:

1. Wrap each protein with a 1- or 2-character terminal token at each end §3.2.
2. Provide the sequence as-is and flipped. The paper does not say how the two orientations are weighted §3.2.
3. Pack several proteins to fill the maximum context length §3.2.
4. Pretrain with Adam and gradient-norm clipping. Learning rate, clipping and length come from tab. 1 §3.3.
5. Watch the global gradient norm. The paper treats a rising norm as a sign of divergence from the expected power-law trend. Lowering the learning rate, raising weight decay and lowering the clipping threshold kept the norm near constant §3.3.
6. Fine-tune on a narrow family: re-initialise both Adam moment estimates, cut the initial learning rate by a factor of 5, and train for at most two epochs §3.3.
7. Generate with temperature and nucleus sampling, then cut each sample at the first terminal token §3.4.
8. Score a sequence by its average log-likelihood under the model for zero-shot fitness §4.3.

The perplexity used in the language-modelling evaluation is stated inline, with no equation number:

ppl(x) = \exp\left(-\frac{1}{n} \sum_{i=1}^{n} \ln p(x_i)\right) §4.1

The paper gives no numbered equations and no algorithm blocks. The pseudocode below is a reading of the code, cited to the repository; it is not a copy of it.

# One ProGen2 decoder block: attention and MLP share one normed input
# [§3.1] [code: progen2/models/progen/modeling_progen.py#L261-L275]
def block(hidden):
    # One pre-norm; both branches read this normalised copy  [§3.1]
    normed = layer_norm(hidden)
    # Parallel residual: x + attn(ln(x)) + mlp(ln(x))  [§3.1]
    hidden = hidden + attention(normed) + mlp(normed)
    return hidden

# Causal multi-head attention, partial rotary embedding
# [§3.1] [code: progen2/models/progen/modeling_progen.py#L147-L222]
def attention(normed):
    # One fused projection without bias; split order is query, value, key
    # [code: progen2/models/progen/modeling_progen.py#L157-L164]
    fused = linear_no_bias(normed)            # (B, N, 3d)
    query, value, key = split_heads(fused)    # (B, H, N, d_head) each
    # Rotary embedding on the first rotary_dim dims of each head, base 10000
    # [code: progen2/models/progen/modeling_progen.py#L38-L44, #L178-L190]
    query = rotate_first_dims(query, rotary_dim)
    key = rotate_first_dims(key, rotary_dim)
    # Scaled dot product with causal mask; scores kept in float32
    # [code: progen2/models/progen/modeling_progen.py#L122-L135]
    scores = query @ transpose(key) / sqrt(d_head)
    scores = mask_future_positions(scores)    # fill with -1e9
    weights = softmax(scores)
    # Output projection without bias  [code: ...modeling_progen.py#L215]
    return linear_no_bias(weights @ value)

# Sample a continuation with temperature and nucleus (top-p) sampling
# [§3.4] [code: progen2/sample.py#L69-L75]
def generate(context, temperature, top_p, max_length):
    # Encode the context; the default context is the token "1"
    # [code: progen2/sample.py#L121]
    tokens = encode(context)
    # Sample each next token from the temperature-scaled distribution,
    # restricted to the smallest set whose probability reaches top_p
    tokens = sample_nucleus(tokens, temperature, top_p, max_length)
    # Cut after the first terminal tag after position 0
    # [code: progen2/sample.py#L78-L87, #L194]
    return truncate(decode(tokens), terminals=["1", "2"])

# Zero-shot score: average log-likelihood over both reading directions
# [§4.3] [code: progen2/likelihood.py#L285-L296]
def fitness_score(sequence):
    # Mean log-likelihood of the sequence read left to right
    forward = mean_log_likelihood(sequence)
    # Same score on the reversed string, tags included  [likelihood.py#L287]
    backward = mean_log_likelihood(reverse(sequence))
    # The code averages the two directions  [likelihood.py#L296]
    return 0.5 * (forward + backward)

# Mean log-likelihood with the softmax restricted to letter tokens
# [code: progen2/likelihood.py#L174-L200]
def mean_log_likelihood(sequence):
    # Next-token logits and targets, shifted by one position  [§4.1]
    logits, targets = next_token_pairs(sequence)
    # Drop a trailing BOS/EOS id, then keep only ids 5..29
    # [code: progen2/likelihood.py#L184-L195]
    logits, targets = restrict_to_letters(logits, targets, 5, 29)
    return -cross_entropy(logits, targets, reduction="mean")

Hyperparameters

Complete model and optimiser table, as printed in the paper. Values are copied from the table; the unit of batch size is not stated tab. 1.

Derived, not stated: heads multiplied by head dimension gives a model width of 1024 (small), 1536 (medium and base), 2560 (large) and 4096 (xlarge) tab. 1.

Other settings, from the text and from the paper's evaluation section:

Results

Author-reported. Setup, in short: language modelling on held-out Test-max90 and Test-max50 sets §3.4 tab. 2; generation of 5,000 sequences from ProGen2-xlarge, with structures predicted and compared with the PDB §4.2; zero-shot Spearman correlation on narrow and wide fitness landscapes and on antibody landscapes §3.4 §4.3.

Headline results:

• Perplexity falls as models grow. Test-max90 goes from 12.9 (151M) to 9.9 (6.4B) tab. 2.
• Narrow landscapes, average Spearman rho: ProGen2-base 0.505 is the best single ProGen2 size; an ensemble column reaches 0.518; ProGen2-small 0.456 is above RITA-XL at 0.443 tab. 3.
• Across 5,000 xlarge-generated sequences, the paper reports a median TM-score of 0.89 and median pLDDT of 73.7 §4.2.
• Antibody landscapes: ProGen2-small has the best binding average (0.44); ProGen2-xlarge has the best general-property average (0.74) tab. 5.

Table 2: Perplexity of held-out sequences tab. 2.

Table 3: Zero-shot fitness, narrow landscapes, average Spearman rho. Bold marks are as printed in the paper. Baseline values are taken from Hesslow et al. (RITA) according to the caption tab. 3.

Table 4: Zero-shot fitness, wide landscapes tab. 4.

Table 5: Zero-shot fitness, antibody landscapes, averaged tab. 5.

Table 6: Zero-shot fitness on indel datasets, from Notin et al. (Tranception) tab. 6, appendix A.

Table 7: Full antibody results tab. 7, appendix A.

Table 8: Two 2.7B models trained on different metagenomic mixes tab. 8, appendix A.

Table 9: Full narrow-landscape results, Spearman rho per dataset tab. 9, appendix A. The average row matches Table 3. Bold marks appear only in the average row of the source.

Figures, described in words because images are not supported yet:

• fig. 1: Universal generations from ProGen2-xlarge §4.2. Panel A plots AlphaFold2 confidence (pLDDT) against structural similarity to the nearest PDB protein (TM-score). Panels B to D show three generated structures: a solenoid fold with 11.2% identity to its nearest PDB homologue; a multi-domain alpha+beta fold with no structural homologue above a low TM-score (0.283, 11.1% identity); and a beta-roll fold resembling the RTX domain (TM-score 0.374, 30.3% identity) §4.2.
• fig. 2: ProGen2-large fine-tuned on two-layer sandwich proteins (CATH 3.30), after one and after two epochs §4.1 §4.2. Panels A and B show that more fine-tuning gives generations closer to natural sequences, and that higher temperature or top-p gives more sequence diversity. Panels C and D show that more restrictive sampling gives structures closer to PDB proteins. Panels E to G show that structural novelty concentrates in ligand-binding regions.
• fig. 3: Antibody generation from ProGen2-OAS §4.2. Unprompted and EVQ-prompted generations are compared on sequence identity to the training set, on aggregation propensity (SAP score) and on solubility (CamSol-intrinsic), across temperature and top-p. Less restrictive sampling gives lower aggregation and higher solubility. Re-ranking with ProGen2-base improves solubility for the top 50% but does not significantly reduce aggregation fig. 3.

Paper vs code

Checked at commit c27a419c234a0997923761e1fe7daffcebf0eaf5. The repository holds inference and likelihood code (progen2/sample.py, progen2/likelihood.py), a PyTorch model and configuration, a tokenizer file, and a README. It does not contain training code, training data, or the per-checkpoint configuration files that the README downloads README.md.

• Layer counts. The paper gives 12, 27, 27, 32 and 32 layers tab. 1. The code default is 28 layers configuration_progen.py, which matches no row. The per-checkpoint value lives in the downloaded checkpoints, which were not fetched here, so the paper's Table 1 is the only stated source.
• Heads and width. The code default of 16 heads and width 4096 configuration_progen.py matches the paper's xlarge row for heads and head dimension (16 and 256) tab. 1.
• Parallel block. The code uses one normalisation for both branches and sums attention, MLP and residual modeling_progen.py. This agrees with the residual form in §3.1.
• Rotary embedding. The paper does not give a rotary dimension. The code applies rotary embedding to the first rotary_dim dimensions of each head, with a default of 64, base 10000, and interleaved pairs modeling_progen.py configuration_progen.py.
• Fused projection order. The code splits the fused projection into query, value and key, in that order modeling_progen.py. The paper does not describe this layout. The code also hard-codes a logical split of 8 groups, with a TODO comment about it modeling_progen.py.
• Sampling defaults. The code defaults to temperature 0.2, top-p 0.95 and max length 256 sample.py. The README example uses temperature 0.8, top-p 0.9 and max length 1024 README.md. The paper reports sweeps over temperature and top-p, not one headline setting §3.4.
• Scoring. The paper defines perplexity but not the zero-shot score used in Tables 3, 4 and 9 §4.1 §4.3. The code's likelihood.py averages the forward and reversed scores and can return sum or mean likelihood.py. The paper does not say whether this produced its tables.
• Softmax normalisation. The sanity check in sample.py uses the full-vocabulary cross entropy sample.py. The scoring in likelihood.py restricts the softmax to 25 letter tokens (ids 5 to 29) likelihood.py. The paper's perplexity formula does not say which is used §4.1.
• Sequence orientation. The code's sanity sequences show both orders: one starts with the tag "2" and ends with "1", and two start with "1" and end with "2" sample.py. The default generation context is "1" sample.py. The paper does not assign the tags to N- or C-terminal positions.
• Oracle values. The code's own sanity check requires the cross entropy of a fixed sequence to be within 0.1 of a per-checkpoint value. On the UniRef-style sequence: small 2.4, medium 1.9, base 1.9, large 1.8, xlarge 1.0. On the OAS sequence the OAS model gives 0.3, and on the BFD sequence the BFD90 model gives 1.3 sample.py. These values are not in the paper and can serve as a check for a re-implementation.
• Model naming. The code's checkpoint names map to the paper's models: progen2-small, -medium, -base, -large, -xlarge, -oas and -BFD90 sample.py. The code's size grouping (151M, 754M, 2B, 6B) differs slightly from the paper's rounding.
• License. The repository is BSD-3-Clause LICENSE.txt. Two model files are adapted from HuggingFace transformers code under Apache-2.0 modeling_progen.py. The README states that code and models are BSD-3 licensed README.md.

Pitfalls

• Layer count is not in the code. Use Table 1 for the paper's sizes tab. 1, and read config.json from the checkpoint before building a model. The code default of 28 layers is not any paper size.
• Training is not reproducible from this repository. No training loop, data pipeline or data-mixing script is included. Learning-rate warm-up shape, final learning rate, weight-decay style, batch-size unit and the number of accelerators are not stated §3.3 tab. 1.
• Fused projection layout. A re-implementation must split the fused query/value/key projection in the same order as the code to load the weights modeling_progen.py.
• Activation and MLP width. The paper does not name its activation or MLP width. The code defaults to GELU (gelu_new) and to four times the width configuration_progen.py modeling_progen.py.
• Attention numerics. Attention scores are computed in float32 and scaled by the square root of the head dimension. The causal mask fills with -1e9, and the padding mask adds -10000 modeling_progen.py. The paper does not mention these choices.
• Scoring ambiguity. The zero-shot score is undefined in the paper: sum or mean, which direction, whether tags are scored, and whether the softmax is restricted. The code makes choices that the paper does not state likelihood.py.
• Orientation and augmentation. The paper says sequences are used "as-is and flipped" but not how often each orientation is seen in training §3.2.
• Sampling counts. "5,000 sequences" for the universal experiment is not clearly per setting or in total §3.4 §4.2. Fine-tuned generation gives 10,000 per checkpoint, derived from the 20,000 total §4.2.
• Non-monotone table. Test-max50 perplexity is 14.3 for medium and 14.4 for large tab. 2. The text says larger models yield "substantially lower" perplexities §4.1.
• Ensemble column. Table 3 reports an ensemble at 0.518 tab. 3. The paper does not define how the ensemble is formed. The text says ProGen2-base is "the best performing ProGen2 size" and that the best model "outperforms or matches" the baselines §4.3 tab. 3.
• Antibody claim. The text says the OAS model "performs poorly" §4.3. On the general-property average, OAS (0.66) beats ProGen2-small (0.61), and it is below the three larger universal models tab. 5.
• Single-landscape emergence claim. The "emergent behavior" on GB1 rests on one landscape, where the top-100 average jumps from 0.24 (large) to 0.85 (xlarge) §4.3 tab. 4. Treat it as an observation, not a general result.
• Baselines are copied. The baseline numbers in Table 3 and Table 9 are taken from Hesslow et al. (RITA), as the caption says tab. 3. They were not re-run here.
• Unexplained baseline. The WaveNet column in Table 6 is not described in the text tab. 6.
• Structure metrics depend on settings. Generation is scored with AlphaFold2 using single-sequence input and one or three recycles, and with TM-score via Foldseek §3.4. Changing those choices changes the reported medians.
• Table and figure numbers are counted in LaTeX source order, because the compiled PDF was not read. Check them against the PDF if numbers disagree.

Lineage

• Builds on: the authors' earlier ProGen model, which is fine-tuned on lysozyme families §2. The architecture follows GPT-3 hyperparameter choices §3.1 §3.3, the GPT-J parallel attention and MLP layout §3.1, and rotary positional encodings §3.1.
• Compared against: RITA-XL, EVE, Tranception (with and without retrieval), MSA Transformer and ESM-1v on zero-shot fitness tab. 3 tab. 9; Tranception and WaveNet on indels tab. 6; IgLM is discussed for antibody generation §2.
• Superseded by: not stated in this card.