[Paper] 通过规范化重新思考具有对称性的扩散模型及其在分子图生成中的应用

Source: arXiv - 2602.15022v1

概述

本文重新审视了扩散式生成模型如何处理对称性——比如原子的置换或分子的旋转。作者提出 canonicalization（标准化）方法：首先将每个数据点放入标准的“姿态”，在这些标准形式上训练普通（非等变）的扩散模型，然后在生成时重新应用随机对称性。这样可以以更简单、更快速且更具表达力的方式生成三维分子图。

关键贡献

• Canonical diffusion framework: 正式理论表明，在数据流形的规范切片上进行训练是正确的（恢复原始不变分布）且更具表达力，相较于直接强制等变性。
• Training efficiency gains: 证明规范化去除了扩散分数中的混合对称项，降低了方差并加速了扩散模型和流匹配模型的收敛。
• Unified view with aligned priors & optimal transport: 展示了这些互补技术在与规范化结合时如何进一步加速学习。
• Practical instantiation for molecules: 实现了一种基于几何谱的规范化器，用于分子图的置换 × SE(3) 组合对称性。
• State‑of‑the‑art results: CanonFlow 模型在 GEOM‑DRUG 基准上超越了现有的等变基线，即使使用更少的扩散步数且计算量相当。

方法论

1. 确定对称群 – 对于分子而言，相关的群是 (S_n)（原子置换）× (SE(3))（三维旋转与平移）。
2. 规范化步骤 – 将每个分子转换为 规范姿态：
   • 计算一个旋转，使分子的几何谱（基于距离矩阵的特征值）对齐。
   • 按确定性的方式对原子排序（例如，按特征向量分量排序），以固定置换。
     这为每个对称轨道提供唯一的代表。
3. 训练无约束生成模型 – 在规范化数据上训练标准的扩散（或流匹配）网络，且不施加任何等变约束。
4. 采样 – 生成规范样本后，抽取一个随机对称变换（随机旋转 + 随机置换）并应用，得到遵循原始不变分布的分子。
5. 增强措施 – 作者加入 对齐先验（使潜在先验匹配规范分布）和 基于最优传输的流匹配，进一步降低训练方差。

结果与发现

• 表达能力：规范模型可以表示任何等变模型能够表示的不变分布，并且由于不受架构对称性约束，往往能够捕获更细致的细节。
• 速度：在扩散得分中去除群混合项可降低梯度方差，使得在性能相当的情况下训练 epoch 减少约 20 %。
• 少步生成：即使仅使用 10 步扩散（相较于常见的 100 步以上），CanonFlow 仍保持高有效性和多样性，适用于实时应用。

实际影响

• 更简化的模型流水线: 开发者可以直接复用现成的扩散库（例如 PyTorch‑Diffusers），无需编写自定义等变层，从而降低工程开销。
• 更快的原型开发: 训练方差降低意味着可以更快速地进行超参数搜索并降低 GPU 成本——这对计算资源有限的初创公司或实验室至关重要。
• 更好地与下游工具集成: 由于生成的分子已经是规范形式，下游任务（例如对接、属性预测）可以更高效地缓存或批处理这些分子。
• 少步采样实现实时设计: 需要快速候选分子生成的药物发现流水线（例如主动学习循环）现在可以在不牺牲质量的前提下进行即时采样。
• 可扩展到其他领域: 任何具有已知对称群的生成任务（点云、蛋白质结构、物理模拟等）都可以采用相同的先规范化方法，可能取代笨重的等变网络。

限制与未来工作

• 规范化器设计：当前的谱对齐对小到中等分子效果良好，但在非常大或高度柔性的结构上可能会遇到特征谱模糊的困难。
• 群覆盖：该框架假设已知且可处理对称群；将其扩展到 SE(3) 之外的连续群（例如缩放、剪切）需要新的规范化技巧。
• 采样偏差：理论上，在生成时随机重新应用对称变换是无偏的，但在实践中有限样本效应可能导致细微的分布偏移——这是需要更严格统计分析的领域。
• 更广泛的基准：虽然 GEOM‑DRUG 是一个强有力的测试平台，但在其他化学数据集（如 QM9、MOSES）以及非化学对称数据上进行评估，将进一步巩固其通用适用性的主张。

底线：通过颠覆传统观念——先规范化，再生成——作者提供了一种更易用、高效且强大的对称感知生成建模方案，对分子 AI 及更广领域立即产生益处。

作者

• Cai Zhou
• Zijie Chen
• Zian Li
• Jike Wang
• Kaiyi Jiang
• Pan Li
• Rose Yu
• Muhan Zhang
• Stephen Bates
• Tommi Jaakkola

论文信息

• arXiv ID: 2602.15022v1
• 分类: cs.LG, cs.AI, math.GR, q-bio.BM
• 出版日期: February 16, 2026
• PDF: 下载 PDF