Efficient Molecular Conformer Generation with SO(3)-Averaged Flow Matching and Reflow

📄 arXiv: 2507.09785v1 📥 PDF

作者: Zhonglin Cao, Mario Geiger, Allan dos Santos Costa, Danny Reidenbach, Karsten Kreis, Tomas Geffner, Franco Pellegrini, Guoqing Zhou, Emine Kucukbenli

分类: cs.LG, physics.chem-ph

发布日期: 2025-07-13

备注: ICML 2025 poster

💡 一句话要点

提出SO(3)平均Flow Matching与Reflow，加速分子构象生成。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 分子构象生成 Flow Matching SO(3)平均 Reflow 蒸馏 计算化学 药物发现

📋 核心要点

  1. 现有分子构象生成方法计算成本高昂，特别是基于扩散或Flow的模型，训练和采样需要大量资源。
  2. 论文提出SO(3)-平均Flow训练目标加速训练，并利用Reflow和蒸馏方法加速推理，提高生成效率。
  3. 实验表明，SO(3)-平均Flow训练的模型达到最先进的生成质量，Reflow方法实现高质量的少步甚至单步生成。

📝 摘要（中文）

快速且准确的分子构象生成对于下游计算化学和药物发现任务至关重要。目前，训练和采样最先进的基于扩散或Flow的模型进行构象生成需要大量的计算资源。本文基于Flow Matching，提出了两种加速三维分子构象生成模型的训练和推理的机制。为了快速训练，我们引入了SO(3)-平均Flow训练目标，与条件最优传输Flow或Kabsch对齐Flow相比，该目标能够更快地收敛到更好的生成质量。我们证明了使用SO(3)-平均Flow训练的模型可以达到最先进的构象生成质量。为了快速推理，我们表明基于Flow模型的Reflow和蒸馏方法能够以高质量进行少步甚至单步分子构象生成。这项工作提出的训练技术展示了使用基于Flow的模型高效生成分子构象的途径。

🔬 方法详解

问题定义：分子构象生成旨在预测分子在三维空间中的稳定结构。现有基于扩散模型或Flow模型的构象生成方法，训练和推理过程计算量大，耗时较长，限制了其在实际应用中的效率。因此，如何降低计算成本，加速分子构象生成过程是本文要解决的核心问题。

核心思路：本文的核心思路是利用SO(3)群的平均操作来改进Flow Matching的训练过程，并结合Reflow和蒸馏技术加速推理。SO(3)平均能够减少模型需要学习的旋转不变性，从而加速收敛并提高生成质量。Reflow和蒸馏则通过减少推理步骤，显著降低计算成本。

技术框架：整体框架包含两个主要部分：快速训练和快速推理。快速训练部分使用SO(3)-平均Flow训练目标，替代传统的条件最优传输Flow或Kabsch对齐Flow。快速推理部分则利用Reflow和蒸馏技术，将原本需要多步迭代的Flow模型转化为少步甚至单步的生成过程。

关键创新：最重要的技术创新点在于SO(3)-平均Flow训练目标。与传统的Flow Matching方法相比，该方法通过在SO(3)群上进行平均，使得模型能够更好地学习旋转不变性，从而加速训练过程并提高生成质量。此外，将Reflow和蒸馏技术应用于分子构象生成，实现了高效的单步或少步推理，也是一个重要的创新。

关键设计：SO(3)-平均Flow训练目标通过对旋转后的分子坐标进行平均来构建Flow场。损失函数基于Flow Matching的损失函数，但加入了SO(3)平均操作。Reflow通过学习一个从噪声到数据的直接映射来减少推理步骤。蒸馏则通过训练一个更小的模型来逼近原始模型的输出，从而降低计算复杂度。具体的网络结构和参数设置在论文中有详细描述，但此处未给出。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，使用SO(3)-平均Flow训练的模型在构象生成质量上达到了最先进水平，并且训练速度显著提升。Reflow方法能够实现高质量的单步或少步构象生成，推理速度提升明显。这些结果验证了该方法在加速分子构象生成方面的有效性。

🎯 应用场景

该研究成果可广泛应用于药物发现、材料科学等领域。快速准确的分子构象生成能够加速虚拟筛选、分子动力学模拟等过程，从而缩短药物研发周期，降低研发成本。此外，该方法还可用于预测蛋白质结构、设计新型材料等。

📄 摘要（原文）

Fast and accurate generation of molecular conformers is desired for downstream computational chemistry and drug discovery tasks. Currently, training and sampling state-of-the-art diffusion or flow-based models for conformer generation require significant computational resources. In this work, we build upon flow-matching and propose two mechanisms for accelerating training and inference of generative models for 3D molecular conformer generation. For fast training, we introduce the SO(3)-Averaged Flow training objective, which leads to faster convergence to better generation quality compared to conditional optimal transport flow or Kabsch-aligned flow. We demonstrate that models trained using SO(3)-Averaged Flow can reach state-of-the-art conformer generation quality. For fast inference, we show that the reflow and distillation methods of flow-based models enable few-steps or even one-step molecular conformer generation with high quality. The training techniques proposed in this work show a path towards highly efficient molecular conformer generation with flow-based models.