🤖 machine learning

FlashMol: High-Quality Molecule Generation in as Few as Four Steps

本文介绍了 FlashMol，这是一种超快生成模型，它通过采用分布匹配蒸馏并结合优化的时间步重排与 Jensen-Shannon 散度正则化，仅需四步即可生成高质量的三维分子构象，在保持或超越传统扩散模型质量的同时实现了高达 250 倍的加速。

原作者： Xinyuan Wei, Zian Li, Shaoheng Yan, Cai Zhou, Muhan Zhang

发布于 2026-05-11

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Xinyuan Wei, Zian Li, Shaoheng Yan, Cai Zhou, Muhan Zhang

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

想象一下，你正在尝试烘焙一块完美的蛋糕。在计算机科学的世界里，“烘焙分子”意味着利用人工智能设计一种新的化学结构，使其可能成为拯救生命的药物。

长期以来，做到这一点最好的方法就像遵循一个需要1000 个微小、谨慎步骤的食谱。你必须混合、检查、调整，然后再混合，如此反复数百次，只为得到一块好蛋糕。这种方法虽然准确，但极其缓慢且昂贵，使得烘焙成千上万块蛋糕以寻找最佳方案变得不可能。

最近，科学家们试图通过将步骤减少到约 12 步或 50 步来加速这一过程。但这里有个陷阱：烘焙出的蛋糕要么烤焦了，要么扁平了。人工智能过于匆忙，丢失了使分子稳定的“化学规则”。

FlashMol 应运而生。你可以将 FlashMol 想象成一位大师级烘焙师，他找到了只需4 步就能烘焙出完美蛋糕的方法。

以下是他们如何实现这一目标的分解，分为三个简单的技巧：

1. “智能调度”技巧（重新调整时间步长）

想象你正试图穿过一个房间。旧的方法是迈出 1000 个大小相等的微小步伐。新的“少步”尝试试图迈出 50 个巨大的跨步，但它们总是绊倒，因为最初的几步跨得太大，而最后几步又太小，无法轻柔落地。

FlashMol 团队意识到，步伐的大小比步伐的数量更重要。

问题：如果在最开始（当分子还只是一团模糊的噪声云）迈出巨大的步伐，你可能会错过重要的细节。如果在最后迈出微小的步伐，则会浪费时间。
解决方案：他们重新设计了调度方案。他们告诉人工智能：“在开始时迈出几个巨大而大胆的步伐以确立大致形状，然后在最后迈出许多微小而谨慎的步伐以完善细节。”
结果：即使在人工智能尚未学习任何新内容之前，仅仅改变步伐大小，就使得 4 步版本能够生成稳定、未破损的分子。这为人工智能提供了一个更好的“起点”。

2. “模仿者”技巧（分布匹配蒸馏）

通常，为了教会一个学生人工智能快速工作，你会让它观察慢速的专家教师人工智能完成所有 1000 个步骤的过程。但学生试图模仿教师走过的路径，如果学生在早期犯了一个错误，整条路径就会出错（就像“传话”游戏一样）。

FlashMol 使用了一种称为分布匹配的不同方法。

学生不再被要求模仿路径，而是被告知：“不要担心你是如何到达那里的。只要确保你烘焙的最终蛋糕看起来与教师的最终蛋糕完全一样即可。”
人工智能会观察教师完成的分子和学生完成的分子，并问：“它们有多不同？”然后它进行调整以使两者匹配。这防止了“传话游戏”式的错误层层累积。

3. “多样性”技巧（添加安全网）

“模仿者”技巧有一个副作用：学生人工智能变得过于保守。它一遍又一遍地烘焙完全相同类型的蛋糕，因为它害怕犯错。在药物发现中，你需要多样性——你需要尝试许多不同的形状，以找到起作用的那一种。

为了解决这个问题，团队添加了一个“多样性助推器”。

他们添加了一条特殊规则，规定：“只要它仍然是一块好蛋糕，尝试稍微不同的形状是可以的。”
这保持了分子的高质量，同时确保人工智能不会只是重复相同的设计 1000 次。

最终结果

通过结合这些技巧，FlashMol 能够在4 步内生成高质量、化学上有效的 3D 分子。

速度：它比旧的 1000 步方法快250 倍。
质量：它生成的分子与缓慢生成的分子一样稳定且有效。事实上，在某些测试中，它甚至比那些试图在 12 步或 50 步内完成的其他“快速”方法还要好。

简而言之：FlashMol 不仅仅是让人工智能跑得更快；它教会了人工智能如何迈出更聪明的步伐，以及如何从结果而非过程中学习，从而能够几乎瞬间设计出复杂的化学结构。

技术摘要：FlashMol

问题陈述

生成化学上有效的三维分子构象是计算药物发现中的关键瓶颈。虽然基于扩散的模型（如 GeoLDM 和 EDM）能够生成高质量、稳定的分子，但它们通常每个样本需要数百到数千次顺序网络评估（NFEs）（例如 1000 次 NFEs）。这种计算成本使得大规模in silico（计算机模拟）筛选变得不切实际。最近的少步方法已将 NFEs 减少到 12–50 次范围，但往往伴随着显著的权衡：样本稳定性和有效性的急剧下降。迄今为止，没有任何现有方法能够在 10 步或更少步数内实现可靠的高质量分子生成，而这一步数范围正逐渐成为图像和视频生成领域的标准。

方法论

作者提出了FlashMol，这是一种超快生成模型，能够在少至 4 步内生成高质量的三维分子构象。FlashMol 利用分布匹配蒸馏（Distribution Matching Distillation, DMD），从 1000 步的教师模型（GeoLDM）中蒸馏出一个少步学生生成器，该方法针对分子领域进行了专门调整。该方法解决了将 DMD 应用于分子生成时固有的两个主要挑战：

1. 通过重排时间步逃离不良初始模式

标准 DMD 依赖于学生模型生成合理的初始样本，以提供有用的优化信号。然而，分子结构对噪声高度敏感；键长和键角的微小偏差可能导致分子在化学上无效。直接应用具有均匀时间步长的 DMD 会导致学生模型产生近 0% 的稳定分子，从而引发梯度爆炸和训练失效。

为解决这一问题，FlashMol 引入了受 EDM 设计空间分析启发的时间步重排策略。该方法不再使用均匀时间步长，而是在高噪声水平下前置较大的去噪步骤，并在低噪声区域（精细几何细节恢复处）保留许多小步骤。

机制：噪声调度 $\sigma_i$ 定义为 $\sigma_i = (\sigma_{max}^{1/\rho} + \frac{i}{N-1}(\sigma_{min}^{1/\rho} - \sigma_{max}^{1/\rho}))^\rho$ 。
影响：该策略为少步学生提供了可行的初始化，使其在蒸馏开始前就能生成稳定（尽管不完美）的分子。这稳定了 DMD 训练过程，防止了梯度爆炸并实现了有效的优化。

2. 通过 $f$ -散度正则化提高多样性

DMD 最小化反向 Kullback-Leibler (KL) 散度（ $KL(p_{fake} || p_{real})$ ），这鼓励学生匹配教师的高概率模式，但往往导致“寻模”行为，使支持集坍缩并降低样本多样性。

为了缓解这一问题，FlashMol 使用Jensen-Shannon (JS) 散度项对 DMD 目标进行正则化。

机制：JS 散度结合了前向和反向 KL 分量。前向 KL 分量引入了“寻均值”行为，抵消了反向 KL 的寻模倾向。
实现：JS 项使用在真实分子和生成分子上训练的 GAN 判别器进行近似。最终损失函数为 $L_{FlashMol} = L_{DMD} + \lambda \cdot L_{JS}$ 。

主要贡献

DMD 在分子生成中的集成：作者成功将分布匹配蒸馏适配到潜在分子生成模型 GeoLDM 中，在极少的步数内实现了高保真样本生成。
时间步重排策略：他们识别出推理时的时间步调度是分子分布蒸馏的关键瓶颈。通过引入受 EDM 启发的重排策略，他们在无需额外训练的情况下显著提高了初始少步生成器的质量，从而实现了有效的蒸馏。
统一的散度目标：他们通过集成 Jensen-Shannon 散度项扩展了原始 DMD 目标，提供了一个统一的框架，在通过反向 KL 实现高保真生成的同时，通过前向 KL 分量保持多样性。
性能：该方法仅用 4 到 8 步即取得了最先进的结果，匹配甚至超越了原始 1000 步的教师模型。

实验结果

实验在QM9（13 万个小分子）和GEOM-DRUG（45 万个较大分子）数据集上进行。

稳定性和有效性：
- 在QM9上，FlashMol 仅用4 步即实现了83.89%的分子稳定性，用8 步达到94.87%。这显著优于重排后的 GeoLDM 基线（4 步时为 12.43%）以及其他少步方法，如 GeoRCG（50 步）和 MOLTD（12 步）。
- 在GEOM-DRUG上，FlashMol 用 4 步实现了84.72%的原子稳定性，用 8 步达到89.98%，超过了数据集层面的 86.5% 原子稳定性。
加速：与 1000 步的 GeoLDM 教师模型相比，FlashMol 在采样速度上实现了高达250 倍的加速，同时保持了高分子质量。
条件生成：在条件生成任务（预测 $\epsilon_{HOMO}$ 、 $\epsilon_{LUMO}$ 等属性）中，FlashMol（32 步）实现的平均绝对误差（MAE）与 1000 步基线相当或更优。
消融研究：
- 重排：没有重排时，4 步生成的稳定性接近于零。有了重排，即使不进行蒸馏，稳定性也跃升至约 94%。
- 散度：仅使用 DMD 提高了稳定性但降低了多样性（有效且唯一）。添加 JS 项在保持高稳定性的同时恢复了多样性。

意义与主张

该论文声称，FlashMol 代表了使基于扩散的分子设计适用于高通量筛选的重要一步。通过将采样成本从 1000 步减少到少至 4 步，FlashMol 消除了阻碍这些模型在工业药物发现流程中部署的主要计算障碍。

作者强调，他们的方法并未引入超越教师模型的新型生成能力，而是通过高效蒸馏释放了现有模型的潜力。他们指出，虽然在多样性和稳定性之间存在权衡（这是分数蒸馏目标的已知特征），但他们的方法成功地将模式坍缩缓解到了允许可靠生成有效化学结构的程度。这项工作表明，类似的蒸馏技术可以应用于其他领域的等变潜在扩散模型，例如蛋白质 - 配体共设计和材料发现。