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

TL;DR

Proust is a novel causal protein language model that combines the strengths of MLMs and generative models, achieving state-of-the-art results in fitness prediction and maintaining generative capabilities.

Contribution

The paper introduces Proust, a 309M-parameter causal PLM with architectural innovations enabling efficient training and superior performance on protein fitness benchmarks.

Findings

Achieves Spearman ρ=0.390 on ProteinGym substitutions

Sets new state-of-the-art on indels, outperforming larger models

Approaches structure-aware methods using sequence alone

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 $\rho = 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