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

2024-10-24

arXiv.org

AI Summary

To address the challenge of high-fidelity discrete encoding for large-scale, all-atom 3D biomolecular structures (proteins, RNA, small molecules), this paper proposes the first end-to-end tokenization framework integrating a quantized variational autoencoder (Q-VAE) with the Mamba state-space model. The method introduces 3D geometry-aware atomic embeddings to directly learn atomic coordinates—bypassing coarse-grained approximations—and achieves sub-angstrom structural compression via Q-VAE while leveraging Mamba’s efficiency in modeling long-range atomic dependencies. Experiments demonstrate reconstruction errors below 1 Å, substantially outperforming SE(3)-invariant IPA-based architectures; the framework scales to systems with up to 100,000 atoms. This work establishes a scalable, high-fidelity structural token representation for large-scale all-atom generative modeling.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Computer Vision: 3D Computer Vision
• Natural Language Processing: (Large) Language Models

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Scalable techniques for the creation, curation, publication, maintenance, and consumption of large, Web-based, structured, reusable, knowledge graphs and ontologies
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications

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.

Problem

• Efficient encoding of large 3D molecular structures
• High-fidelity atom-level representation of biomolecules
• Scalable modeling for systems with 100,000 atoms

Innovation

• Quantized auto-encoders for atom-level tokenization
• Mamba state space model for efficient scaling
• High-fidelity reconstruction below 1 Angstrom accuracy