Bio2Token

A structure tokenizer that emits one token per residue has already discarded the thing most binding questions turn on: where the side-chain atoms actually sit. The usual response is to tokenize coarsely and rebuild atoms afterwards — ESM-3 decodes all-atom coordinates from residue-level codes, and backbone-only tokenizers pair with a separate inverse-folding stage — but each hand-off is another place for geometric error to enter. Tokenizing at atomic resolution directly is the obvious fix and carries an equally obvious obstacle: a mid-sized ribosomal RNA runs to tens of thousands of heavy atoms, well past what a quadratic-attention encoder holds on ordinary hardware.

Bio2Token, released by Flagship Pioneering in October 2024, gets there by dropping attention entirely. Its quantized autoencoder is built from bidirectional Mamba blocks — a selective state space model whose cost grows linearly in sequence length — so the encoder can accept a whole complex as one point cloud of heavy atoms, assign a discrete token to every atom, and decode the tokens back to coordinates. There is no SE(3)-invariant featurization and no invariant point attention anywhere in the architecture; the model is instead trained under random rotation augmentation and learns rotational robustness rather than having it imposed.

The input is deliberately impoverished: a structure enters as bare 3D positions of heavy atoms, carrying no residue or atom-type labels, so the learned codebook encodes geometry alone. That is what lets one vocabulary span molecule classes. Alongside the domain-specific mol2token, protein2token and rna2token variants, the unified bio2token model trains across small molecules, proteins and RNA at once and outperforms the specialists on macromolecules. Unlike residue-level tokenizers such as ProTokens and AIDO.StructureTokenizer, its token sequence is one-to-one with atoms.

#Key Features

• One token per heavy atom: The autoencoder preserves sequence length end to end, so every input atom has a corresponding token and reconstructed position — side chains and nucleotide bases included, not inferred by a downstream module.
• A single codebook across three molecule classes: The unified model tokenizes proteins, RNA and small molecules with the same 4,096-entry vocabulary, learned once from a mixed corpus.
• Linear-time encoding: Bidirectional Mamba layers avoid the quadratic memory of attention, reaching structures of nearly 100,000 atoms on conventional GPUs where an invariant point attention decoder is intractable past a few hundred residues.
• Finite-scalar quantization: FSQ projects latents into an integer hypercube and rounds, sidestepping the codebook collapse that makes vector quantization hard to train.
• Complexes without complex training: Multi-chain and protein-RNA assemblies appear only at inference; neither kind was present in any training set.

#Technical Details

All four checkpoints share one architecture — four encoder and six decoder bidirectional Mamba layers, a 4,096-entry FSQ codebook, 1.2M parameters — and differ only in training corpus. Training combines CATH 4.2 (18k protein structures, 40–500 residues), a 100k-cluster subset of the AlphaFold Database drawn from FoldSeek's sequence-structure clusters, RNA3DB, and a DFT conformer dataset of 1.9M small molecules with 16M conformations. The loss equally weights total-structure RMSE after Umeyama-Kabsch alignment and an intra-residue inter-atomic distance term. On held-out proteins, bio2token reconstructs at 0.56 Å all-atom RMSE on the CATH 4.2 test split and 0.58 Å and 0.59 Å on CASP14 and CASP15, with TM-scores of 0.98–0.99; ESM-3's tokenizer reports 1.3 Å and 1.7 Å on the same CASP sets, and a published backbone tokenizer 1.89 Å against bio2token's 0.52–0.55 Å backbone error. RNA3DB test structures reconstruct at 0.66 Å (TM-score 0.96), including a 4,269-nucleotide chain of 90,441 atoms. On small molecules mol2token reaches 0.2 Å, though only 41.7% of reconstructions pass every chemical validity check.

#Applications

The intended use is as a frozen front end: a structure serialized into discrete tokens can be consumed by a language model, discrete diffusion, or any generative method that operates over vocabularies, without the generator itself having to handle all-atom geometry. Because the tokens are atom-level and molecule-class agnostic, one pipeline covers ligand conformers, protein folds, RNA and their assemblies. The repository ships the four pretrained checkpoints with scripts that tokenize and reconstruct a raw PDB file directly.

#Impact

Bio2Token is a preprint awaiting peer review, and it is a component rather than an application: it produces a representation, and the generative models meant to consume it are future work. Its practical argument is that atomic-resolution structure tokenization is cheap — 1.2M parameters, roughly 127,000 macromolecules of training data, with code, weights and preprocessed data released under an MIT license. Sub-Ångström RMSE does not guarantee chemical validity: reconstructions can imply covalent bonds that should not exist or miss ones that should. Accuracy degrades for atoms far from the coordinate origin, beyond the spatial extent seen in training. And because tokens encode position without atom identity, any downstream model must supply chemistry from another channel.