[Paper Review] Learning the Language of Protein Structure

Benoit Gaujac, Jérémie Donà+4 more|arXiv (Cornell University)|May 24, 2024

Biomedical Text Mining and Ontologies5 citations

TL;DR

Introduces a vector-quantized autoencoder to discretize protein structures into token-like codes, enabling sequence-model based generation of designed protein backbones. Demonstrates competitive de novo structure generation with a GPT model trained on the learned codebooks.

ABSTRACT

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 Å. 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.

Motivation & Objective

• Convert 3D protein structures into a discrete sequential representation to enable sequence modeling.
• Learn a discrete, potentially low-dimensional latent space with a finite codebook for protein structures.
• Achieve high reconstruction fidelity (RMSD ~1-5 Å) while providing a scalable token vocabulary.
• Demonstrate downstream utility by training a GPT model on the learned codebooks to generate novel, designable protein structures.

Proposed method

• Encode backbone coordinates (N-Cα-C-O) with a 3-layer MPNN-based graph encoder producing embeddings for downsampled residues.
• Discrete the latent embeddings using Finite Scalar Quantization (FSQ) to form a codebook (up to 64k codes).
• Decode from quantized latent codes to 3D backbone coordinates using a structure module inspired by AlphaFold’s framing, optimized with Frame Align Point Error (FAPE) loss.
• Train on ~310k PDB entries with augmentation-like downsampling and rotation-invariant features to ensure frame alignment and robust reconstruction.
• Evaluate reconstruction via RMSD and TM-score, and demonstrate generative capability by training a decoder-only GPT on tokenized structures.
• Provide public code and compare generation quality against specialized diffusion models (FrameDiff, RFDiffusion).

Experimental results

Research questions

• RQ1 Can protein structures be effectively discretized into a finite token vocabulary without catastrophic information loss?
• RQ2 Does a discrete latent space enable effective use of standard sequence models (e.g., GPT) for de novo protein structure generation?
• RQ3 What are the reconstruction fidelity and designability trade-offs when varying codebook size and downsampling in a vector-quantized autoencoder?
• RQ4 How do generated structures compare to reference data in designability, novelty, and diversity?

Key findings

• A codebook of 64k with no downsampling (r=1) achieves ~1.59 Å RMSD and TM-score 0.95 on test structures.
• Quantization does not collapse the codebook; larger codebooks improve reconstruction while maintaining competitive fidelity.
• Downsampling and smaller codebooks increase reconstruction error but keep TM-scores above 0.5 on average, indicating preserved global folds.
• A simple decoder-only GPT trained on tokenized structures can generate novel, designable backbones with competitive self-consistency metrics compared to diffusion baselines.
• Generated structures exhibit non-trivial secondary structure elements and align well with ESM-predicted structures for designability assessments.
• Compared to FrameDiff and RFDiffusion, the GPT-based approach yields competitive designability and novelty/diversity profiles, with trade-offs depending on sampling strategy.