Source-linked AI summary
Learning the Language of Protein Structure
Benoit Gaujac, Jérémie Donà, Liviu Copoiu, Timothy Atkinson, Thomas Pierrot, Thomas D. Barrett
TL;DR
Protein structures are continuous and three-dimensional, and no established methodology readily applies sequence modeling to them. This paper learns quantized structure representations and uses them with a simple GPT, which generates diverse, structurally viable proteins and is competitive with some diffusion-based models, though it does not yet match RFDiffusion.
Problem
Protein structures are continuous and three-dimensional, and no established methodology readily applies sequence modeling to them.
Source: §1 · p. 2Method
The authors learn quantized representations of protein structures so sequence-based language models can operate on them.
Source: §1 · p. 2Results
The simple GPT generates novel, diverse, structurally viable structures and is competitive with some diffusion-based models, but does not yet match RFDiffusion.
Source: §1 · p. 2, §3.2 · p. 8, §5 · p. 10Takeaways & Limitations
Tokenized structures provide a sequential representation that sets the stage for applying standard sequence-modeling techniques to protein structures.
Source: §5 · p. 10The simple GPT does not yet match seminal approaches like RFDiffusion.
Source: §5 · p. 10
Abstract
from arXiv · showhide
Representation learning and \emph{de novo} generation of proteins are pivotal computational biology tasks. Whilst natural language processing (NLP) techniques have proven highly effective for protein sequence modelling, structure modelling presents a complex challenge, primarily due to its continuous and three-dimensional nature. Motivated by this discrepancy, we introduce an approach using a vector-quantized autoencoder that effectively tokenizes protein structures into discrete representations. This method transforms the continuous, complex space of protein structures into a manageable, discrete format with a codebook ranging from 4096 to 64000 tokens, achieving high-fidelity reconstructions with backbone root mean square deviations (RMSD) of approximately 1-5 \AA. To demonstrate the efficacy of our learned representations, we show that a simple GPT model trained on our codebooks can generate novel, diverse, and designable protein structures. Our approach not only provides representations of protein structure, but also mitigates the challenges of disparate modal representations and sets a foundation for seamless, multi-modal integration, enhancing the capabilities of computational methods in protein design.
1 INTRODUCTION
Protein structures are continuous and three-dimensional, making them difficult to model directly with sequence-based methods. The paper proposes quantized representations to bring structures into sequence-modeling workflows and reports that a simple GPT can generate diverse, structurally viable protein structures.
1 INTRODUCTION
- The broader motivation is to model interconnected biological modalities, including sequences, structures, functional descriptors, and experimental assays.Transformer-based sequence modeling has also been applied to biological sequences.Source: §1 · p. 1
- No established methodology readily applies sequence modeling to protein structures, despite sequence-space methods being used across other modalities.Large multimodal models combine modalities in sequence space using pretrained encoders.Source: §1 · pp. 1–2
- Protein structures are continuous and three-dimensional, complicating direct use of transformer models built primarily for discrete data.Structure generation commonly uses methods for continuous variables, such as diffusion and flow matching.Source: §1 · p. 2
- The paper aims to map protein structures into quantized representations so sequence-based language models can be used.Its objectives include discrete sequential structural data, a discrete potentially low-dimensional latent space, and low reconstruction error.Source: §1 · p. 2
- A simple GPT trained on the learned codebook generates novel, diverse, and structurally viable protein structures.The authors present this as evidence for the efficacy and practicality of their learned representations.Source: §1 · p. 2
2 METHOD
The method encodes protein backbone coordinates as a downsampled sequence, quantizes that representation, and decodes it into a structure. Training combines a frame-aware reconstruction loss with filtered, cluster-sampled protein data.
2.1 PROTEIN STRUCTURE AUTOENCODER
- The autoencoder maps backbone coordinates to a downsampled sequence of finite codes, then decodes those codes into predicted backbone coordinates.Figure 1 depicts graph encoding, quantization, and decoding to estimate backbone-atom positions.Source: §2.1 · p. 2, Figure 1 · p. 3
- The encoder uses an MPNN with cross-attention to downsample representations while a local attention mask limits aggregation to nearby sequence nodes.Its graph features use relative distances and orientations that are invariant to rotation and translation.Source: §2.1 · p. 3
- Finite Scalar Quantization rounds bounded low-dimensional encodings into discrete values whose level combinations define codebook indices.The authors use FSQ to address the expense and training instability associated with explicit vector quantization, including codebook collapse.Source: §2.1 · pp. 3–4
- The decoder uses the AlphaFold-2 structure module, trained from scratch, to estimate a point cloud from the quantized sequence.It constructs per-residue and pairwise representations and parameterizes the point cloud using frames.Source: §2.1 · p. 4
2.2 TRAINING
- Frame Align Point Error compares predicted and true coordinates in their respective local frames, and clamping the loss at 10 improves training stability.The loss supports joint optimization of frames and coordinates.Source: §2.2 · p. 4
- Training uses about 310,000 PDB entries, samples inversely by sequence-cluster size, filters chains under 50 residues, and crops longer structures to 512 residues.The authors allocate 90% of clusters to training and divide the remaining 10% into validation and test sets.Source: §2.2 · pp. 4–5
- The models are trained for 250 epochs with AdamW on 128 TPU v4-8 devices, using batch size 1024 and a learning rate warmed from 10^-6 to 10^-3.The encoder uses a three-layer message-passing network; the AlphaFold structure module uses six structure layers.Source: §2.2 · p. 5
3 EXPERIMENTS
Experiments assess reconstruction quality of the discrete autoencoder and test its use in GPT-based de novo protein-structure generation.
3.1 AUTOENCODER EVALUATION
- At K = 64000 and r = 1, average reconstruction achieved 1.59 Å RSMD and a TM-score of 0.95, approaching experimental resolution.Table 1 reports average test-set reconstruction results and uses medians for CASP-15 because of its limited dataset size.Source: §3.1 · p. 5, Table 1 · p. 6
- Larger codebooks reduce reconstruction errors at a fixed downsampling ratio, while greater downsampling or smaller codebooks worsens accuracy; quantization costs about 0.5–1 Å versus continuous autoencoders.Across settings, reconstruction remains within a few angstroms and average TM-scores exceed 0.5; the authors also report that increasing codebook size indicates no codebook collapse.Source: §3.1 · pp. 5–6
- For r = 2 and K = 64000, reconstructions capture global structure and preserve local α-helices and β-sheets.Figure 3 shows reconstructed structures superimposed on their targets from different rotation angles.Source: §3.1 · p. 6, Figure 3 · p. 7
3.2 De novo PROTEIN STRUCTURE GENERATION
- The generation experiment trains a 344M-parameter, 20-layer decoder-only GPT with next-token prediction on r = 1, K = 4096 structural tokens, then decodes generated tokens into structures.The parameter choice balances reconstruction performance, dataset size, and parameter efficiency.Source: §3.2 · p. 7
- 9.01% (70 out of 777) of GPT samples met cathTM < 0.5 and scRMSD < 2 Å, versus 8.18% (53 out of 648) for FrameDiff and 58.58% (379 out of 647) for RFDiffusion.The paper describes GPT's designability as competitive with FrameDiff, while noting that novelty and diversity results are more nuanced and depend on the sampling strategy.Source: §3.2 · p. 8
- Generated samples show varied secondary structure, and their ESM-predicted structures align closely with the samples.Figure 4 visualizes the generated structures in green and their self-consistent ESM-predicted structures in blue.Source: §3.2 · p. 8, Figure 4 · p. 9
4 RELATED WORKS
Related work covers learning from protein structures, discrete structural representations, and generative approaches, highlighting distinctions in scope and representation.
4 RELATED WORKS
- Protein-structure learning includes inverse folding and binding-site estimation, as well as methods that use structure to study biological processes.The cited examples connect inverse folding and binding-site estimation to drug design.Source: §4 · p. 9
- A VAE using distance matrices to predict 3D coordinates offers straightforward latent-space sampling, but sampling can yield subpar quality.The passage describes supervision through comparison of distance matrices.Source: §4 · p. 9
- Foldseek quantizes local residue-level geometry but lacks global structural representation, limiting use in structure generation or binding prediction.Other approaches combine structure tokens with sequence tokens or cluster latent representations on a fixed reference dataset.Source: §4 · p. 9
- Prior discrete structural approaches include VQ-VAE representations and diffusion in their latent space; Gao et al. (2023) is described as limited by reconstruction performance.Liu et al. (2024a) concurrently combines finite scalar quantization with a specialized transformer-based autoencoder for proteins, RNA, and small molecules.Source: §4 · p. 9
- FoldToken jointly quantizes sequence and structure, whereas this work focuses on structural information and complete-structure de novo generation.The passage contrasts FoldToken's emphasis on backbone inpainting and antibody design with this work's application.Source: §4 · pp. 9–10
- Protein-structure generation literature includes diffusion and flow-matching methods, alongside approaches that tokenize structures and train language models on combined sequence and structure tokens.The passage notes that one symmetry-preserving parameterization works only on small protein crops.Source: §4 · p. 10
5 CONCLUSION
The work maps protein geometry to discrete integer sequences that decode to near-native conformations, providing a compressed representation for sequence-based modeling. FSQ supports improved reconstruction with larger codebooks or more tokens, while a proof-of-concept GPT generates backbones but remains behind RFDiffusion.
5 CONCLUSION
- The method maps protein geometry to discrete integer sequences that can be decoded into near-native conformations while compressing the representation relative to 3D coordinates.The authors identify enabling standard sequence-modeling techniques for protein structures as the work’s primary contribution.Source: §5 · p. 10
- FSQ mitigates codebook collapse, and larger vocabularies or increased token usage offer routes to improved reconstruction by reducing the information bottleneck.The conclusion presents expressiveness and high-fidelity reconstruction as prerequisites for sequence-based structure modeling.Source: §5 · p. 10
- A proof-of-concept GPT trained on tokenized PDB entries generates protein backbones and is competitive with some recent diffusion models, though it does not match RFDiffusion.The authors describe further sequence-modeling development as a promising direction.Source: §5 · p. 10
A.1 ARCHITECTURES
The architecture compresses residue-level representations through resampling and finite scalar quantization, then decodes the discrete codes into protein structures. Local attention masks constrain the resampling step, while a pairwise representation supports the structure decoder.
A.1 ARCHITECTURES
- Finite scalar quantization projects residue embeddings into a six-dimensional latent space, rounds the channels to discrete values, and projects the quantized codes back to the original dimensionality.The codebook contains 4096 entries with levels [4, 4, 4, 4, 4, 4], or 64000 entries with levels [8, 8, 8, 5, 5, 5].Source: §A.1 · p. 16
- Cross-attention resampling can change sequence length, and local masks restrict updates to direct residue neighbors to encourage local information and exclude distant embeddings.The encoder preserves residue order but does not use amino-acid identities.Source: §A.1 · pp. 16–17
- The decoder reuses AlphaFold’s Structure Module, deriving pairwise features from projected quantized codes and relative positional encodings.The pairwise module forms an outer product of left and right projections before combining it with relative positional encoding.Source: §A.1 · p. 17
- The overall pipeline applies a graph neural network, downsamples and quantizes residue embeddings, upsamples the codes, forms pairwise features, and decodes coordinates.The listed algorithm returns both quantized output and reconstructed coordinates.Source: §A.1 · p. 17
A.2 TRAINING DETAILS AND METRICS
The described pipeline represents residue order and local geometric relationships in its inputs, then maps backbone coordinates through encoding, resampling, quantization, and decoding. The compression factor compares the storage required for backbone positions with that required for latent codes.
A.2 TRAINING DETAILS AND METRICS
- Edge features use residue-local coordinate frames built from backbone-atom directions and concatenate geometric and positional information.The local frame uses directions from the alpha-carbon to the nitrogen and carboxyl carbon atoms, plus their plane normal.Source: §A.2 · p. 17
- The overall algorithm takes backbone coordinates as input, encodes and downsamples residue features, quantizes them, upsamples them, and decodes reconstructed coordinates.It also constructs a pairwise representation before the structure-decoding step.Source: §A.2 · p. 18
- The compression factor is the ratio of bits for backbone-position coordinates to bits for latent codes, assuming 64-bit floats for positions and 16-bit unsigned integers for codes.The stated code storage assumption follows from the largest codebook size being 64000, below 2^16.Source: §A.2 · p. 18
A.3 ADDITIONAL RESULTS: STRUCTURES AUTOENCODING
Figure 6 presents RMSD distributions on CASP-15 across downsampling ratios for a fixed codebook size of 4096.
A.3 ADDITIONAL RESULTS: STRUCTURES AUTOENCODING
- Figure 6 plots CASP-15 RMSD distributions across downsampling ratios with the codebook fixed at 4096.Source: Figure 6 · p. 18
A.4.1 TRAINING DETAILS
The GPT prior is trained on padded structural-token sequences with a 512-token maximum, using a 344M-parameter decoder-only transformer and AdamW optimization. The training set contains 70M actual structural tokens; the appendix also tracks reconstruction metrics across downsampling ratios.
A.4.1 TRAINING DETAILS
- The 344M-parameter decoder-only transformer has 20 layers, 16 heads per layer, dimension 1024, and 10% training dropout.It uses pre-layer normalization and follows Hoffmann et al. (2022).Source: §A.4.1 · p. 18
- Figures 7 and 8 track TM-Score and RMSD distributions on CASP-15 as downsampling ratio changes, with codebook size fixed at 4096.The supplied captions do not state the direction or magnitude of any change.Source: Figure 7 · p. 19, Figure 8 · p. 19
- Sequences are padded to length 514, allowing at most 512 actual structural tokens, and padding-token loss is masked.The preparation prepends <bos> and appends <eos>.Source: §A.4.1 · p. 19
- Optimization uses AdamW with β1 = 0.95, β2 = 0.9, weight decay 0.1, and a 1,000-step warm-up to learning rate 5.10−5.The batch size is 65,792 tokens, and embedding, residual, and attention dropout are each 10%.Source: §A.4.1 · p. 19
- The training corpus contains 70M actual structural tokens; the authors suggest predicted structures such as AlphaFold 2 could improve latent-prior training.This is presented as a prospective opportunity, not a measured result.Source: §A.4.1 · p. 19
A.4.2 STRUCTURE GENERATION METRICS
The appendix defines designability, novelty, and diversity measures for generated structures, then describes sampling choices and additional evaluations across structure lengths. Temperature 0.6 is selected as a trade-off between designability and diversity.
A.4.2 STRUCTURE GENERATION METRICS
- Figures 10–12 show held-out RMSD and TM-Score distributions for codebook sizes 432, 1,728, and 4,096 at downsampling ratio 1.Figure 9 tracks TM-Score distributions as downsampling changes; the supplied captions do not report outcomes.Source: Figure 10 · p. 20, Figure 9 · p. 20, Figure 11 · p. 21
- The designability score is the best of eight ProteinMPNN-sequence/ESMFold reconstructions, with designability defined by scTM>0.5 or scRMSD<2 Å.ProteinMPNN samples sequences at temperature 0.1; ESMFold folds them without recycling.Source: §A.4.2 · pp. 19–20
- Novelty is measured against the s40 CATH dataset, and a structure is novel when its maximum CATH TM-score is below 0.5.The procedure retrieves 1,000 Foldseek hits before TM-align comparison.Source: §A.4.2 · p. 20
- Diversity is the number of structural clusters normalized by sample count, using pairwise TM-score clustering at a threshold of 0.6.The clustering uses MaxCluster; more clusters indicate a more diverse sample set.Source: §A.4.2 · p. 20
A.4.3 SAMPLING
- Temperature sampling trades sample diversity against designability: higher temperatures broaden length distributions, while lower temperatures yield higher scTM scores for samples around lengths 100–300.The authors use temperature 0.6 for Table 2 as a satisfactory trade-off; Figure 18 presents length and designability by temperature.Source: §A.4.3 · pp. 21–22, Figure 18 · p. 24
- The model samples the joint length-and-structure distribution, unlike diffusion baselines modeling p(s|l); evaluation bins 40,000 generated structures by length and caps each bin at 10.The bins have width 5 and centers from 100 to 500; baseline tests sample eight backbones per length in that range.Source: §A.4.3 · pp. 21–23
A.4.4 ADDITIONAL RESULTS
- Additional figures examine designability by length, novelty by length and versus designability, novel-domain designability scores, and ProteinMPNN per-residue negative log-likelihood by sequence length.The captions identify these evaluation axes but do not state their outcomes.Source: Figure 19 · p. 24, Figure 20 · p. 25, Figure 21 · p. 25