From Static Structures to Ensembles: Studying and Harnessing Protein Structure Tokenization

TL;DR

This work tackles the problem of unifying protein sequence and structure in a generative framework by using discrete structure tokens learned via a VQ-VAE. A GPT-style model predicts structure tokens conditioned on the amino-acid sequence, and the authors demonstrate that success hinges on using rich sequence embeddings, with ESM3 outperforming ProGen2. They uncover substantial semantic redundancy in the ESM3 codebook, where distinct tokens map to nearly identical local geometries, enabling a synonym dictionary and a synonym-swap perturbation to generate realistic conformational ensembles that correlate strongly with MD-derived dynamics (e.g., $r=0.84$ for RMSF and $W_2$ distance of $1.83$). This leads to a near-instantaneous, training-free method to sample protein dynamics, offering a scalable alternative for modeling flexibility in design and analysis. The approach highlights fundamental properties of discrete structural representations and opens avenues for fast, multimodal protein modeling.

Abstract

Protein structure tokenization converts 3D structures into discrete or vectorized representations, enabling the integration of structural and sequence data. Despite many recent works on structure tokenization, the properties of the underlying discrete representations are not well understood. In this work, we first demonstrate that the successful utilization of structural tokens in a language model for structure prediction depends on using rich, pre-trained sequence embeddings to bridge the semantic gap between the sequence and structural "language". The analysis of the structural vocabulary itself then reveals significant semantic redundancy, where multiple distinct tokens correspond to nearly identical local geometries, acting as "structural synonyms". This redundancy, rather than being a flaw, can be exploited with a simple "synonym swap" strategy to generate diverse conformational ensembles by perturbing a predicted structure with its structural synonyms. This computationally lightweight method accurately recapitulates protein flexibility, performing competitively with state-of-the-art models. Our study provides fundamental insights into the nature of discrete protein structure representations and introduces a powerful, near-instantaneous method for modeling protein dynamics. Source code is available in https://github.com/IDEA-XL/TokenMD.

Figures (5)

• Figure 1: The VQ-VAE (left) discretizes continuous protein structures into a finite set of "Structure Tokens". These tokens are then used in an autoregressive language model (right) that predicts a sequence of structural tokens conditioned on the amino acid sequence.
• Figure 2: The training curve of the GPT model for protein structure prediction with different sequence embeddings.
• Figure 3: t-SNE visualization and the distance matrix of the code vectors. The ESM3 codebook shows distinct, well-defined clusters, while the t-SNE of the AIDO.st codebook zhang2024balancing vectors are uniformly distributed in the 2D space.
• Figure 4: Perturbation of the structure tokens for exploring the conformational ensemble space.
• Figure 5: Protein ensembles for 6uof_A generated by token perturbation and MD, and the C$\alpha$ RMSFs indexed by the residue id (Pearson $r=0.81$).