Skip to content

Protein Structure Tokenization: Benchmarking and New Recipe

Conference: ICML 2025 arXiv: 2503.00089 Code: GitHub Area: Protein Structure / Tokenization Evaluation Keywords: Protein Structure Tokenization, VQ-VAE, codebook collapse, Benchmarking, ESM3

TL;DR

Proposes StructTokenBench—the first comprehensive evaluation framework for protein structure tokenizers (PSTs). It assesses existing methods across four dimensions: downstream effectiveness, sensitivity, distinctiveness, and codebook utilization efficiency. It further introduces AminoAseed, a strategy that significantly improves the quality of VQ-VAE-type tokenizers through codebook reparameterization and Pareto-optimal configurations (achieving a 6.31% improvement and a 124% increase in utilization compared to ESM3).

Background & Motivation

Protein structure tokenization (PST) encodes 3D protein structures into discrete or continuous representations, serving as the foundation for protein language modeling and multimodal models. Existing PST methods fall into two categories:

VQ-VAE-type: FoldSeek, ProTokens, ESM3—compressed into a discrete codebook via reconstruction objectives.

Inverse folding (IF): ProteinMPNN, MIF—compressed by predicting sequences that can fold into target structures.

However, the capabilities and limitations of these methods remain unclear due to the lack of a unified evaluation framework. In particular, the codebook collapse issue is severe—up to 70% of the 4096 codes in ESM3 are inactive during inference.

Method

StructTokenBench Evaluation Framework

Evaluates PST quality across four complementary dimensions:

1. Downstream Effectiveness: 12 supervised tasks (24 test splits), including binding site prediction, catalytic site prediction, conserved site prediction, repeat motif prediction, epitope prediction, structural flexibility prediction, and remote homology detection.

2. Sensitivity: Detects subtle differences between highly similar protein conformations, measured by the correlation between representation similarity and TM-score.

3. Distinctiveness: The diversity between codebook vector pairs, avoiding redundant token-to-substructure mappings.

4. Codebook Utilization Efficiency: Utilization Rate (UR), Perplexity, and Marginal Vocabulary Utility (MUV).

AminoAseed Method

1. Codebook Reparameterization

Applies a learnable linear transformation to fixed orthogonal vector bases:

\[\mathbf{Q} = \text{Linear}(\mathbf{C})\]

where \(\mathbf{C}\) is randomly initialized as approximately orthogonal and fixed during training. Unlike the vanilla VQ-VAE, all parameters of the linear layer receive gradient updates, avoiding distribution drift for unselected codes.

2. Pareto-optimal Codebook Configuration

Balances codebook size \(K\) and dimension \(D\) under total capacity \(K \times D\) constraints: - \(K > 2^{10}\): Downstream performance decreases. - \(K < 2^8\): Insufficient expressiveness. - \(K = 2^9 = 512\): The optimal balance between utilization and performance.

VQ-VAE Optimization Objective

\[\mathcal{L} = \log p(\tilde{\mathbf{x}}|\mathbf{q}_k) + \|sg(\mathbf{z}) - \mathbf{q}_k\|_2^2 + \beta\|\mathbf{z} - sg(\mathbf{q}_k)\|_2^2\]

Key Experimental Results

Downstream Effectiveness (Average AUROC %)

Method Type Average AUROC
MIF IF 79.82
ProteinMPNN IF 75.92
ESM3 VQ-VAE 69.24
AminoAseed VQ-VAE 72.43 (+4.74%)
VanillaVQ VQ-VAE 68.30 (-0.86%)

AminoAseed vs ESM3 in 24 Test Splits

Across the 24 supervised task test splits, AminoAseed yields an average improvement of 6.31%, with improvements up to 17.33% in certain tasks (e.g., CatBio-SupFam).

Codebook Utilization Efficiency

Model Utilization Rate UR (%)
FoldSeek Highest
ESM3 ~30% (approx. 1200 active out of 4096)
AminoAseed 2.24× of ESM3 (124% increase)

Key Ablation Findings

  1. Vector quantization affects model expressiveness primarily due to optimization challenges, rather than differences between discrete and continuous formats.
  2. Structural tokens retain most of the information of amino acid tokens but are more sensitive to noise.
  3. Reconstruction quality correlates inconsistently with codebook quality—both require separate evaluations.
  4. Scaling up the VQ-VAE encoder yields sub-exponentially diminishing returns.

Highlights & Insights

  1. First Unified Benchmark for Protein Structure Tokenization: 4 dimensions + 12 tasks + 24 splits, filling a critical evaluation gap.
  2. Thorough Root Cause Analysis of Codebook Collapse: Unselected codes during training are not updated \(\rightarrow\) distribution drift \(\rightarrow\) further non-selection (positive feedback loop).
  3. Simple and Effective Reparameterization Solution: A single linear layer ensures all codes participate in gradient updates.
  4. Biological Plausibility of \(K=512\): Reconciles with the findings of heuristic tertiary structure methods like TERMs (~600 substructures describe 50% of the PDB).
  5. Complementarity of IF vs VQ-VAE Types: No single method dominates, each has its strengths.

Limitations & Future Work

  1. AminoAseed performs worse on physicochemical property prediction (FlexBFactor), indicating that the relationship between codebook quality and downstream tasks is not always positive.
  2. Evaluation focuses on backbone structures as input, not covering all-atom structures.
  3. Pre-trained only on 10% of filtered PDB data, which might limit the model's upper bound.
  4. Codebook utilization evaluation does not apply to IF-type methods (due to continuous representations).
  • VQ-VAE PST: FoldSeek, ProTokens, ESM3
  • IF PST: ProteinMPNN, MIF
  • Codebook collapse: EMA update, k-means initialization
  • Protein Benchmarks: TAPE, ProteinGLUE

Rating

⭐⭐⭐⭐ (4/5)

The design of the benchmarking framework is outstanding, with a comprehensive and scientifically sound four-dimension evaluation. Although the AminoAseed method is simple, its effectiveness is substantial. The ablation and scaling studies are deep, providing a wealth of valuable insights. The only minor drawback is that AminoAseed has yet to close the gap between VQ-VAE and IF methods.