技術的サマリー：リーマン平均フロー

問題定義

拡散モデルとフローモデルは、リーマン多様体上の生成モデリングにおける支配的なパラダイムとなり、タンパク質のバックボーン生成や DNA 配列設計などのタスクに成功裏に適用されている。しかし、これらの手法は通常、推論中に常微分方程式（ODE）または確率微分方程式（SDE）を数値的に積分するために、数十から数百のニューラルネットワーク評価（関数評価、NFE）を必要とする。この計算負荷はボトルネックとなり、特に生成モデルが反復ループ内のプロポーザル機構として機能する、特性誘導最適化のための科学的設計パイプラインにおいて顕著である。一貫性モデルやフローマップ学習などの数ステップ生成技術はユークリッド空間で探求されてきたが、高次元のリーマン多様体設定では未探求であり、かつ不安定であることが多い。

手法：リーマン平均フロー（RMF）

著者は、多様体上で直接フローマップを学習し、1 回のフォワードパスのみで高品質な生成を可能にするフレームワークである**リーマン平均フロー（RMF）**を導入する。この手法の核心は、任意の時間点 $s$ と $t$ 間のフローマップと関連する多様体平均速度の 3 つの同等な特徴付けを導出・利用することにある。

1. 理論的基盤

このフレームワークは、平均速度 $u_{s,t}$ を、時間 $t-s$ にわたって $x_s$ から $x_t$ への測地線に沿って維持された場合、点を正しく輸送する定数速度として定義する。著者は、訓練目的の基礎となる 3 つの恒等式を導出した。

• オイラー的恒等式: 定義関係をソース時間 $s$ について微分することで導出される。これは、平均速度を積分曲線に沿った速度場の共変微分および対数写像の微分と関連付ける。
• ラグランジュ的恒等式: 定義関係をターゲット時間 $t$ について微分することで導出される。これは、平均速度を瞬間速度に適用された対数写像の微分および平均速度の偏時間微分と関連付ける。
• 半群恒等式: フローマップの半群性質（$\Phi_{r,t} \circ \Phi_{s,r} = \Phi_{s,t}$）から導出される代数的条件である。これは、$s$ から $t$ への直接フローと、中間時間 $r$ を通じたフローの合成との間の整合性を強制する。

2. 訓練目的と安定化

著者はこれらの恒等式を回帰目的に変換する。重要な設計選択として、ストップグラディエント演算子の配置がある。従来の研究（例：Generalized Flow Maps）が部分的にストップグラディエントを適用するのに対し、RMF は回帰ターゲット全体にストップグラディエントを適用する。この定式化は以下の利点をもたらす。

• 逆伝播中の高階微分（ヤコビアン - ベクトル積を通じた勾配など）の評価を不要にする。
• メモリ使用量をほぼ 2 倍削減する。
• 高次元設定における訓練の安定性を向上させる。

高次元における最適化の難しさを解決するため、本論文は以下の安定化技術を提案する。

• 時間サンプリング: オイラー的および半群目的には順序付きの時間ペア（$s \le t$）、ラグランジュ的目的には順序なしのペアを使用する。
• 適応的損失重み付け: 微分依存項の分散を処理するため、回帰ターゲットの大きさに基づいて損失を再重み付けする。
• 時間微分制御: 低周波の時間埋め込みを使用して、ネットワーク出力における時間微分の大きさを制限する。

3. パラメータ化

本論文は、平均速度に対する 3 つのパラメータ化方式を評価する。

• $v$-予測: 接ベクトル（速度）を予測する。
• $x_t$-予測: 終点を直接予測する。
• $x_1$-予測: 現在の状態と時間ギャップを与えられた場合、$t=1$ における軌道の終点を予測する。
  著者は、$x_1$-予測が高次元多様体に対して最も堅牢であることを発見した。これは既存の科学的アーキテクチャ（タンパク質向けの SE(3) フレームを出力するものなど）と互換性があり、次元が $D=2048$ までスケールするのに対し、$x_t$-予測は失敗し、$v$-予測は次元が増加するにつれて劣化する。

4. 報酬誘導推論

このフレームワークは、モデルの再訓練なしに報酬誘導生成をサポートする。フローマップが中間ステップから終端状態を予測する能力を活用し、著者は報酬先読み戦略を提案する。現在のノイズ状態 $x_t$ における報酬の勾配を使用するのではなく、予測された終端点 $\Phi_{t,1}(x_t)$ における勾配を使用する。このアプローチは、生成を特定の目標に誘導する上で、単純な誘導や無誘導生成よりも効果的であることが示された。

主要な結果

著者は、RMF を 2 つの生物学的タスクで検証した。プロモーター DNA 設計（$\mathbb{R}^{1024 \times 4}$ 内の単体上の配列）とタンパク質バックボーン生成（$N$ が最大 128 の $SE(3)^N$ 上）である。

• 効率性: RMF は、最大10 倍少ない関数評価で、最先端の多ステップベースライン（Fisher FM、FrameFlow、FrameDiff など）と同等のサンプル品質を達成する。
  • タンパク質バックボーン生成において、RMF は 5 ステップで 82% の設計可能性を達成し、1 ステップでも 35% を達成するのに対し、ベースラインはこれらの少数ステップ領域で急激に低下するか完全に崩壊する。
  • DNA 設計において、RMF の変種は 1 ステップのみで Fisher FM の 100 ステップのパフォーマンスに匹敵する。
• 安定性: 提案された安定化技術（特に半群目的と$x_1$-予測の組み合わせ）は、他のフローマップ手法（GFM など）が収束に失敗するか高誤差のサンプルを生成する高次元空間での安定した訓練を可能にする。
• 報酬誘導: $x_1$ 先読み戦略を用いた報酬誘導推論は、単純な誘導や無誘導生成と比較して、ターゲット調節プロファイル（DNA）および二次構造組成（タンパク質）との整合性を大幅に向上させる。

意義と主張

本論文は、リーマン多様体上の高速サンプリングのための原理的な理論的基盤と実践的なガイドラインを確立すると主張している。その意義は以下の点にある。

1. ギャップの埋め合わせ: 数ステップ生成能力をユークリッド空間から複雑で高次元のリーマン多様体に拡張する。これは科学的応用にとって重要な一歩である。
2. スケーラビリティ: 高階微分を回避する恒等式の導出と$x_1$-予測の利用により、この手法は以前の数ステップ多様体生成モデルが到達できなかった次元（$D \approx 2000$）および配列長にスケールする。
3. 科学的ワークフローにおける効率性: 推論コストの削減（最大 10 倍）は、創薬や合成生物学における特性誘導最適化などの反復的科学的設計ループ内のプロポーザル機構として、生成モデルをより実用的にする。
4. 報酬先読み: フローマップが中間ステップから終端状態を予測することで効率的な報酬誘導設計を可能にすることを示しており、これは多様体上で以前に示されたことがない能力である。

著者は、RMF をすべての生成モデリングの代替として位置づけるのではなく、計算コストが制限要因となる幾何学的に制約された科学的ドメインにおいて、効率的かつ高品質な生成を可能にする専門的なフレームワークとして位置づけている。