FlashMol: High-Quality Molecule Generation in as Few as Four Steps
画像由来の分布マッチング蒸留(DMD)を 3D 分子生成へ初適用し最大 250× 高速化 ｜ arXiv:2605.07020 [cs.LG] (2026)
🎯 1000 ステップ GeoLDM 教師を 4〜8 ステップに蒸留し、教師に匹敵/凌駕する品質を最大 250 倍の高速化で達成する（少数ステップ高品質 3D 分子生成）。
① 背景と課題

拡散ベース 3D 分子生成(EDM/GeoLDM)は妥当・安定な構造を生むが 1 サンプルあたり数百〜千 NFE を要し、ハイスループット in silico スクリーニングには遅すぎる。近年の少数ステップ手法は NFE を 12〜50 に削減するが安定性を大きく損ない、10 ステップ以下で信頼できる生成は存在しなかった。FlashMol は DMD を分子に初適用してこれを解く。

拡散 3D 分子生成は 1 サンプルあたり数百〜千 NFE を要し、大規模スクリーニングに遅すぎる
軌跡蒸留（progressive/consistency）は軌跡長が伸びると誤差が累積する
DMD をそのまま適用すると ①初期化敏感で勾配爆発、②mode-seeking で多様性低下（特に QM9）

→ reverse-KL の DMD を分子へ移植しつつ、EDM 由来リスペースで初期化を安定化＋JS 正則化で多様性を確保

② 手法: FlashMol（Geometric DMD）
Geometric DMD（DMD2 三モデル） G_θ 少数ステップ生成器 K-step consistency sampling (4–8 NFE) μ_real 凍結教師スコア 1000-step GeoLDM μ_fake 代理スコア 拡散損失で 5 step 更新 DMD 勾配 s_real − s_fake で G 更新 + リスペース + JS 正則化
L = L_DMD + λ·L_JS （λ=0.1）
EDM リスペース(ρ=2.25)・SE(3)不変勾配(命題1)・GeoLDM 潜在拡散をバックボーン

リスペースは高ノイズ側に大ステップ・低ノイズ側に小ステップを配分し結合長/角度の詳細を回復。公開実装なし。

③ 本研究で示したこと（要点）
  • 分布マッチング蒸留(DMD)を 3D 分子生成へ初適用（軌跡蒸留の誤差累積を回避）
  • EDM 由来リスペースタイムステップ(ρ=2.25)で初期化を安定化（勾配爆発を防止）
  • JS ダイバージェンス正則化で reverse-KL の mode-seeking を相殺し多様性を確保
  • 1000 ステップ GeoLDM 教師に匹敵/凌駕しつつ最大 250× 高速化
④ 主な結果 (a) QM9 安定性
83.9 12.4 4 NFE 94.9 70.8 8 NFE FlashMol GeoLDM-respaced QM9 分子安定性 (%)

8 NFE で 94.87%（respaced 70.78%）、4 NFE でも 83.89%（respaced 12.43%）。

④ 主な結果 (b) 高速化
250× 4 step 200× 5 step GeoLDM(1000step) 比 高速化

教師に匹敵/凌駕する品質を最大 250 倍速で達成。

④ 主な結果 (c) vs 既存少数step
94.87 FlashMol(8) 92.53 MOLTD(12) 89.08 GeoRCG(50) 分子安定性 vs NFE (%)

最小 NFE で最良の安定性（括弧内が NFE）。GEOM-DRUG は 8 step で原子安定性 89.98%。

④ 主な結果 (d) アブレーション
ρ = 2.25 が最適
リスペース指数に敏感。ρ=5 では分子安定性 35.23% へ急落
DMD + JS が最良
V&U: DMD 単独 64.84 → +JS 72.70（mode-seeking を相殺）
リスペース効果
蒸留なしでも 16 NFE で 94.61%（DDIM 一様は 41.84%）
⑤ テイクホームメッセージ
⚡ 250×高速化
1000 ステップ GeoLDM を 4〜8 ステップに蒸留。
🎯 初期化安定
EDM リスペース(ρ=2.25)で勾配爆発を防ぎ低ノイズ詳細を回復。
🌈 多様性確保
JS 正則化で reverse-KL の mode-seeking を相殺。
🧪 教師超え
8 NFE で QM9 94.87%、GEOM-DRUG 原子安定性はデータセット水準超。
少数ステップ生成比較 (QM9)
手法NFE分子安定性
GeoLDM-respaced870.78%
GeoRCG5089.08%
MOLTD1292.53%
FlashMol（本研究）894.87%
本研究のインパクト
  • lib/molgen: GeoLDM ベースの 3D 分子生成器を DMD2 三モデルループで蒸留し、JobManager/MolgenYaml の生成バックエンドを 4〜8 NFE 化。ハイスループット in silico スクリーニングの生成段を桁違いに高速化
  • lib/docking・lib/fep: 少数ステップでも結合長・角度を保つため、生成→ドッキング→FEP の一気通貫スクリーニングのスループットを底上げ
  • リスペース＋JS 正則化は他の幾何拡散(EDM/GEOM 系)へも移植可能な汎用蒸留レシピ。ただし ρ はデータセット依存で再チューニングが必要、教師性能の上限は超えない
FlashMol: High-Quality Molecule Generation in as Few as Four Steps. arXiv:2605.07020 [cs.LG]. 2026. https://arxiv.org/abs/2605.07020