No Generation without Representation: Efficient Causal Protein Language Models Enable Zero-Shot Fitness Estimation

TL;DR

Proust is a 309M-parameter causal protein language model that closes the gap between generation-capable PLMs and fitness-predictive MLMs. By integrating architecture-agnostic innovations borrowed from LLMs—GQA-S2 attention with shared K/V projections, cross-layer value residuals, key-offset induction heads, Canon depthwise convolutions, and a stable Muon optimizer—Proust matches MLM performance on substitutions and sets new state-of-the-art on indels while dramatically reducing training compute. It also preserves autoregressive generation and offers interpretable insights into when retrieval augmentation helps, guided by per-position entropy patterns. Across ProteinGym, EVEREST, and related benchmarks, Proust demonstrates strong, robust performance with favorable test-time tradeoffs, and the work highlights practical avenues for scaling and retention of generation in protein modeling.

Abstract

Protein language models (PLMs) face a fundamental divide: masked language models (MLMs) excel at fitness prediction while causal models enable generation, forcing practitioners to maintain separate architectures. We introduce \textbf{Proust}, a 309M-parameter causal PLM that bridges this gap through architectural innovations adapted from recent LLM research, including grouped-query attention with shared K/V projections, cross-layer value residuals, and depthwise causal convolutions. Trained on 33B tokens in 40 B200 GPU-hours, Proust achieves Spearman $ρ= 0.390$ on ProteinGym substitutions, competitive with MLMs requiring 50--200$\times$ the compute. On indels, Proust sets a new state-of-the-art, outperforming models up to 20$\times$ larger. On EVEREST viral fitness benchmarks, it approaches structure-aware methods using sequence alone. These powerful representations position Proust in a sweet spot as it also retains native generative capabilities that MLMs lack by design. Interpretability analysis reveals that per-position entropy variance predicts, to an extent, when retrieval augmentation helps and hurts. Such insights can grow in both quantity and quality at scale and inform capabilities such as test-time scaling. Code and weights are available at https://github.com/Furkan9015/proust-inference

Figures (5)

• Figure 1: ProteinGym substitutions: accuracy vs. training compute. Proust (star) achieves competitive Spearman $\rho$ at 229$\times$ fewer FLOPs than E1-600M. Marker shape indicates model family and color indicates training objective (CLM red, MLM teal, GLM purple). Single-sequence models are shown with solid markers and retrieval-augmented models (E1) with hollow markers.
• Figure 2: ProteinGym indels: accuracy vs. training compute. Proust achieves the highest Spearman $\rho$ among all compared models while using the fewest training FLOPs. Causal models handle indels naturally since insertions and deletions do not disrupt autoregressive scoring.
• Figure 3: EVEREST viral fitness benchmark. Points are colored by assay type and horizontal bars indicate mean $\rho$ per model. Proust shows lower variance across assays than baselines, avoiding failure modes but missing peaks on stability assays where structure-aware models (SaProt) excel.
• Figure 4: Effect of homolog depth on ProteinGym substitution performance. We select the $N$ most similar homologs (by sequence identity, $\geq$50% coverage). Mean Spearman $\rho$ (red, left axis) improves with depth. The number of assays with sufficient homologs (gray, right axis) decreases at higher depths, as not all proteins have deep MSAs.
• Figure 5: Learning rate transfer from 50M to 309M. Validation perplexity curves for the 50M proxy models and the full 309M model. Legend entries are formatted as modelsize_learningrate_weightdecay.