🤖 machine learning

SymDrift: One-Shot Generative Modeling under Symmetries

SymDrift 是一种新颖的单样本生成建模框架，它通过最优对齐与不变嵌入引入一种对称性感知漂移场，以解决物理系统中的对称性挑战，在分子生成任务中实现了最先进的性能，同时相较于多步基线方法将计算开销降低了高达 40 倍。

原作者： Samir Darouich, Vinh Tong, Lluís Pastor-Pérez, Tanja Bien, Loay Mualem, Mathias Niepert

发布于 2026-05-08

📖 1 分钟阅读☕ 轻松阅读

原作者： Samir Darouich, Vinh Tong, Lluís Pastor-Pérez, Tanja Bien, Loay Mualem, Mathias Niepert

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

想象一下，你正在教一个机器人构建完美的分子三维模型，就像构建由原子组成的微小、精巧的雕塑。问题在于，这些分子就像旋转的陀螺或拼图块：如果你旋转它们或交换两个相同的原子，它们仍然是完全相同的分子。在物理学中，这被称为对称性。

长期以来，试图学习这些形状的 AI 模型不得不采取非常缓慢、循序渐进的方法。它们会猜测一个形状，检查它，调整它，再次检查，并重复数百次才能将其修正。这就像试图通过拍一张照片、检查它、将相机移动一点点、再拍一张，并重复此过程直到获得完美角度来寻找照片的最佳拍摄角度。这虽然有效，但极其缓慢且昂贵。

最近，一种名为"漂移模型"（Drifting Model）的新型 AI 被发明出来。将其想象为一台“单次快门”相机。它不再拍摄数百张照片，而是试图在单次点击中捕捉完美的画面。它学习一种“漂移场”——一种无形的风，能将随机的原子云直接推入正确的形状，一步到位。这要快得多，但在处理对称分子时存在一个重大缺陷。

问题：困惑的指南针

关键在于：标准的“漂移模型”会被对称性搞糊涂。

想象一下，你正在教一个机器人识别人脸。如果你给它看一张人的照片，然后给它看同一个人旋转 90 度后的照片，一个聪明的机器人会知道这是同一张脸。但标准的漂移模型会将旋转后的版本视为完全不同的一个人。

由于训练数据（分子照片）通常只显示一种随机朝向，模型学习到的“漂移”是指向该特定朝向的。它没有意识到，无论分子如何转动，“风”的方向应该是一样的。这就像指南针只有当你站在北极点时才指向北方，但只要你移动哪怕一英寸，它就会疯狂旋转。这导致模型生成的结构混乱且不正确，因为它在与分子的自然对称性作斗争。

为了解决这个问题，研究人员过去不得不向机器人展示同一分子数千次，以各种可能的方向旋转，只是为了让它学会这种模式。这就像为了教机器人人脸长什么样，而从宇宙中的每一个角度拍摄同一个人的照片。这太慢且太昂贵了。

解决方案：SymDrift

本文的作者 SymDrift 想出了一个巧妙的方法，无需拍摄数千张照片就能修正指南针。他们引入了两种新策略，使“漂移”本身具备对称性感知能力：

“最佳拟合”对齐（显式对齐）：
想象你有一块拼图和一个孔。与其强行将拼图塞进去，不如先旋转和翻转拼图，直到它完美契合，然后再将其推入。SymDrift 在数学上就是这样做的。在计算推动原子的“风”（漂移）之前，它会瞬间找到与当前猜测最匹配的目标分子的最佳旋转和排列。它在移动之前将它们完美对齐。这确保了模型学习的是正确的形状，而不仅仅是特定的朝向。
“抗对称”地图（不变嵌入）：
想象你试图向一个看不见的朋友描述一个形状。与其给出像“向上移动 3 英寸，向左移动 2 英寸”这样的坐标，不如描述各部分之间的距离：“鼻子到耳朵的距离是 2 英寸，耳朵之间的距离是 4 英寸。”无论你怎么转动头部，这些距离都不会改变。
SymDrift 利用了这一思想。它将三维分子转换为原子之间距离的列表。由于无论怎样旋转或打乱分子，这些距离都是相同的，模型便在这张“抗对称”地图上进行学习。它不会感到困惑，因为地图本身不关心旋转或交换相同的部分。

结果：快速且准确

论文表明，SymDrift 对于两项特定任务是一个游戏规则改变者：

构象生成：创建分子可以采取的不同三维形状（就像一个人伸展手臂或弯曲膝盖）。
过渡态生成：预测分子在化学反应正中间瞬间呈现的不稳定形状（就像球在滚下山坡前处于山顶最高点的快照）。

为什么这很重要：

速度：SymDrift 比旧的缓慢方法快达40 倍。它只需单步即可生成分子，而不是数百步。
准确性：它产生的结果与缓慢方法一样好，有时甚至更好。
效率：它节省了巨大的计算能力。

作者在标准的分子和化学反应数据集上测试了这一点。他们发现，SymDrift 几乎可以瞬间生成高质量的分子结构，使其成为筛选数百万种潜在药物或探索复杂化学反应等需要速度的任务的强大工具。

简而言之，SymDrift 教会 AI 理解分子形状的本质，忽略旋转和交换带来的干扰，使其能够在单次闪电般的快门中捕捉完美的三维画面。

技术摘要：SymDrift

问题陈述

物理系统（如分子）的生成建模需要学习在整体对称性下保持不变的分布（例如 $\mathbb{R}^3$ 中的旋转和相同原子的置换）。尽管等变扩散和流匹配模型能够有效融入这些不变性，但它们通常依赖耗时的多步采样过程，使其在虚拟药物筛选等高吞吐量应用中计算成本过高。

近期，“漂移”（drifting）模型作为一种高效的一次性（one-shot）替代方案应运而生，它将生成问题表述为分布匹配问题，学习一个其稳态点与目标分布相匹配的向量场。然而，作者指出，将漂移模型应用于对称域时存在一个关键局限：一个等变生成器通常不会产生与从对称化目标分布获得的漂移场相同的漂移场。

具体而言，当在有限的经验分布 $p$ （该分布对对称群 $G$ 不具备不变性）上进行训练时，学习到的漂移场会偏向于每个等价类的任意代表元。理论分析（定理 1）表明，由等变生成器导出的聚合向量场 $\hat{V}_p(x)$ 与由完全对称化分布诱导的向量场 $V_{p_G}(x)$ 在数学上是截然不同的。这种不匹配导致模型无法恢复与对称化目标相关的正确漂移，除非执行昂贵的数据增强（对称化），否则会导致训练崩溃或生成效果次优。

方法：SymDrift

为了在不产生显式数据增强计算成本的情况下解决这一特定于对称性的挑战，作者提出了SymDrift，这是一个使漂移场本身具备对称感知能力的框架。该方法引入了两种互补策略来构建近似对称化目标分布的漂移：

1. 通过显式对齐实现的对称化漂移

该策略在坐标空间中运行。对于每个生成的样本 $x$ 和目标样本 $y$ ，该方法显式地寻找最优群变换 $g^*(x, y)$ ，以最小化距离 $\|x - gy\|$ 。

机制：通过将标准样本对 $(x, y)$ 替换为最优对齐的样本对 $(x, g^*(x, y)y)$ 来构建漂移场。
实现：旋转对齐利用 Kabsch 算法高效计算。然而，确定相同原子的最优置换仍然具有计算挑战性（阶乘复杂度）。作者指出，虽然迭代方法（在匈牙利算法和 Kabsch 对齐之间交替）提供了一种实用的近似，但它们可能陷入局部极小值并产生高昂的开销。

2. 在 $G$ 不变嵌入空间中进行漂移

为了绕过显式对齐的计算瓶颈，作者提出在固定的 $G$ 不变嵌入空间中计算漂移。

嵌入构建：他们定义了一个基于类型条件成对距离的非参数嵌入 $\phi(x)$ 。对于分子构型，成对距离根据原子对的相互作用类型（例如 C-C、C-H）进行分组，并在每个组内进行排序。
性质：该嵌入被构建为在全局旋转（$SO(3) $）和相同原子置换（$ \hat{S}_N$）下保持不变。通过将输入映射到该空间，核相似度度量内在尊重物理系统的对称性，从而在漂移计算过程中无需显式对齐即可消除歧义。
灵活性：该框架并不局限于这种特定的嵌入；任何可微的不变嵌入均可被利用。

训练目标最小化生成器预测与由对称感知漂移场定义的不动点之间的距离，使用停止梯度（stop-gradient）操作将目标视为冻结的参考。

主要贡献

SymDrift 框架：一个用于高度对称分布一次性生成建模的新颖框架，显式地将对称性融入漂移场的设计中。
理论洞察：证明了仅靠等变生成器不足以让漂移模型恢复对称化目标分布，从而需要在漂移计算中进行显式的对称性处理。
双重策略：开发了两种实现对称感知能力的互补方法：(i) 显式坐标空间对齐和 (ii) 不变嵌入空间。
效率：通过实现一次性推理同时保持高精度，显著降低了计算开销。

实验结果

作者在两个标准基准上评估了 SymDrift：分子构象生成和过渡态（TS）生成。

分子构象生成（GEOM-QM9）：
- SymDrift 在准确性指标（AMR）和覆盖率（COV）方面，在一次性推理基线中取得了最先进的性能。
- 它优于现有的单次方法，如 GeoMol 和 AvgFlow。
- 虽然多样性略低于使用额外能量模型的 EnFlow，但 SymDrift 的速度比 EnFlow 快 3.5 倍。
- 与多步基线（例如使用 50 次网络评估的 ET-Flow）相比，SymDrift 在保持相当准确性的同时，推理成本降低了 40 倍。
过渡态生成（RDB7）：
- 使用更大的模型变体（1870 万参数），SymDrift 实现了与 25 步 GoFlow 基线相当的 RMSD 和 DMAE 指标。
- 一次性推理比多步 GoFlow 基线快 87 倍。
- 这证明了该方法适用于速度至关重要的反应网络大规模探索。
消融研究：
- 实验证实，若对等变架构应用基础坐标漂移而不进行对称性处理，会导致完全的训练崩溃（0% 覆盖率）。
- 随着对称性处理的引入，性能逐步提升，其中不变嵌入空间产生了最佳结果，优于容易陷入局部最优的迭代对齐策略。

意义与主张

本文声称，SymDrift 成功弥合了一次性生成模型的效率与对称域所需的物理正确性之间的差距。通过解决标准漂移方法中固有的漂移场不匹配问题，SymDrift 实现了：

高吞吐量应用：推理时间减少高达 40 倍，使该方法在虚拟药物筛选和大规模反应网络探索等应用中极具前景，而这些应用目前受限于迭代采样。
训练效率与推理速度的权衡：虽然由于需要成对对齐或嵌入计算，该方法增加了训练成本，但鉴于推理速度提升了数量级，这被视为一种有利的权衡。
通用性：该框架适用于任何涉及 $G$ 不变分布的生成任务，为数据增强或复杂的多步蒸馏提供了一种原则性的替代方案。

作者对局限性保持谦逊，指出所提出的不变嵌入并非最大表达力（丢弃了关联信息），且与标准流匹配相比训练时间有所增加。然而，他们论证道，对于高吞吐量的科学应用，推理速度的提升足以证明这些成本的合理性。