Learning Latent Protein Languages for Autoregressive Generation

AI summary

Overview

Research area: Machine learning for protein generation — specifically, autoregressive transformer language modeling of protein sequences and 3D backbone structures.

Technical level: Advanced. The paper assumes familiarity with autoregressive transformers, vector quantization (VQ-VAE), protein language models (ESM-2), and protein structure benchmarks.

Scope: The paper introduces two learned discrete "latent languages" for proteins — one for amino acid sequences (PLL) and one for backbone structures (SLL) — and tests whether autoregressive transformers trained on these token spaces scale and generate better than models trained directly on amino acid tokens.

What This Paper Is About

Autoregressive transformers dominate generative AI in text, images, audio, and video, but they are comparatively weak for protein sequence and structure generation, where diffusion and flow-matching models generally set the state of the art. The authors argue this gap is partly a representation problem: amino acid tokens carry residue identity but not contextual semantics, and raw 3D coordinates are not discrete. Their goal is to build learned, decodable latent token languages for both protein modalities and then apply the standard next-token prediction recipe unchanged.

Key Contributions

1. PLL (Protein Latent Language): A contextual, residue-aligned sequence tokenizer with a 4,096-state alphabet, built on a frozen ESM-2 encoder using a VQ-RAE-style construction, decodable back to amino acid sequences with one token per residue.
2. SLL (Structure Latent Language): An adaptation of GCP-VQVAE Lite for backbone structures, adding auxiliary sequence and confidence supervision (amino acid identity, pLDDT, and frozen ESM-2 representation heads) while remaining decodable to backbone coordinates.
3. Scaling and generation evidence: PLLM shows a steeper fitted compute-scaling exponent than an amino acid autoregressive protein language model (PLM) under matched downstream training, and substantially less low-complexity drift in unconditional generation.
4. Inference-time scaling for structure prediction: SLLM is applied to sequence-to-structure prediction, combining DeepConf-adapted confidence filtering with token-space consensus to select candidates before coordinate decoding, with latent-token sampling reported as roughly 1,000 times faster than MSA-based AlphaFold2 for long proteins in their measurements.

Main Findings

• Steeper scaling for the latent sequence language: Under matched downstream sequence training, PLLM has a fitted compute-scaling exponent of 0.038 versus 0.020 for the amino acid PLM — roughly a 1.9x steeper fitted slope. SLLM has a fitted exponent of 0.049 on structure-token data, with different training exposure. The authors note the SLLM exponent is numerically close to an exponent reported for text language models against minimum training compute adjusted for batch efficiency, but describe that comparison as qualitative because corpora, tokenizers, compute accounting, and frontier construction differ.

• Less low-complexity drift: Across a 12-temperature unconditional sweep, 5,154/12,000 amino acid samples (43.0%) versus 2,383/12,000 PLLM samples (19.9%) had decoded residue-composition entropy below a heuristic 1.5-bit threshold. PLLM therefore reduces this diagnostic failure mode by 54% relative to the amino acid model, though low-complexity samples remain.

• Better match to training-data statistics at moderate temperatures: At T = 0.6 to T = 1.0, PLLM gives the strongest joint match between decoded sequence length and residue-composition entropy relative to the UniRef50 source distribution.

• PLL reconstruction quality: The final Stage 1 PLL checkpoint reaches ΔPPL% = 0.47, median normalized MSE of 0.044, reconstruction KL of 0.003, cosine similarity of 0.970, and full code activation on the validation split. Stage 2 achieves 99.81% micro token-level residue recovery on the UniRef50 validation split, with all standard amino acids above 99.17% class-wise recovery.

• SLL improves semantic utility of structure tokens: The selected recipe ("GCP-VQVAE Lite 2") improves over the base tokenizer by 24.2% in functional-site AUROC and 35.0% in physicochemical Spearman in the supervised PST ablation. In the combined ablation row, functional AUROC is 72.37 and physicochemical Spearman 32.30, versus 58.26 and 23.93 for the base.

• SLL trades a small amount of reconstruction for utility: SLL's average TM-score (0.9456) is 2.4% below the original GCP-VQVAE Lite tokenizer and close to AIDO (0.9456 vs 0.9470), while its average RMSD (1.4722 Å) is lower than AIDO's (1.8189 Å) over CASP14, CASP15, CASP16, CAMEO-2024, and the original zero-shot set.

• Better autoregressive target: In a matched Prot2Token-style sequence-to-structure setup, changing only the target tokenizer from the original GCP-VQVAE Lite to SLL reduces best validation perplexity by 34% and reaches a comparable training-loss regime 49% faster when progress is measured by training epoch.

• Backbone generation: diversity and novelty, not quality leadership: On structure-design ProteinBench averaged over target lengths 50, 100, 300, and 500, SLLM records scTM 0.5390, scRMSD 3.3974, TM similarity 0.4558, cluster ratio 0.9175, and Max TM 0.5536. La-Proteina has the highest mean scTM (0.9677) and lowest mean scRMSD (0.7871). Among generators with mean scTM > 0.5, SLLM has the highest cluster ratio and lowest Max TM.

• Filtered folding comparison favors the amino acid model on pLDDT: In the entropy-filtered T = 0.5 sequence comparison, the amino acid model has higher mean pLDDT than PLLM (65.26 vs 59.86), lower pairwise TM, and lower Max TM, while PLLM has the higher cluster ratio (0.9700 vs 0.9250). The authors note these describe different sample populations than the drift statistics.

• Structure prediction improves with more samples: On the held-out CAMEO-2024 set, both lDDT-CA and RMSD improve overall as the number of sampled SLL trajectories per input sequence increases. A DeepConf-style selector using only internal SLL-token confidence and token-space consensus recovers a meaningful fraction of the oracle best@k gain. The evaluated checkpoint sees 4.32 B structure-token targets.

Methodology in Plain English

The authors keep the standard decoder-only transformer and next-token prediction objective but change what the model predicts.

For sequences, they run a frozen ESM-2 150M encoder over a protein, then learn a transformer latent encoder plus a 4,096-entry codebook to assign each residue position a discrete code (chosen by negative Euclidean distance). Stage 1 trains this bottleneck to reconstruct the frozen ESM-2 representations themselves (loss of 10·MSE + 1·cosine + 0.5·commitment), and Stage 2 freezes the encoder and codebook and trains a fresh transformer decoder with a linear 21-class head (20 amino acids plus an unknown class X) to recover residue identities from the codes alone. At generation time, no ESM-2 is needed.

For structures, they keep the GCP-VQVAE Lite backbone interface: a GCPNet encoder produces geometric residue embeddings, a transformer contextualizes them, and vector quantization assigns one code per residue from a 4,096-token codebook; a transformer decoder with a six-dimensional rotation head maps codes back to backbone coordinates. They add three auxiliary training heads — two pre-decoder heads predicting amino acid identity and per-residue pLDDT, and one post-decoder head predicting frozen ESM-2 representations — used only during tokenizer training. The final recipe is selected by supervised performance on the Protein Structure Tokenization benchmark, not by reconstruction alone.

Once both token languages are fixed, they pretrain separate autoregressive transformers on amino acid tokens, PLL tokens, and SLL tokens using the same format (begin token, content, end token). Conditional generation prepends special tokens such as a length token. For sequence-to-structure prediction, a Profluent-E1 600M encoder supplies continuous residue-level context, SLL tokens are generated autoregressively, and the fixed SLL decoder produces coordinates. Multiple sampled SLL trajectories are scored by internal confidence and token-space consensus before decoding.

Sequence work uses UniRef50; SLL tokenizer experiments use AFDB representative structures; SLLM generation uses the larger AFDB ∩ cap UniRef50 structure-token corpus from GCP-VQVAE. Everything runs in PyTorch 2.8 on four nodes of eight NVIDIA A100 80 GB GPUs, using FlashAttention-2, mixed precision, AdamW, and cosine learning-rate decay.

Why This Matters

Impact on research: The paper reframes autoregressive protein modeling as a tokenization problem rather than an architecture problem. It provides a matched comparison showing that a learned contextual latent alphabet changes compute-scaling behavior and reduces a known pathological failure mode of left-to-right protein generation, and it positions a structure tokenizer's value in terms of semantic utility and autoregressive compatibility rather than coordinate reconstruction alone. The authors also explicitly state that their sequence comparison does not isolate which component — ESM-2 semantics, vector quantization, vocabulary expansion, or the learned tokenizer transform — produces the effect.

Real-world applications (as directions the work supports, given the paper's framing):

• Protein sequence design pipelines that generate diverse, non-repetitive amino acid sequences without drifting into low-complexity repeats.
• De novo backbone generation where diversity and novelty relative to PDB references matter more than self-consistency score alone.
• Faster sequence-to-structure prediction for long proteins, where the paper reports latent-token sampling roughly 1,000 times faster than MSA-based AlphaFold2 in their measurements.
• Inference-time scaling for structure prediction using model-internal confidence, avoiding reliance on ground-truth labels during candidate selection.

Industry relevance: The approach is a drop-in target vocabulary for standard decoder-only transformers, meaning existing autoregressive training and sampling infrastructure can be reused. The authors release code (github.com/mahdip72/latent_protein_languages) and a project page (mahdip72.github.io/latent-protein-languages.github.io/). Computational work was supported by NVIDIA and the NAIRR Pilot program under project NAIRR240246.

Future Directions

1. Unified multimodal pretraining: Each pretrained model here operates on a single token stream. The authors state that joint PLL–SLL co-generation, and compatibility between the independently learned languages on joint tasks, remain untested.
2. Closing the self-consistency quality gap: SLLM's strengths are diversity and novelty, while its backbone self-consistency quality remains below specialized generators such as La-Proteina. Improving quality-centered metrics is an explicit open problem.
3. Alternative uses of pretrained protein representations: The paper did not compare PLL against continuous-target generation or similarly resourced ESM-2 fine-tuning and distillation approaches.
4. Multichain complexes: Extending SLL to multichain complexes would require modeling relative chain poses, which the authors state they have not evaluated.

Target Audience

Researchers and practitioners working on protein language models, generative protein design, or discrete tokenization for scientific modalities. It is most useful for readers who already understand autoregressive transformers and vector quantization, and who want a matched empirical comparison of target representation choices rather than another new architecture. Readers primarily interested in state-of-the-art backbone quality metrics should note the paper positions its results in diversity and novelty rather than quality leadership.

Authors’ abstract

Autoregressive transformers remain comparatively weak for protein sequence and structure generation. We study the role of target representation: amino acid tokens encode residue identities without explicit contextual semantics, while backbone coordinates require a discrete representation in our framework. We introduce two learned latent protein languages. Protein Latent Language (PLL) maps sequences to a 4,096-state contextual alphabet built on a frozen ESM-2 encoder, with one token per residue. Structure Latent Language (SLL) adapts GCP-VQVAE Lite with auxiliary sequence and confidence supervision while retaining decoding to backbone coordinates. We separately pretrain autoregressive transformer models on PLL and SLL tokens using next-token prediction, yielding PLLM and SLLM. Under matched downstream sequence training, PLLM has a fitted compute-scaling exponent of 0.038 versus 0.020 for the amino acid autoregressive model. In unconditional sequence generation, PLLM reduces the fraction of samples below a heuristic 1.5-bit residue-composition entropy threshold by 54% relative to the amino acid model across sampling temperatures. For sequence-to-structure prediction, replacing the original GCP-VQVAE Lite tokenizer with SLL reduces best validation perplexity by 34% under matched training. For long proteins, latent-token sampling is approximately 1,000 times faster than MSA-based AlphaFold2 in our measurements. In backbone generation, SLLM compares favorably with other generative models on diversity and novelty. We also observe early signs that using SLLM's internal token confidence for inference-time sampling can improve sequence-to-structure prediction quality beyond a single decoded sample. These results position learned latent protein languages as a promising substrate for autoregressive transformer scaling and inference-time sampling in protein generation.