Bio2Token: All-atom tokenization of any biomolecular structure with Mamba

Andrew Liu, Axel Elaldi, Nathan Russell, Olivia Viessmann

Track: Machine learning: computational method and/or computational results

Nature Biotechnology: Yes

Keywords: all-atom modelization, long context, auto-encoder, tokenization

TL;DR: We develop quantized auto-encoders to tokenize small molecules, proteins and RNA on the all-atom level with mamba.

Abstract: Efficient encoding and representation of large 3D molecular structures with high fidelity is critical for biomolecular design applications. Despite this, many representation learning approaches restrict themselves to modeling smaller systems or use coarse-grained approximations of the systems, for example modeling proteins at the resolution of amino acid residues rather than at the level of individual atoms. To address this, we develop quantized auto-encoders that learn atom-level tokenizations of complete proteins, RNA and small molecule structures with reconstruction accuracies well below 1 Angstrom. We demonstrate that a simple Mamba state space model architecture is efficient compared to an SE(3)-invariant IPA architecture, reaches competitive accuracies and can scale to systems with almost 100,000 atoms. The learned structure tokens of bio2token may serve as the input for all-atom generative models in the future.

Anonymization: This submission has been anonymized for double-blind review via the removal of identifying information such as names, affiliations, and identifying URLs.

Presenter: ~Axel_Elaldi1

Format: Yes, the presenting author will attend in person if this work is accepted to the workshop.

Funding: No, the presenting author of this submission does *not* fall under ICLR’s funding aims, or has sufficient alternate funding.

Submission Number: 54