AminoAseed

AminoAseed is a VQ-VAE-based protein structure tokenizer that converts 3D backbone geometry into a sequence of discrete tokens. Discrete structure tokens matter because they let structural information be consumed by the same machinery that works on text and protein sequence — autoregressive language modeling, masked prediction, and multimodal models that jointly reason over structure, sequence, and functional text. AminoAseed was developed by Xinyu Yuan, Zichen Wang, Marcus Collins, and Huzefa Rangwala; the work was performed at Amazon, with Yuan also affiliated with Mila and the Université de Montréal.

AminoAseed arrived alongside StructTokenBench, an evaluation framework introduced in the same paper. The two are distinct artifacts: StructTokenBench is a benchmark that measures the quality and efficiency of structure tokenizers with an emphasis on fine-grained local substructures rather than global fold, while AminoAseed is the tokenizer the authors built after that benchmark exposed a specific failure mode.

That failure mode is codebook under-utilization. Existing structure tokenizers train a discrete codebook of structural "words," but in practice a large fraction of codebook entries receive negligible gradient signal and go unused, so the effective vocabulary is far smaller than the nominal one. AminoAseed is the recipe that fixes this, and the resulting tokenizer outperforms the structure tokenizer inside ESM-3, the strongest prior VQ-VAE structure tokenizer, across the benchmark's downstream task suite.

#Key Features

• Codebook gradient rescaling: The training objective is modified so that codebook entries receive stronger and better-distributed gradient updates, pulling entries that would otherwise stay dormant into active use.
• Size-dimension balance: Rather than scaling the vocabulary indefinitely, AminoAseed jointly tunes codebook size against embedding dimension, trading nominal capacity for capacity that is actually reachable during training.
• Local substructure focus: The tokenizer is built and evaluated around fine-grained local geometry — the level at which residues actually interact — rather than global fold similarity, which is what most earlier structure benchmarks measured.
• Drop-in frozen encoder: Once trained, the tokenizer encodes arbitrary new protein structures into fixed discrete tokens with no per-task retraining, so downstream models can treat structure as just another token stream.
• Paired with an open benchmark: StructTokenBench ships in the same repository, so competing tokenizers can be compared on identical tasks and metrics rather than on self-reported numbers.

#Technical Details

AminoAseed builds on ESM-3's local-frame paradigm, using geometric self-attention layers that are invariant to rotation and translation of the input structure. The released configuration uses a codebook of 512 entries at dimension 1024. Training used 48,316 protein chains from the RCSB Protein Data Bank filtered to under 40% sequence identity — a 10% downsample following OpenFold2 filtering criteria — split 90/10 between training and validation and run for 108,530 steps on 8 NVIDIA A100 GPUs with the Adam optimizer (learning rate 1e-4, weight decay 0.01).

Evaluated on StructTokenBench's 24 supervised downstream tasks, AminoAseed averages a 6.31% performance improvement over ESM-3's structure tokenizer. Two diagnostic metrics explain where that gain comes from: sensitivity — how well the tokenizer distinguishes structurally distinct local environments — rose 12.83%, and codebook utilization rose 124.03%, more than doubling the fraction of the vocabulary that carries information. The paper's broader benchmarking result is that no single existing tokenizer dominates across all evaluation axes, which is why the authors report utilization and sensitivity separately rather than collapsing to one score.

#Applications

Structure tokenizers are infrastructure rather than end-user tools, and AminoAseed is most useful to teams building multimodal protein models. Encoding structures into discrete tokens lets a language model consume backbone geometry directly, which is the mechanism behind structure-aware generation, structure-conditioned function prediction, and joint sequence-structure-text models. Groups fine-tuning downstream predictors — binding-site identification, functional-site annotation, conservation prediction — can substitute AminoAseed tokens for a structure encoder and inherit the improved local sensitivity without retraining the tokenizer. StructTokenBench itself gives method developers a standard yardstick for new tokenization schemes.

#Impact

The paper's contribution is as much diagnostic as constructive: by showing that codebook under-utilization is a systematic weakness across structure tokenizers, it reframed a design choice that had largely been treated as a hyperparameter. AminoAseed demonstrates that a straightforward change to gradient flow and codebook geometry recovers substantial downstream performance without a larger model or more data. Code and weights are released under an MIT license, and the accompanying benchmark makes the comparison reproducible. The main caveat is scope: AminoAseed is a tokenizer, not a generative or predictive model on its own, and its reported gains are measured through downstream probes rather than end-to-end applications.