技術要約：Reward Transport – Flow Matchingにおけるノイズ空間のアライメントによる特性制御

1. 問題提起

Flow Matching (FM) および関連する拡散ベースの生成モデルにおいて、結合（coupling）（どのノイズベクトルをどのデータポイントとペアリングするかを決定する規則）は、通常、計算上のヒューリスティックとして扱われます。標準的なアプローチでは、単純化のために独立したペアリングを用いるか、あるいは軌道を直線化し分散を抑えるためにミニバッチ最適輸送 (OT) を用います。これらのパラダイムでは、結合は訓練を容易するためだけに最適化されるのであり、特定の構造的特性を学習されたフロー場にエンコードするためのものではありません。

その結果、制御可能な生成（例：特定のlogPやQEDスコアを持つ分子の生成）を実現するには、通常、以下のような推論時の明示的なメカニズムが必要となります：

• Classifier-free guidance: ターゲットとなる特性値を入力として必要とする。
• 勾配ガイダンス (Gradient guidance): サンプリングの各ステップで特性予測器の勾配を計算する。
• 強化学習 (RL): オラクルに対して生成器を微調整するが、これはサンプリング中にオラクルのクエリを必要とすることが多い。

これらの手法は、推論時のオーバーヘッド、外部の報酬モデルやオラクルへの依存、あるいは複雑な訓練パイプラインを導入します。本論文は、結合自体をアライメント・インターフェースとして機能させ、訓練中に制御可能な構造をフロー場に直接埋め込むことで、追加の計算やオラクル呼び出しなしに制御を可能にできるのではないか、という問いを投げかけています。

2. 手法：Reward Transport

著者らは、ノイズとデータの結合を、タスクの構造をフロー場に注入するためのメカニズムとして再定義する手法である Reward Transport を提案しています。

2.1 特性整合型最適輸送結合

ノイズとデータを独立してペアリングする代わりに、Reward Transport はターゲットとなる分子特性 $y(x)$（例：logP, QED）に基づいた単調結合 (monotone coupling) を構築します。

1. スカラーキー: 各ノイズベクトル $z$（トークン埋め込みのシーケンス）に対し、スカラーソートキー $s(z)$ を計算します。主要な選択肢は、平均プーリングされたノイズのユークリッドノルムです：$s(z) = \|\bar{z}\|_2$。
2. ランクに基づくペアリング: ノイズベクトルの集合 $\{z_n\}$ と、特性 $\{y_n\}$ を持つ分子の集合 $\{x_n\}$ が与えられたとき、両方の集合をそれぞれのキー（$s$ および $y$）によってソートします。その後、ランクごとにペアリングを行います。すなわち、$k$ 番目に小さい $s$ を持つノイズベクトルは、$k$ 番目に小さい $y$ を持つ分子とペアになります。
3. 訓練目的関数: フローマッチングモデル $v_\theta$ は、これらの整合されたペアを輸送するように訓練されます。損失関数は、エンドポイント予測（連続埋め込み空間内）に対する標準的なMSEと、デコードされたトークンが有効な分子を形成することを保証するためのクロスエントロピー（CE）損失を組み合わせたものです。

2.2 方向条件付け (Direction Conditioning)

推論時にモデルがノイズ空間の座標 $s$ にアクセスできるようにするために、Direction MLP が導入されます。この2層のMLPは、正規化されたスカラー $s$ を $d$ 次元の埋め込みへと写像し、これを時間および位置埋め込みと共にトランスフォーマーの入力に加算します。これにより、モデルは単調結合で使用された特定のノイズ座標に基づいて、自身の速度場を条件付けることが可能になります。

2.3 推論メカニメント

推論時、ユーザーはスカラー方向信号 $s^*$（ターゲットの特性値ではなく、ノイズ空間の座標）を指定します。モデルは、$s(z) \geq s^*$ を満たすノイズベクトルから開始して、フロー場を積分することで分子を生成します。

• 理論的保証 (Proposition 1): 学習されたフローが $s$ と $y$ の間のランク対応関係を完全に保持する極限（結合保存極限）において、条件付け $s \geq s^*$ は、クロスエントロピー法 (CEM) の1ステップの選択と同等の分布を誘導します。具体的には、生成分布を、閾値 $s^*$ によって決定される上位 $\epsilon$ 分画に切り詰めます。
• ゼロ・オーバーヘッド: この制御には、オラクル、報酬モデル、勾配計算、または拒絶サンプリングを必要としません。

3. 主な貢献

1. アライメント・インターフェースとしての結合: 本論文は、フローマッチングにおけるノイズとデータの結合を、単なる訓練上のヒューリスティックから、分布レベルのアライメント・インターフェースへと再定義しています。特性に基づく結合を選択することで、学習されたフロー場が、ノイズ空間の座標と特性値が対応するように整理されることを示しています。
2. 単一ノブによる分子ステアリング: 本手法は、単一のスカラーノイズ座標を用いて、分子特性（logPおよびQED）に対する単調な制御を実現します。極めて重要な点は、同じノブが異なるターゲットに対して逆の構造的応答を生み出すことです。すなわち、増加する $s$ は、logP（親油性）に対してはより大きな分子を生成しますが、QED（薬物らしさ）に対してはより小さな分子を生成します。これは、汎用的なサイズバイアスではないことを示しており、結合がターゲット固有の構造を書き込んでいることを裏付けています。
3. 離散シーケンスに対するReward Transport: 著者らはこれを可変長のSELFIESトークンシーケンスに実装し、離散データを生成するフローモデルにおける2つの具体的な失敗モードを特定し、対処しました：
   • 潜在的な信号消失 (Implicit Signal Erasure): Pre-LayerNormアーキテクチャでは、大きさの情報（$\|\bar{z}\|$ にエンコードされたもの）が消失します。Direction MLPは、学習された方向パターンを注入することでこれを解決します。
   • 長さの混同 (Length Confound): パディングトークンに対するアンマスクMSE損失がない場合、モデルは高 $s$ が単にノルム制約を満たすためにシーケンス長を増やすというショートカットを学習してしまいます。アンマスクMSEは、パディング位置における生のノイズを回帰することを強制し、大きさの信号を固定することで、このショートカットを防ぎます。

4. 実験結果

手法は ZINC-250K および GuacaMol データセットを用いて評価されました。

• 制御性能:
  • logP: ZINC-250Kにおいて、$s$ を -3 から +7 までスイープすると、平均logPは 1.50 から 5.44 へと変化しました（+137%）。$s$ と平均特性の間のスペアマン相関係数は $\rho=1.000$ でした。効果量（Cohen's $d$）は 1.63 に達しました。
  • QED: 同じノブがQEDを単調に変化させました（例：ZINCにおいて $s=5$ で +0.059）。$\rho=1.000$ でした。
  • 構造的特異性: logPについては、$s$ の増加に伴い重原子数が増加しました（12 $\to$ 23）。QEDについては、$s$ の増加に伴い重原子数が減少しました（23 $\to$ 16）。
• 汎化性能: GuacaMol (1.6M 分子) でも結果が再現され、logPではさらに強力な効果量（$d=2.31$）を示し、有効な範囲内で一貫した単調性（QED）を示しました。
• ベースラインとの比較:
  • Conditional Flow Matching（特性値が明示的な入力となる手法）と比較して、Reward Transport は同様の単調性（$\rho=1.000$）を達成しますが、明示的なターゲットへの適合ではなく、分布の切り詰めを通じて動作します。
  • $\epsilon$-prediction diffusion と比較すると、本手法は、$\epsilon$-prediction においてはベイズ最適ターゲットが結合に依存しないため、結合レベルのアライメントが構造的に欠如していることを示しています。
• 訓練ダイナミクス: ランク保存（$s$ と生成された $y$ の間の相関）は、微調整中（1エポック以内）に急速に現れ、1.0 未満（例：logPで 0.57）でプラトーに達します。これは Proposition 1 で述べられている「分布近似」のギャップと一致しています。

5. 意義と主張

本論文は、Reward Transport が、フローマッチングにおける結合を、それ自体がアライメント構造を運ぶことができる、明確な分布レベルの制御チャネルとして確立したと主張しています。

• 効率性: 推論時のオラクル呼び出し、報酬モデル、または反復的な最適化を排除し、RLやガイダンスベースの手法と比較して計算コストを大幅に削減します。
• メカニズム: 生成モデルの「操縦可能性（steerability）」が、必ずしも明示的な条件付け入力に依存する必要はないことを示しています。それは、訓練時の結合を通じて、モデルの幾何学的な構造の中に焼き付けることができるのです。
• 限界と範囲: 著者らはその範囲について謙虚です。本手法は、1次元の分布的なノブ（分布の切り詰め）を提供するものであり、個々のサンプルに対する正確な値の割り当てを行うものではありません。また、サイズや極性といったグローバルな分子特性に関連するスカラー特性にはうまく機能しますが、類似のサイズを持つ複雑な構造を区別できない合成容易性（SA）のようなトポロジカルな特性には適していません。さらに、本手法が $\epsilon$-prediction ターゲットとは互換性がないことも示されており、このインターフェースがネットワークの予測対象に依存していることを浮き彫りにしています。

要約すると、本研究は、制御可能な生成のパラダイムを「推論時の軌道のガイダンス」から「訓練分布のアライメント」へとシフトさせ、フローマッチングにおける特性制御のための、原理に基づいたゼロ・オーバーヘッドのメカニズムを提示しています。