Adaptive Protein Tokenization

This paper introduces Adaptive Protein Tokenization, a global tokenization method that overcomes the limitations of existing local approaches by progressively adding detail to protein representations, thereby improving performance in generative, reconstruction, and classification tasks while enabling adaptive inference and novel applications like zero-shot protein shrinking.

The Gist

Imagine you are trying to describe a complex 3D sculpture to a friend over the phone so they can build an exact replica.

The Old Way (Local Tokenization):
Most previous AI models tried to do this by describing the sculpture piece-by-piece, like a mosaic. They would say, "Here is a piece of clay for the nose, here is a piece for the ear, here is a piece for the chin." They focused on small, local neighborhoods of the protein.

• The Problem: If you make a tiny mistake describing the nose, the whole head might end up crooked. Also, if the sculpture is huge (like a giant protein complex), you have to send thousands of tiny pieces, which is slow and inefficient. If you miss a piece, the whole picture falls apart.

The New Way (Adaptive Protein Tokenization):
The authors of this paper, Rohit Dilip and colleagues, propose a smarter way to describe the sculpture. Instead of sending thousands of tiny mosaic tiles, they send a series of "global snapshots" that get more detailed with each step.

Think of it like zooming in on a map:

1. Token 1: "This is a large, round object." (The big picture).
2. Token 2: "It has a long, curved shape." (Adding context).
3. Token 3: "It has a spiral on the left side." (Adding detail).
4. Token 4: "The spiral has a specific texture." (Adding fine detail).

This is called Adaptive Tokenization. The AI doesn't just look at one spot; every new token adds a layer of detail to the entire protein structure.

How It Works (The Magic Ingredients)

The paper introduces a tool called APT (Adaptive Protein Tokenizer). Here is how it functions using simple analogies:

• The "Coarse-to-Fine" Hierarchy: Imagine listening to a song. First, you hear the bass drum (the rhythm/global shape). Then you hear the melody. Finally, you hear the high-hat cymbals (the tiny details). APT works the same way. The first few tokens describe the overall shape of the protein. As you add more tokens, the AI fills in the fine details.
• The "Smart Stop" Button: In the old way, you had to send all the tiles to get a good picture. With APT, the AI can decide to stop sending tokens once it has enough information for the job at hand. If you just need to know the general shape of the protein, you can stop after 16 tokens. If you need to build a precise drug, you might use 64 tokens. This saves time and reduces errors.
• The "Noise-Canceling" Feature: Because the AI builds the protein from the "big picture" down to the details, it is less likely to make a mistake that ruins the whole thing. If a small detail is slightly off, the overall shape remains correct.

What They Proved (The Results)

The team tested this new method against the current best models (like ESM3 and DPLM2) in three main ways:

1. Rebuilding (Reconstruction): When they tried to rebuild a protein from its tokens, APT was just as accurate as the best existing models, but it could do it with fewer tokens.
2. Creating New Proteins (Generation): They asked the AI to invent brand new proteins. The new method created proteins that were more "designable" (meaning they could be physically built in a lab without falling apart) than the old methods. In fact, it achieved a success rate of about 87%, beating the previous leaders significantly.
3. Understanding Function (Representation): They tested if the AI could understand what a protein does just by looking at its token description. They found that APT's "global view" was better at classifying protein types than models that only looked at local pieces.

Cool New Tricks Enabled by This Method

Because the AI separates the size of the protein from the description of its shape, the authors demonstrated two specific "magic tricks":

• Protein Shrinking: You can take a large protein and tell the AI, "Keep the same shape, but make it smaller." The AI successfully shrank proteins (like hemoglobin) while keeping their core structure intact. This is useful for making drugs that can fit into cells more easily.
• Affinity Maturation (Making Things Stickier): If you have a protein that barely sticks to a target (a "weak binder"), you can use the AI to generate variations that stick better. The AI can "search" through possibilities to find a version that works much better, essentially evolving the protein in seconds.

The Bottom Line

This paper presents a new way for computers to "speak" about protein shapes. Instead of describing them as a pile of tiny, fragile bricks, it describes them as a series of evolving blueprints. This makes the AI faster, more accurate, and better at creating new, usable proteins for science and medicine.

Note: The paper focuses on the computational method and the ability to generate and shrink protein structures. It does not claim these proteins are currently being used in human clinical trials or as approved medicines, but rather that the method makes them "designable" (capable of being built and tested).

Technical Summary

Technical Summary: Adaptive Protein Tokenization

Problem Statement

Current approaches to protein structure tokenization rely on locality, where tokens are created by pooling information from spatially neighboring residues along a protein sequence. While these methods achieve high-fidelity reconstruction, the authors identify two critical limitations:

1. Error Accumulation in Generation: Generative models based on local tokenization suffer from error propagation. A single missampled token can lead to significant deviations in the predicted structure, causing discrete approaches to underperform compared to continuous diffusion models on generative metrics.
2. Poor Scalability: The number of tokens scales linearly with protein size. This makes modeling large protein complexes computationally expensive and limits the ability to compress proteins effectively.
3. Representation Limitations: Local tokenizers often require mean-pooling operations to generate fixed-size representations for downstream tasks, which can discard critical global information and scale poorly with protein size.

Methodology

The authors propose Adaptive Protein Tokenization (APT), a global tokenization method where successive tokens contribute increasing levels of detail to a global representation, rather than corresponding to specific local neighborhoods.

Architecture

APT is implemented as a diffusion autoencoder with a discrete bottleneck:

• Input: A sequence of raw $C_\alpha$ coordinates (normalized to zero-mass) of length $L$.
• Encoder: A bidirectional attention transformer maps raw positions to a latent sequence $c \in \mathbb{R}^{L \times d}$.
• Discretization: The latent sequence is discretized using Finite Scalar Quantization (FSQ) with levels $(8, 5, 5, 5)$, resulting in an effective codebook size of 1,000.
• Adaptivity Mechanism: To enforce adaptivity, the model employs nested dropout. During training, an upper cutoff $U$ is uniformly sampled from $[1, \min(L, k_{max})]$. Tokens beyond this cutoff are dropped. This encourages the model to place critical global information (low-frequency components) in the first few tokens and higher-frequency details in subsequent tokens.
• Decoder: A diffusion decoder, trained using a flow-matching objective, reconstructs the atomic coordinates conditioned on the discrete tokens. The decoder utilizes stochastic equivariance, learning symmetries rather than enforcing them explicitly.
• Size Prediction: The protein size is regressed from the first token using a cross-entropy loss, allowing the model to decouple protein size from the conditioning sequence length.

Training and Inference

• Training: The model is trained on $\approx 473,000$ synthetic AlphaFold2 predictions from the Foldseek-clustered AFDB database. Training involves two stages: (1) training the autoencoder with random rotations to learn stochastic equivariance, and (2) training an autoregressive (GPT-style) model over the APT tokens.
• Inference Strategies:
  • Tail Dropout: During generation, the model can stop generating tokens early. Any prefix of the token sequence is a valid conditioning signal. The tail (high-frequency details) is delegated to the diffusion decoder.
  • Entropy-Based Stopping: The authors introduce stopping criteria based on token entropy (finite cutoff, spline-based minimum entropy) to determine the optimal point to halt generation, balancing representation fidelity and generative capacity.
  • Classifier Annealing: To mitigate artifacts caused by shifting away from the true data manifold, the authors employ classifier annealing, interpolating between conditional and unconditional guidance fields during decoding.

Key Contributions

1. Global Adaptive Tokenizer: Introduction of a tokenizer where tokens represent global descriptors in a coarse-to-fine hierarchy, rather than local neighborhoods.
2. Performance Validation: Demonstration that the model matches or outperforms state-of-the-art (SOTA) local tokenizers (e.g., DPLM2, ESM3, Kanzi) on generative, representation learning, and reconstruction tasks.
3. Inference-Time Techniques: Development of entropy-based sampling and classifier annealing to manage sample complexity and mitigate error exposure, leading to improved designability.
4. Zero-Shot Applications: Demonstration of decoupling protein size from conditioning, enabling zero-shot applications such as protein shrinking (reducing residue count while maintaining structure) and affinity maturation via inference-time scaling.

Results

• Reconstruction: On CATH, CAMEO, and AFDB test sets, APT achieves RMSD and TM-scores comparable to SOTA continuous diffusion models. Notably, using only 32–64 tokens yields RMSDs $< 2$ Å, affirming the compressibility of protein structures.
• Generation: An autoregressive model trained on APT tokens achieves a designability of 0.871 (fraction of samples with scRMSD $< 2$ Å), significantly outperforming discrete diffusion models like DPLM2 (0.486) and ESM3. This is achieved without "best-of-N" sampling techniques.
• Representation Learning: On the CATH classification task, non-linear (MLP) probing on APT tokens outperforms probing on representations from ESM3 and DPLM2. Crucially, APT provides fixed-size global representations without mean-pooling, and even highly compressed representations (16 tokens) outperform larger models.
• Applications:
  • Protein Shrinking: Successfully generated smaller versions of proteins (e.g., hemoglobin, beta-barrel) with preserved global and local structure by conditioning the diffusion process on fewer residues.
  • Affinity Maturation: Demonstrated the ability to evolve weak binders into strong binders using beam search guided by reward functions (e.g., iPAE) in the latent space.

Significance

The paper claims that Adaptive Protein Tokenization provides a task-aware approach to scaling biological tasks. By moving from local to global tokenization, the method resolves the trade-off between reconstruction fidelity and generative capacity. The ability to compress proteins into fixed-size vectors without information loss via mean-pooling offers a practical pathway for modeling large protein complexes and performing inference-time scaling. The authors position this work as a step toward generative models capable of reasoning across multiple length scales, addressing a key challenge in frontier bioengineering.

Note: The authors acknowledge limitations, specifically that APT tokens are effective for global tasks but less suitable for local tasks like motif scaffolding, and that the current tokenizer does not account for sidechains or amino acid composition, which are critical for certain functions.