Bio2Token: All-atom tokenization of any biomolecular structure with Mamba
作者: Andrew Liu, Axel Elaldi, Nathan Russell, Olivia Viessmann
分类: cs.LG, cs.AI
发布日期: 2024-10-24 (更新: 2025-04-08)
💡 一句话要点
Bio2Token：利用Mamba实现生物分子结构的全原子Token化
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 生物分子结构 原子级Token化 Mamba模型 量化自编码器 分子设计 表征学习 三维分子
📋 核心要点
- 现有生物分子表征学习方法在处理大型复杂结构时，常采用粗粒度近似，牺牲了原子级别的精度。
- Bio2Token提出一种基于量化自编码器的原子级Token化方案，能够精确表示蛋白质、RNA和小分子结构。
- 实验表明，基于Mamba的Bio2Token模型在效率和精度上均优于传统方法，并能扩展到更大规模的生物分子系统。
📝 摘要（中文）
高效且高保真地编码和表示大型3D分子结构对于生物分子设计应用至关重要。然而，许多表征学习方法限制于建模较小的系统，或使用系统的粗粒度近似，例如以氨基酸残基的分辨率而非单个原子的水平对蛋白质进行建模。为了解决这个问题，我们开发了量化自编码器，学习完整蛋白质、RNA和小分子结构的原子级Token化，重建精度远低于1埃。我们证明，与SE(3)-不变的IPA架构相比，一个简单的Mamba状态空间模型架构是高效的，达到了具有竞争力的精度，并且可以扩展到具有近10万个原子的系统。bio2token学习到的结构Token可以作为未来全原子生成模型的输入。
🔬 方法详解
问题定义：现有生物分子表征方法，如基于氨基酸残基的粗粒度模型，无法捕捉原子级别的细节信息，限制了其在精确分子设计中的应用。此外，处理大型生物分子结构时，计算复杂度高，难以扩展。
核心思路：论文的核心思路是利用量化自编码器学习生物分子结构的原子级Token化表示。通过将三维原子坐标映射到离散的Token空间，实现对复杂结构的有效编码和压缩，同时保留原子级别的精度。
技术框架：Bio2Token框架主要包含两个阶段：编码阶段和解码阶段。在编码阶段，量化器将原子坐标映射到离散的Token。在解码阶段，解码器利用Token重建原始的原子坐标。论文使用Mamba状态空间模型作为编解码器的核心架构，Mamba模型擅长处理序列数据，能够有效地捕捉原子之间的空间关系。
关键创新：该方法的关键创新在于使用Mamba状态空间模型进行原子级Token化。Mamba模型相比于传统的Transformer模型，具有更高的计算效率和更好的长程依赖建模能力，使其能够处理更大规模的生物分子结构。此外，该方法通过量化自编码器学习原子坐标的离散表示，实现了对生物分子结构的有效压缩。
关键设计：论文采用量化变分自编码器（VQ-VAE）作为核心框架。量化器的作用是将连续的原子坐标映射到离散的Token空间。损失函数包括重建损失和量化损失。重建损失用于衡量重建原子坐标与原始坐标之间的差异，量化损失用于优化量化器的性能。Mamba模型的具体参数设置（如层数、隐藏层维度等）需要根据具体的生物分子系统进行调整。
🖼️ 关键图片
📊 实验亮点
实验结果表明，Bio2Token能够以远低于1埃的重建精度对蛋白质、RNA和小分子结构进行原子级Token化。与SE(3)-不变的IPA架构相比，基于Mamba的Bio2Token模型在效率和精度上均具有竞争力，并且能够扩展到具有近10万个原子的系统。这表明Bio2Token在处理大型生物分子结构方面具有显著优势。
🎯 应用场景
Bio2Token具有广泛的应用前景，可用于生物分子设计、药物发现和蛋白质工程等领域。通过将生物分子结构转化为Token序列，可以利用自然语言处理领域的先进技术，如生成模型和预训练模型，进行分子性质预测、分子生成和分子优化。该方法有望加速新药研发和生物材料设计。
📄 摘要（原文）
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.