🤖 machine learning

Autoregressive latent diffusion for 3D molecule generation

本文介绍了 KRONOS，这是一个潜自回归扩散框架，它在单一架构内统一了无条件和片段条件的 3D 分子生成，在保持与扩散模型竞争力的同时，实现了自回归方法中的最先进性能。

原作者： Federico Ottomano, Gaopeng Ren, Yingzhen Li, Kim E. Jelfs, Alex M. Ganose

发布于 2026-07-13

📖 1 分钟阅读☕ 轻松阅读

原作者： Federico Ottomano, Gaopeng Ren, Yingzhen Li, Kim E. Jelfs, Alex M. Ganose

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下你是一位乐高建筑大师，正试图从零开始构建一个全新的、复杂的分子。长期以来，构建分子的最佳方式就像是使用扩散模型（diffusion model）：想象一下从一大团混乱的、漂浮着的乐高积木云开始，然后慢慢地、神奇地摇晃桌子，直到这些积木自动拼凑成一个完美的形状。这种方法效果极佳，但它有一个缺陷：你必须在开始之前准确地告诉机器你需要多少块积木。如果你猜错了，整个结构就会崩塌。

随后出现了自回归模型（autoregressive models），它们更像是通过一次放置一块积木来搭建一座塔。你放一块，再放下一块，接着下一块。这种方法非常出色，因为你不需要预先知道最终的大小，甚至可以从一个特定的“种子”零件（比如已知的药物骨架）开始，围绕它构建新的结构。然而，这些“逐块搭建”的构建者往往难以达到“云端摇晃”式扩散模型的高质量，而且在尝试围绕特定种子进行构建时，有时会感到困惑。

于是，KRONOS诞生了——这是一个由伦敦帝国理工学院研究人员推出的新框架。请不要把 KRONOS 看作是一个处理原始、杂乱积木的建造者，而要把他看作是一位使用压缩的、神奇蓝图的建筑师。

神奇蓝图（潜空间/Latent Space）

K-RONOS 并不是试图直接预测每一个原子的精确 3D 坐标和化学类型（这就像是在放置每一块积木时都要描述它的颜色和纹理），它首先将分子转化为一个“潜空间”。你可以把它想象成一种秘密的、压缩的语言，其中整个分子由一个单一的、平滑且连续的标记（token）来表示。这就像是将一个复杂的 3D 雕塑转化为一行优雅的代码，其中包含了所有的形状和化学信息。

KRONOS 通过逐个预测这些压缩后的标记来构建分子，就像阅读句子一样，一个词接一个词。但其中的玄机在于：为了预测下一个标记，它使用了扩散过程（diffusion process）。这就像建筑师看着当前的蓝图，想象下一个部件的一个略微“嘈杂”或模糊的版本，然后通过“去噪（denoising）”来确定它究竟应该是什样子。这结合了两者的优点：既有循序渐进构建的灵活性，又拥有扩散模型那种高质量的“摇晃”能力。

“中间填空”技巧（Fill-in-the-Middle）

KRONOS 最酷的特性之一是它如何处理“部分”分子。想象一下你有一个拼图，你想填补缺失的部分，但缺失的部分并不总是在末尾；有时它们在中间，或者你想围绕一个特定的块进行构建。

研究人员使用了一种受“中间填空”（Fill-in-the-Middle, FIM）启发的策略。在训练期间，他们不仅教模型从左向右阅读，还教它观察一个分子，将其切分成随机的块（片段），隐藏其中的一些块，并要求模型填补空白。有时模型必须从头开始构建整个结构（无条件生成），有时它必须完成一个特定的片段（片段条件生成）。

论文指出，通过这种方式，KRONOS 学会了在不需要两个不同大脑的情况下精通这两项任务。这就像训练一位厨师既能从零开始做一顿全餐，又能完成别人已经开始做的菜肴，而且在同一个厨房里，完全不会感到困惑。

结果：它有多好？

团队在两个主要数据集上测试了 KRONOS：QM9（包含 97,609 个小分子的集合）和 GEOM-Drugs（一个包含超过 110 万个类药性构象的大规模集合）。

在 QM9 上： KRONOS 实现了 97.3% 的 3D 有效性（3D Validity） 和 99.07% 的 3D 原子稳定性（3D AtomStable）。这意味着它生成的几乎每个分子在化学上都是正确且稳定的。在所有测试的“逐块搭建”（自回归）方法中，KRONOS 脱颖而出。它甚至能与像 JODO 这样强大的“云端摇晃”（扩散）模型相媲美，后者的有效性为 96.3%。
在 GEOM-Drugs 上： 这是更难的测试。KRONOS 实现了 94.7% 的 3D 有效性 和 86.45% 的原子稳定性，超越了其他自回归模型，并与扩散方法保持了竞争力。
物理现实检查： 研究人员不仅检查分子看起来是否正确，还通过一个名为 GFN2-xTB 的量子化学模拟来查看它们在物理上是否稳定。在 QM9 上，生成的分子与优化后的完美结构相比，中值 RMSD（均方根距离） 仅为 0.074 Å（埃）。这表明 KRONOS 构建的分子已经非常接近其自然的、稳定的形状。

KRONOS 不是什么

需要注意的是，这篇论文并没有声称 KRONOS 是一个完全“置换不变（permutation-invariant）”的模型，像某些其他先进架构（如 NEAT）那样。这意味着它观察片段的顺序会对它产生一定影响，而真正的置换不变模型则完全不在乎顺序。然而，论文认为，“中间填空”的训练足以让模型变得鲁棒，而无需这种额外的复杂性。

此外，KRONOS 的质量受限于其“神奇蓝图”（它所使用的预训练自编码器）的质量。如果蓝图不完美，KRONOS 就无法构建完美的分子。作者建议，如果未来有人发明了更好的蓝图，KRONOS 也会随之变得更好。

总结

KRONOS 表明，你并不需要在“逐块构建分子（灵活、长度可变）”与“使用强大的扩散过程（高质量）”之间做选择。通过在压缩的“潜空间”中工作并使用巧妙的“中间填空”训练技巧，它成功实现了两者的兼得。它创建的分子在化学上有效、几何上精确，并且可以围绕特定的起点进行构建，同时还能保持构建任意大小分子的可能性。这是一个极具前景的新型药物研发工具，它表明，有时构建未来的最佳方式是学会从“中间”入手。

技术摘要：KRONOS —— 用于 3D 分子生成的自回归潜空间扩散模型

问题陈述

三维（3D）分子生成是药物发现和分子设计中的一项关键任务。目前最先进的方法主要依赖于扩散模型（diffusion models）和流匹配（flow matching）。虽然这些方法实现了高质量的生成，但它们通常需要预先指定分子大小（原子数量），这限制了其在变长生成方面的灵活性。

相反，**自回归（AR）**模型天然支持变长生成，并能对部分分子结构进行条件化（例如，对已知支架进行扩展），这在药物化学中具有极高的价值。然而，现有的 AR 模型面临两个主要挑战：

性能差距： 从历史上看，AR 模型在生成质量方面一直落后于扩散模型，特别是在几何保真度和化学有效性方面。
目标平衡： 在单一架构内实现强力的无条件生成（从头设计）和片段条件化生成（分子补全）之间的平衡是非常困难的。许多 AR 模型在不进行架构修改或性能权衡的情况下，难以在这些范式之间实现泛化。

此外，现有的潜空间分子生成器几乎完全依赖于扩散模型，这使得在潜空间中进行自回归生成的潜力在很大程度上仍未被探索。

方法论：KRONOS

作者引入了 KRONOS，这是一个潜自回归扩散框架，它直接在预训练统一自编码器（Unified Autoencoder, UAE, Luo et al., 2025）的潜空间中生成 3D 分子。该方法结合了三个核心组件：

1. 潜自回归生成

KRONOS 不在原始数据空间中操作（即同时预测离散原子类型/键和连续坐标），而是在预训练 UAE 的连续潜空间中运行。

表示： 一个 3D 分子被编码为一组潜标记序列 $Z = \{z_1, \dots, z_n\}$ ，其中每个标记共同编码了原子身份、键合环境和 3D 几何结构。
因子分解： 模型学习一个从左到右的自回归分布：
$p_\theta(z_{1:n}) = \prod_{i=1}^n p_\theta(z_i | z_{<i})$
变长特性： 由于潜空间缺乏天然的序列结束标记，KRONOS 采用了一个专门的停止分类器（stop classifier）（一个轻量级 MLP），该分类器与生成器一起训练，用于动态预测序列终止，从而消除了对预定义原子计数的需求。

2. 基于扩散的标记预测

为了对连续潜标记的条件分布 $p_\theta(z_i | z_{<i})$ 进行建模，KRONOS 利用了**扩散损失（Diffusion Loss）**公式（Li et al., 2024）。

一个因果 Transformer 将前序标记编码为上下文表示 $c_i$ 。
扩散去噪网络 $\epsilon_\theta$ 预测添加到当前潜标记 $z_i$ 中的噪声，该预测以 $c_i$ 和扩散时间步为条件。
这种方法通过去噪目标取代了离散标记预测，从而能够灵活地对连续变量进行建模，而无需离散化。

3. 混合训练策略（片段条件化生成）

为了在单一模型中支持无条件生成和片段条件化生成，作者引入了一种受 填空中间（Fill-in-the-Middle, FIM） 范式启发的训练策略。

片段分解： 分子使用基于 BRICS 的方案分解为具有化学意义的片段。
序列构建： 在训练期间，每个分子会被转化为以下两种形式之一：
- 默认顺序序列（用于无条件生成）。
- 片段条件化序列，其中随机的一组片段作为“种子”（上下文），剩余的片段作为生成目标。
机制： 模型被训练为在以种子片段为条件的条件下，自回归地生成目标片段。这是通过为片段和其中的原子分配单一的确定性顺序来实现的，而不是像顺序无关（Order-Agnostic）AR 模型那样对所有排列进行边缘化。这使得模型能够在不改变架构的情况下学习分子补全能力。

核心贡献

KRONOS 框架： 一种新型的潜自回归扩散框架，在统一的潜空间中共同建模分子图拓扑结构和几何结构。
统一训练策略： 一种混合训练方法，通过受 FIM 范式启发的手段，使单一左向右自回归模型能够同时支持无条件生成和片段条件化生成。
变长生成： 通过停止分类器实现的统一框架，自然地支持变长生成，避免了对预先指定大小的需求。
性能： 证明了潜自回归生成可以达到自回归方法的领先性能，同时在与扩散模型竞争时保持竞争力，且在启用片段条件化时不会牺牲无条件生成的质量。

实验结果

作者在 QM9 和 GEOM-Drugs 数据集上对 KRONOS 进行了评估。

无条件生成

QM9： KRONOS 实现了最高的 3D 有效性（97.3%） 以及在自回归方法中最强的综合性能。在包括 3D 原子稳定性、键长和键角 MMD 在内的各项指标上，它与领先的扩散模型（如 JODO, EDM）保持了竞争力。
GEOM-Drugs： KRONS 在所有评估的基准模型（包括自回归和扩散模型）中取得了最强的综合性能。
物理合理性： 使用半经验几何优化（GFN2-xTB），KRONOS 生成的构象相对于优化后的结构表现出较低的均方根距离（RMSD）和能量差（ $\Delta E$ ），表明它们位于物理上合理的构象空间区域附近。

片段条件化生成

在一个包含约 1.5 万个分子-片段对的数据集上，KRONOS 在有效性和稳定性指标的 Mean@3 和 Best@3 指标上显著优于自回归基准模型 QUETZAL。
至关重要的是，启用片段条件化生成对无条件生成性能的影响微乎其微，这证明了可以在单一架构内支持这两种范式。

意义与主张

论文声称 KRONOS 为现有的数据空间自回归和基于扩散的方法提供了一个极具前景的替代方案。其意义在于：

弥合差距： 它显著缩小了自回归模型与扩散模型之间的性能差距，在达到自回归方法领先水平的同时，也能匹配扩散模型的基准。
灵活性： 它将自回归模型的优势（变长、部分条件化）与扩散模型的高质量生成相结合，且全部在单一的潜空间公式内实现。
实际应用价值： 通过支持无需架构修改的片段条件化生成，它为药物化学任务提供了实用的工具，因为在这些任务中，分子通常是通过扩展已知支架来设计的。

作者指出了局限性，包括对预训练潜空间表达能力的依赖，以及对固定 BRICS 分解方案的使用，这表明未来的工作可以探索如何将片段抽象与生成模型共同学习，或扩展该框架以进行属性引导的设计。