EmbedOpt: Robust Inference-Time Steering of Protein Diffusion Models via Embedding Optimization
Minhuan Li, Jiequn Han, Pilar Cossio, Luhuan Wu (Flatiron Institute) — arXiv:2602.05285, 2026年2月
🎯 推論時に条件付き埋め込み空間を最適化することで、cryo-EM / NMR 制約に整合したタンパク質構造を安定生成
① 背景と課題

タンパク質の立体構造決定は生物物理学の逆問題として定式化できる。cryo-EM マップや NMR 距離制約 y から、物理的に妥当かつ観測値に整合する構造 x₀ を復元することが目標である。AlphaFold3 に代表される事前学習済み条件付き拡散モデルが強力な事前分布 p(x₀|c) を提供し、後験 p(x₀|y,c) からのサンプリング（Diffusion Posterior Sampling, DPS）が主流となっている。

DPS は事前分布と尤度の重複が小さい困難ターゲットで尤度を強く重み付けする必要があり、数値的に不安定なサンプリング景観を生む
DPS の性能はハイパーパラメータ α に極めて敏感。α ≥ 0.1 で構造が物理的に崩壊（MolProbity スコア急増）

→ 座標空間ではなく 埋め込み空間（共進化情報を圧縮した意味空間）で誘導することで、事前分布と実験制約のミスマッチを解消する

② EmbedOpt アルゴリズム概要
xₜ, cₜ ▶ denoiser ▶ x̂₀ ▶ R(x̂₀) ▶ ∇cₜR ▶ cₜ₋₁ ← cₜ + α·RMSnorm(∇cₜR) ▶ xₜ₋₁
  • 各拡散ステップで埋め込み (s, z) を1回の勾配上昇で更新
  • RMS 正規化で single / pair 埋め込みのスケール差を補正
  • 更新後の cₜ₋₁ で改めて denoiser を呼び出し座標 xₜ₋₁ を算出
  • 事前学習パラメータは一切変更しない（推論時のみ介入）
理論保証
αₜ・∆σₜ が十分小さければ各ステップで surrogate 報酬が単調増加（Trust-region 命題）
③ DPS との比較（幾何学的観点）
特性DPSEmbedOpt
更新空間座標 xₜ埋め込み c = (s,z)
JacobianJ⊤xt（ノイズ座標感度）Jct J⊤ct（埋め込み幾何）
α 安定域～1桁（α≥0.1で崩壊）～2桁（広いプラトー）
Hard target 性能劣る（prior 偏り強）優れる（prior 再シフト）
構造品質（高α）物理的崩壊正常を維持
denoiser 呼出回数1回/ステップ2回/ステップ
④ 主な結果 (a) cryo-EM マップ相関
Best-Achieved マップ相関係数（15ターゲット平均、高いほど良い） Prior (no guidance) 0.43 DPS (tuned) 0.61 EmbedOpt ✓ 0.74 困難ターゲット（初期相関<0.5）で特に顕著な改善

cryo-EM マップ適合でEmbedOptがDPSを明確に上回る

④ 主な結果 (b) ハイパーパラメータ感度
学習率 α ごとのマップ相関（安定域の広さ比較） 0.001 0.01 0.1 1.0 10 EmbedOpt（安定） DPS（α≥0.1で崩壊） 相関係数 EmbedOpt 安定域（2桁）

EmbedOptは2桁の学習率範囲で高性能・安定。DPSはα=0.1超で急劣化

⑤ 限界点・残る課題
  • AF3 スタイルの埋め込みアクセス可能なモデルに限定（汎用拡散モデル非対応）
  • 各拡散ステップで denoiser 2回呼出し → DPS 比でメモリ・計算コスト増
  • 複数実験制約の同時充足（マルチ目的報酬設計）は未検討
  • 粒子フィルタリングとの統合・収束保証の強化が今後課題
⑥ ケムインフォ活用シナリオ（lib/md・lib/docking）
lib/md 初期構造生成
NMR/cryo-EM 制約を条件として EmbedOpt で AF3 から精密構造生成 → RMSDAnalyzer・HBondAnalyzer の入力品質向上
lib/docking ポーズ精度
Induced-fit 系でターゲットポケットを EmbedOpt で cryo-EM 密度整合させてから UniDockRunner に投入
lib/molgen 生成条件制御
特定コンフォーメーションのポケット形状を条件として JobManager による分子生成に組み込む
実装ポイント
Protenix (OSS) バックボーン。PyTorch retain_grad() で埋め込み勾配を取得。推論時のみ介入（fine-tuning 不要）
⑦ 実装優先度・まとめ
  • Priority: High — MD 初期構造品質改善への直接応用
  • target_module: lib/md（lib/docking・lib/molgen にも波及）
  • EmbedOpt の単一ステップ更新は memory-efficient（BPTT 不要）
  • OSS Protenix で即再現可能（ByteDance AML GitHub）
X 投稿用（120字）
推論時に埋め込み空間を最適化するEmbedOpt。AlphaFold3拡散モデルをcryo-EM/NMR制約に整合させ、DPSより安定かつ頑健。ハイパーパラメータ2桁分の安定域。 #StructurePrediction #DiffusionModels
arXiv:2602.05285 | Flatiron Institute | 2026-02-05 | Category: machine_learning | lib/md, lib/docking, lib/molgen