AIDO.StructureTokenizer

AIDO.StructureTokenizer, abbreviated AIDO.St, converts a protein's three-dimensional structure into a sequence of discrete tokens and converts those tokens back into atomic coordinates. It was developed by GenBio AI with MBZUAI and Carnegie Mellon University, and posted to bioRxiv in December 2024 as the structural module of the AIDO family of foundation models — distinct from the sequence language model AIDO.Protein, which is described in separate preprints.

Tokenizing structure solves a specific problem. Databases such as the AlphaFold Protein Structure Database hold hundreds of millions of predicted structures, and searching, indexing, or generating over them is awkward with continuous coordinates. Turning a backbone into a short string of symbols makes structure searchable with text-retrieval machinery and, more importantly, consumable by the transformer stacks already trained on amino acid sequence.

The paper's contribution is not the idea of a structure tokenizer — it explicitly compares against Foldseek, ProToken, and ESM-3 — but the trade-off it characterizes. The authors find that an encoder producing highly local tokens retrieves homologs well but reconstructs structures poorly, while an encoder carrying more global context reconstructs well and retrieves badly. Their central result is that balancing the two, rather than maximizing either, yields tokens that align best with a protein sequence language model and give the highest downstream structure prediction accuracy.

#Key Features

• Equivariant encoder, invariant decoder: The encoder is an Equiformer that respects rotation and translation of the input backbone; the decoder reconstructs coordinates from tokens alone with an architecture adapted from ESMFold, so it needs no reference frame.
• 512-entry codebook: A vector-quantized bottleneck maps continuous latents to one of 512 discrete structural tokens, giving a fixed, small alphabet that a sequence model can be trained over directly.
• All-atom reconstruction: The decoder recovers side chains, not just the backbone it was given — the tokens carry enough information to place full residues.
• Structure retrieval and homology detection: Token similarity supports searching structure databases, benchmarked against Foldseek, ProToken, and ESM-3.
• Separately released encoder: The encoder is published as its own small checkpoint, so a workflow that only needs tokens — for indexing or for feeding a sequence model — can skip the reconstruction stack entirely.

#Technical Details

The model is a VQ-VAE with three parts. The equivariant encoder is roughly 6 million parameters and maps amino-acid-agnostic backbone coordinates into a latent space that preserves rotational and translational symmetry. A discrete codebook quantizes those latents into 512 structural tokens. The invariant decoder carries the bulk of the roughly 300 million total parameters and reconstructs full atomic structures — including side chains — from the token string, using an architecture derived from ESMFold's structure module.

Evaluation covers three axes. Reconstruction fidelity is measured on a preprocessed CASP15 dataset, released as a Hugging Face dataset alongside the checkpoints, with reconstructed structures Kabsch-aligned to their inputs for comparison. Retrieval and homology detection are benchmarked against Foldseek, ProToken, and ESM-3. Structure prediction is the downstream test that ties the two together: the authors align structure tokens with a protein sequence language model and show that the encoder configuration striking the best balance between locality and reconstruction also produces the most accurate predicted structures, which is the paper's argument for why neither objective should be maximized alone.

Both checkpoints ship through ModelGenerator, GenBio's fine-tuning and inference framework, with encoding and decoding driven by mgen predict against a YAML configuration and a CSV registry of input PDB or mmCIF files. Its experiments/GB.StructureTokenizer directory holds the encode and decode configs, a script that extracts the 512-entry codebook, and a notebook predicting structure tokens from sequence with a 16B protein model. One naming detail matters for finding the artifacts: GenBio has renamed its Hugging Face checkpoints from the AIDO.* prefix used in the paper to GB.*, so the tokenizer and encoder now live at genbio-ai/GB.StructureTokenizer and genbio-ai/GB.StructureEncoder.

#Applications

The tokenizer is infrastructure rather than an end-user tool. Its uses are indexing and searching large structure collections without pairwise structural alignment, supplying a discrete structural channel to multimodal protein models that already handle sequence, and serving as the decoder in structure-prediction pipelines where a sequence model emits structure tokens that are then rendered into coordinates. Groups building on the AIDO stack get it as a drop-in module; groups outside it get a standalone way to move between coordinates and tokens.

#Impact

The most transferable result here is the locality–reconstruction trade-off itself, stated and measured rather than assumed. It gives designers of future structure tokenizers a concrete axis to tune, and an argument that the right target is alignment with a sequence model rather than reconstruction accuracy in isolation. Practical caveats: the preprint carries a CC-BY-NC-ND licence and was presented at the NeurIPS 2024 MLSB workshop rather than in a peer-reviewed journal, the checkpoints are released under GenBio's own community licence rather than a standard open-source one, and decoding structures is slow enough that the authors flag it, with multi-GPU inference not yet supported.