Adaptive Protein Tokenizer (APT)

• Adaptive Protein Tokenizer (APT) is a hierarchical protein tokenization framework that encodes global structure using an adaptive, coarse-to-fine sequence of tokens.
• It leverages a bidirectional Transformer encoder, finite scalar quantization, and nested dropout to allocate global features and fine-grained details efficiently.
• Empirical results demonstrate that APT achieves superior reconstruction and generative performance while enabling efficient representation learning and zero-shot optimization.

Adaptive Protein Tokenizer (APT) is a protein-structure tokenization method that replaces local neighborhood tokenization with a global, coarse-to-fine representation in which successive tokens contribute increasing levels of detail to a single global representation. Introduced in "Adaptive Protein Tokenization" (Dilip et al., 6 Feb 2026), APT is motivated by the observation that existing protein structure tokenizers typically assign one discrete token per local spatial neighborhood, a design that preserves fine-scale geometry but incurs linear scaling with protein length, error accumulation in generative models, and sequence-reduction requirements for global tasks. APT encodes a protein structure into a small ordered set of tokens $\{c_1,c_2,\dots,c_K\}$ such that good reconstructions and generative sampling can be achieved with $K\!\ll\!L$, while any prefix $\{c_1,\dots,c_k\}$ remains a valid global representation (Dilip et al., 6 Feb 2026).

1. Conceptual shift from local pooling to global tokenization

Contemporary protein-structure tokenizers are described as relying on graph neural networks or diffusion transformers that assign one discrete token per local spatial neighborhood, such as a sliding window or a residue-centric subgraph (Dilip et al., 6 Feb 2026). The reported limitations are threefold. First, a protein of length $L$ residues produces $O(L)$ tokens, so compressing or jointly modeling large assemblies becomes prohibitively expensive. Second, if an autoregressive or discrete diffusion model mis-samples even one local token, downstream reconstruction can deviate dramatically because subsequent local neighborhoods no longer align to the true backbone, leading to cascading errors. Third, global prediction tasks such as stability, function, and classification require pooling or aggregating $O(L)$ residue-level embeddings into a single vector of fixed dimension, a process that can destroy spatial information and impose arbitrary, length-dependent biases (Dilip et al., 6 Feb 2026).

APT’s alternative is a global tokenization paradigm in which the token sequence is ordered by information content. Early tokens encode low-frequency, global structure; later tokens refine the representation with progressively higher frequency detail. The paper explicitly connects this hierarchy to multiresolution signal transforms, citing Fourier and wavelets as inspiration (Dilip et al., 6 Feb 2026).

This suggests that APT should be understood not merely as a compression mechanism, but as a representational reparameterization of protein structure. A plausible implication is that the token index itself becomes semantically meaningful: earlier positions correspond to coarse global descriptors, whereas later positions concentrate fine structural corrections.

2. Mathematical formulation and adaptive information content

Given a protein backbone as an $L\times3$ array of C$\alpha$ coordinates $X\in\mathbb{R}^{L\times3}$, APT first applies a bidirectional Transformer encoder $E_\phi$ to mean-centered coordinates with positional encoding:

$K\!\ll\!L$0

The latent/channel dimension is given as, for example, $K\!\ll\!L$1 (Dilip et al., 6 Feb 2026).

The resulting continuous latents are quantized with finite scalar quantization (FSQ). The paper specifies $K\!\ll\!L$2 product-quantization levels with sizes $K\!\ll\!L$3, yielding an effective codebook of size $K\!\ll\!L$4:

$K\!\ll\!L$5

To force the hierarchy to be coarse-to-fine, APT uses nested dropout. During training, an integer cutoff $K\!\ll\!L$6 is sampled, and all tokens $K\!\ll\!L$7 are zeroed out. The decoder must then reconstruct $K\!\ll\!L$8 using only the first $K\!\ll\!L$9 quantized tokens, which drives the model to allocate global features into earlier positions and fine details into later positions (Dilip et al., 6 Feb 2026).

At inference time, the autoregressive prior predicts

$\{c_1,\dots,c_k\}$0

and APT computes per-token Shannon entropy

$\{c_1,\dots,c_k\}$1

This entropy curve serves as an adaptive measure of how much additional information remains at each level of the hierarchy. The paper defines two stopping rules:

$\{c_1,\dots,c_k\}$2

and, after fitting a smooth spline $\{c_1,\dots,c_k\}$3,

$\{c_1,\dots,c_k\}$4

By dropping tokens $\{c_1,\dots,c_k\}$5, APT trades representational fidelity for reduced generative complexity, with fewer sequential decisions and fewer opportunities for mis-sampling (Dilip et al., 6 Feb 2026).

3. Architecture and training objectives

APT combines a Transformer encoder, FSQ discretization, and a diffusion decoder. The encoder is specified as a bidirectional Transformer with 2 layers, $\{c_1,\dots,c_k\}$6 channels, 4 attention heads, relative positional encodings, and layer normalization. The quantizer uses four scalar quantization levels with sizes $\{c_1,\dots,c_k\}$7 arranged in a product code. The decoder is a diffusion transformer (“DiT”) with 8 layers and $\{c_1,\dots,c_k\}$8 that predicts the flow field $\{c_1,\dots,c_k\}$9 under a v-prediction (flow-matching) objective. The model shares adaptive layer-norm parameters across layers and uses stochastic symmetry via random rotations instead of enforcing SE(3) equivariance in the encoder (Dilip et al., 6 Feb 2026).

The high-level Stage 1 training loop includes random rotation and mean-centering of $L$0, continuous encoding, FSQ quantization, nested dropout, a regression target for true size $L$1, a flow loss, a size loss, and optimization of the combined objective:

$L$2

$L$3

$L$4

For the autoregressive prior, the paper gives the standard cross-entropy objective

$L$5

Stage 2 generative sampling uses an autoregressive Transformer prior, entropy-based stopping, optional size regression from $L$6, diffusion ODE/SDE decoding, and classifier annealing:

$L$7

The output is a sampled structure $L$8 (Dilip et al., 6 Feb 2026).

A central architectural consequence is the elimination of explicit sequence reduction. Because each token is a global descriptor, any prefix of length $L$9 suffices to represent the whole protein. For downstream tasks, one picks a fixed $O(L)$0 such as 16, 32, or 64 and directly uses the flattened token indices as a $O(L)$1 one-hot or embedding vector, with no pooling over residues (Dilip et al., 6 Feb 2026).

4. Reconstruction and generative performance

The empirical validation covers reconstruction, generative design, and representation learning. On reconstruction, the paper reports that with the full tail, $O(L)$2, APT matches or slightly outperforms leading continuous and discrete tokenizers. The reported results are:

Dataset	Metric
CATH test	RMSD = 0.90 Å, TMscore = 0.941
CAMEO test	RMSD = 0.90 Å, TM = 0.941
AFDB test	RMSD = 1.17 Å, TM = 0.929

For CATH test, the comparison stated in the paper is “vs. Kanzi 1.09 Å/0.937” (Dilip et al., 6 Feb 2026). The same section reports that truncating to 32–64 tokens yields RMSD $O(L)$3 Å and TMscore $O(L)$4, which the paper characterizes as sufficient for downstream diffusion generation. Reconstruction Fréchet distance (rFID) is reported to improve monotonically with $O(L)$5 (Dilip et al., 6 Feb 2026).

On generative quality, designability is measured as the fraction of samples with self-consistent RMSD $O(L)$6 Å after ProteinMPNN $O(L)$7 ESMFold. The reported values are:

Method	Designability
Baseline discrete diffusion (DPLM2)	0.486
APT-AR with 16 tokens	0.482
APT-AR + classifier annealing ($O(L)$8)	0.771
APT-AR + finite-entropy sampling & guidance	0.871

The paper additionally reports that scRMSD drops from approximately 3 Å for 16 tokens without guidance to approximately 1.3 Å with entropy cutoff plus guidance. Diversity, Novelty, and generation speed are described as competitive with or superior to prior discrete models (Dilip et al., 6 Feb 2026).

These results support the paper’s claim that adaptive stopping based on information content is not merely a compression heuristic. In the reported experiments, entropy-based truncation directly affects designability and structural self-consistency.

5. Representation learning and CATH probing

APT is also evaluated as a representation model. After tokenizer training, the tokenizer is frozen and small heads, either linear or MLP, are trained on top of the first $O(L)$9 tokens to predict CATH labels via cross-entropy (Dilip et al., 6 Feb 2026). The reported outcome is that fixed-token probing without mean pooling yields 1$O(L)$0–2$O(L)$1 better top-1 accuracy than DPLM2 and ESM3 of comparable scale. An MLP head on 16–64 tokens is reported to reach 60–75% top-1 accuracy at the CATH topology (T) level, outperforming much larger local models (Dilip et al., 6 Feb 2026).

The significance of these results lies in the representational format. Local tokenizers generally require residue-level aggregation for global classification; APT uses a fixed-length prefix of global descriptors directly. This suggests that some of the gains are attributable not only to token quality, but also to avoiding sequence-length reduction operations that can obscure spatial organization.

A plausible implication is that APT occupies an intermediate regime between compression and task representation: the same token sequence can be decoded generatively or consumed discriminatively with little modification. The paper presents this dual-use property as a consequence of the global-descriptor design (Dilip et al., 6 Feb 2026).

6. Zero-shot shrinking and affinity maturation

The paper includes two case studies that emphasize inference-time manipulation of token sequences rather than retraining. In protein shrinking, a natural protein such as human hemoglobin of 141 AA is tokenized with APT, decoded via diffusion while forcing a smaller output length $O(L)$2, and then redecoded and aligned to the original structure for TMscore evaluation. The reported results are 0.94 TMscore at 90% length (126 AA) and 0.70 TMscore at 80% length (112 AA). The paper states that structure and global fold remain largely intact even with 20% residue deletion, and frames this as a zero-shot path to rational miniaturization (Dilip et al., 6 Feb 2026).

For affinity maturation via inference-time scaling, APT-AR is combined with beam search and two reward formulations. The first is a CATH classifier reward trained directly on APT tokens, with beam search maximizing

$O(L)$3

The paper states that samplers adapt trajectories to produce, for example, “all-$O(L)$4” designs. The second is an iPAE binder-affinity reward, in which a candidate sequence is co-folded with a target via AlphaFold-Multimer after sequence generation with ProteinMPNN, and the reward is defined as

$O(L)$5

The paper reports that beam search improves predicted affinity over 20 diffusion steps (Dilip et al., 6 Feb 2026).

Because APT tokens remain valid at any prefix length, search can be carried out with few autoregressive steps and truncated decodings. The paper presents this as enabling fast, zero-shot optimization of complex, global biochemical objectives (Dilip et al., 6 Feb 2026).

7. Interpretation, scope, and relation to prior tokenization practice

APT is presented as a direct response to the limitations of local protein tokenization rather than as a generic tokenizer redesign. Its defining claim is that compression should be decoupled from protein length and organized by a coarse-to-fine ordering of global information (Dilip et al., 6 Feb 2026). In this framing, the central innovation is not only the use of FSQ, a Transformer encoder, or a DiT decoder, but the combination of nested dropout and entropy-aware truncation to make token prefixes semantically complete.

The paper explicitly contrasts this with tokenizers based on local pooling. Local tokenizers preserve fine-scale geometry but inherit linear token growth, error accumulation under autoregressive or discrete diffusion sampling, and a dependence on downstream pooling for global tasks. APT addresses these issues by front-loading global information, allowing any prefix to function as a valid global representation (Dilip et al., 6 Feb 2026).

A common misconception would be to interpret APT as merely a shorter token sequence. The paper’s reported behavior is broader than token-count reduction alone: it includes information-content-based stopping rules, fixed-token probing without mean pooling, entropy-guided gains in designability, and inference-time operations such as forced-length shrinking and reward-guided beam search (Dilip et al., 6 Feb 2026). This suggests that APT is best viewed as a hierarchical tokenization framework whose adaptive prefixes unify reconstruction, generation, and representation learning within a single discrete interface.