🤖 AI

NEAT: Neighborhood-Guided, Efficient, Autoregressive Set Transformer for 3D Molecular Generation

本文介绍了 NEAT，一种新颖的邻域引导自回归集合变换器，它通过一种与顺序无关的训练策略确保原子级别的置换不变性，从而在 QM9 和 GEOM-Drugs 数据集上实现了最先进的 3D 分子生成质量，并具备显著的速度优势。

原作者： Daniel Rose, Roxane Axel Jacob, Johannes Kirchmair, Thierry Langer

发布于 2026-05-07

📖 1 分钟阅读☕ 轻松阅读

原作者： Daniel Rose, Roxane Axel Jacob, Johannes Kirchmair, Thierry Langer

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

想象一下，你正在教一个机器人如何构建复杂的三维结构，比如由不同颜色的球（原子）和棍子（化学键）连接而成的分子模型。目标是让机器人从零开始创造新的、有效的结构，而不仅仅是复制现有的结构。

本文介绍了一种名为NEAT（邻域引导、高效、自回归集合 Transformer）的新型机器人架构师。以下是其工作原理的简明解释：

问题：“排队”带来的困惑

大多数以前的机器人试图通过严格的排队顺序来构建分子，就像学生一个接一个地进入教室。它们必须决定：“谁先？谁后？”

问题所在：在分子中，原子实际上并没有“第一”或“第二”的位置。它们作为一个三维簇同时存在。强迫它们排队，就像试图按特定顺序列出雪花晶体来描述雪花一样。这是武断且令人困惑的。如果机器人搞错了顺序，分子就会失效。
旧方法：以前的方法试图通过发明一种“规范”（标准）顺序来解决这个问题，但这使机器人产生了偏见。它学会了列表的顺序，却不如学会分子的形状。

解决方案：“邻域”策略

NEAT 改变了游戏规则。它不再将分子视为一条队列，而是将其视为一个人群。

工作原理：想象你正在一块砖一块砖地建造房屋。NEAT 不是问“列表中下一块砖是哪块？”，而是查看你已建成的当前墙体，并问：“站在墙边紧邻的邻居是谁？”
神奇之处：NEAT 观察当前结构的“边界”。它知道任何与当前结构接触的原子都是下一个可添加的有效候选者。它不在乎顺序，只在乎邻域。
结果：这使得机器人具有“置换不变性”。无论你从左边、右边还是倒过来观察分子，机器人看到的邻居都是一样的，并构建出相同的合法结构。

构建过程：两步同时进行

当 NEAT 决定添加一个新原子时，它会同时做两件事：

它是什么？ 它猜测原子的类型（例如，“这应该是一个碳球”）。
它在哪里？ 它使用一种称为**流匹配（Flow Matching）**的特殊数学技术，来确定该原子在三维空间中确切应该悬浮的位置。这就像一股河流，将一片叶子（原子）从随机起点轻轻推送到分子中完美且稳定的位置。

为什么 NEAT 很特别？

论文声称 NEAT 是一个重大升级，原因有三：

速度快：因为它不需要像其他机器人那样模拟成千上万个噪声步骤（这就像试图通过几个小时敲凿石块来雕刻雕像），NEAT 构建分子的速度要快得多。论文称其速度比竞争对手快5 到 16 倍。
灵活性强：因为它不依赖固定的排队顺序，你可以给它一个“启动套件”（部分分子），让它完成工作。无论你怎么递给它启动套件，它都能完成结构。如果启动套件不符合其他机器人特定的“排队”规则，它们往往会失败。
准确度高：在标准测试（称为 QM9 和 GEOM-Drugs 的数据集）中，NEAT 构建的分子在化学上是有效且独特的，其表现与现有最佳方法相当或更优。

局限（注意事项）

作者诚实地指出了 NEAT 目前尚不能做到的事情：

尺寸限制：它只能构建与其训练期间所见分子一样大的分子。如果它只学习了小分子，就无法外推构建巨大的分子。
化学键的模糊性：它依赖于知道哪些原子相连（化学键）来确定邻域。如果你给它一个化学键不清晰的奇怪晶体结构，它可能会感到困惑。
没有“魔法”三维规则：它没有内置的三维旋转物理规则（SE(3)-等变性），但作者发现，只需随机旋转训练数据就足以让它表现良好。

总结

简而言之，NEAT 是一种教导计算机设计三维分子的新方法。它不再强迫计算机遵循僵化的人工列表，而是教导它们观察直接邻居并自然地构建。它更快、更灵活，并且与当前最先进的方法一样准确。

技术摘要：NEAT——面向 3D 分子生成的邻域引导、高效、自回归集合 Transformer

问题陈述

当前发现生物活性分子的计算工作流严重依赖虚拟筛选或从头（de novo）生成。虽然虚拟筛选受限于化学库的有限规模以及不断增长的计算成本，但生成式建模提供了一种有前景的替代方案。然而，生成 3D 分子结构面临着独特的挑战。分子本质上是三维的，其性质往往由构象决定。

现有的最先进 3D 生成模型主要利用扩散和流匹配技术。这些方法在消息传递更新中往往存在二次缩放问题，缺乏对可变分子大小的内在适应性（通常需要在初始化时确定固定的原子数量），并且在去噪过程中难以管理动态的原子插入。

自回归模型提供了一种灵活的替代方案，天然支持可变原子数量，且每个 token 仅需一次编码器调用。然而，标准的自回归 Transformer 需要 token 的序列顺序。在分子图中，原子没有内在的序列顺序。先前的工作通过使用规范原子排序来解决这一问题，但这些方法对于原子而言不具备置换不变性。它们使下一个 token 的预测偏向于特定的排序惯例，限制了模型的泛化能力，或难以执行前缀补全等任务（在这些任务中，起始原子顺序可能与最终分子的规范顺序不一致）。相反，确保置换不变性的图神经网络（GNN）方法往往生成的分子化学有效性较低。

方法论

作者提出了 NEAT（Neighborhood-Guided, Efficient, Autoregressive Set Transformer，邻域引导、高效、自回归集合 Transformer），这是首个尊重原子置换不变性且不依赖规范原子排序的用于 3D 分子生成的自回归 Transformer 模型。

核心架构

NEAT 将分子图视为原子的集合而非序列。该架构由三个主要组件组成：

数据增强模块（邻域引导）： 模型不使用序列位置编码，而是采用一种采样策略来模拟分子生长。给定一个分子图，采样一个连通子图（源集合）。目标集合定义为源集合的“边界节点”（1 跳邻居）。这提供了一组有效的下一步 token 延续，无论原子是以何种顺序添加的。
集合 Transformer 主干： 源集合使用集合 Transformer 进行编码（基于 NanoGPT/GPT-2，采用双向注意力机制且无位置嵌入）。原子类型通过线性层嵌入，位置使用傅里叶特征进行编码。加性池化将这些节点表示聚合为集合级表示（ $z_n$ ）。该设计确保了源集合中原子顺序的置换不变性。
流匹配头： 集合级表示 $z_n$ $z_{n}$ 用于预测下一个 token，分解为原子类型和位置： $p(a_{n+1}, x_{n+1} | z_n) = p(a_{n+1} | z_n) \cdot p(x_{n+1} | a_{n+1}, z_n)$ $p (a_{n + 1}, x_{n + 1} ∣ z_{n}) = p (a_{n + 1} ∣ z_{n}) \cdot p (x_{n + 1} ∣ a_{n + 1}, z_{n})$ 。
- 原子类型： 通过线性层和 Softmax 建模为分类分布。
- 原子位置： 通过条件流匹配（CFM）建模。模型学习一个向量场，将随机初始位置分布传输到真实的原子位置。

训练策略

邻域引导： 在训练期间，模型不是由单个下一个 token 监督，而是由所有有效边界节点的集合进行监督。这消除了对规范化的需求。
教师强制（Teacher Forcing）： 原子类型预测和位置恢复被解耦。模型预测类型，流匹配头学习基于该类型的位置。
增强： 为了提高对误差累积的鲁棒性，在训练期间向部分源集合原子注入高斯噪声，并相对于含噪源集合重新居中位置。
损失函数： 总损失是原子类型的交叉熵损失和位置的流匹配损失之和。通过求解线性分配问题来耦合源分布和目标分布，最小化传输成本并避免路径交叉。

推理

生成从一个随机初始原子开始。在每一步，当前原子集合被编码，预测下一个原子类型，并采样一个随机初始位置。然后积分流匹配向量场（使用欧拉法或欧拉 - 丸山法）以确定最终的 3D 坐标。此过程重复进行，直到预测到停止 token。

主要贡献

置换不变的自回归生成： NEAT 引入了一种新颖的训练策略，利用邻域引导和集合 Transformer 实现 3D 分子生成，该生成对原子排序具有不变性，克服了自回归模型中规范排序的局限性。
效率： 通过将输入视为集合并使用每个 token 一次编码器调用，NEAT 避免了基于扩散的方法中消息传递的二次缩放，以及某些流匹配基线中每个原子生成步骤所需的多次积分步骤。
前缀补全能力： NEAT 的与顺序无关的特性使其能够进行高质量的前缀补全（围绕给定骨架生成分子），而无需前缀与特定的规范排序对齐，这是 QUETZAL 等基于序列的模型难以完成的任务。
最先进性能： 在标准基准测试（QM9 和 GEOM-Drugs）中，该模型取得了与自回归模型以及扩散/流匹配基线相比具有竞争力或更优越的结果。

实验结果

该模型在 QM9（小有机分子）和 GEOM-Drugs（类药分子）数据集上进行了评估。

生成质量： 在 QM9 上，NEAT 在 xyz2mol 有效性和唯一性方面优于先前的自回归模型（G-SchNet, G-SphereNet, Symphony, Geo2Seq）。它实现了 95.2% 的 xyz2mol 有效性和 91.6% 的唯一性，在查找有效性和唯一性方面排名第二。在 GEOM 上，NEAT 在自回归模型中实现了最高的 查找（lookup） 有效性（99.7%），以及第二好的 xyz2mol 有效性（93.9%）和唯一性（91.6%），其性能与 FlowMol3 等最先进的流匹配模型相当。
前缀补全： 在一项从 100 个多样化前缀生成 1,000 个分子的任务中，NEAT 显著优于 QUETZAL。NEAT 实现了 97.8% 的查找有效性和 92.7% 的 xyz2mol 有效性，而 QUETZAL 分别仅为 27.8% 和 28.1%。这证明了 NEAT 能够泛化独立于原子排序惯例的结构特征。
运行时效率： NEAT 表现出显著的速度优势。其运行速度比 SemlaFlow 快约 5 倍，比 FlowMol3 快 16 倍。在最坏情况下，NEAT 的复杂度为 $O(BN^3)$ （其中 $N$ 为原子数量），而扩散/流模型的缩放比例为 $O(BTN^2)$ （其中 $T$ 为积分步数）。由于对于类药分子而言 $N$ 通常小于 $T$ ，因此 NEAT 提供了实际的速度优势。

意义与主张

该论文声称，NEAT 为利用 Transformer 架构处理复杂药物设计任务（特别是条件 3D 分子生成）奠定了坚实基础。通过消除对规范原子排序的依赖，该模型实现了灵活的任务，如基于骨架的生成（前缀补全），而这些任务对于自回归方法此前而言是困难的。

作者指出，虽然当前架构并非显式的 SE(3) 等变，但在训练期间使用随机旋转增强有效地弥补了这一点，导致预测仅存在微小方差。该工作表明，将 NEAT 与 SE(3) 等变 Transformer 集成是未来研究的一个有前景的方向。

作者强调，虽然该模型生成了高质量的 3D 结构，但分子生成只是复杂流程中的一步，该流程需要广泛的验证、合成和测试。他们还承认了局限性，例如邻域引导对边信息的依赖（这可能限制其在无机晶体中的应用），以及无法外推到训练期间见过的最大原子数量之外。