Multi-Marginal Flow Matching with Adversarially Learnt Interpolants

技術要約：敵対的学習を用いた補間によるマルチ・マージナル・フロー・マッチング

問題提起
本論文は、真の軌跡（グラウンドトゥルース）へのアクセスがない状態で、離散的な時点におけるサンプリングされた観測値（スナップショット）のみから、システムの時系列ダイナミクスを学習するという課題に取り組んでいる。これは、生物学（例：シングルセルRNAシーケンシング、細胞追跡）や医学などの分野における中心的なタスクである。形式的には、初期周辺分布 $q_0$ をターゲットとなる $q_1$ へと押し出すような常微分方程式（ODE） $dx_t = v_t(x_t)dt$ を支配する時系列ベクトル場 $v_t$ を推論することが目標である。ただし、これは中間的な周辺制約 $p_{t_i} = q_{t_i}$ （$0 = t_1 < \dots < t_K = 1$ の一連のタイムステップ）を満たす必要がある。

既存の手法（例えば、連続する一対の周辺分布に対して標準的なフロー・マッチング（FM）を適用する方法や、区分線形または三次スプライン補間を用いる方法）は、しばしば非平滑な補間曲線をもたらす。このような非平滑性の問題は、学習目的関数における勾配分散の増大を招き、特に高次元の設定やデータにノイズが含まれる場合に、結果として得られるベクトル場の周辺化を困難にする。

手法：ALI-CFM
著者らは、補間の学習とベクトル場の学習を分離した2段階の手法である ALI-CFM（Adversarially Learnt Interpolants with Conditional Flow Matching：敵対的学習を用いた補間による条件付きフロー・マッチング）を提案している。

1. 補間の敵対的学習 (ALI):
   特定のデータ点を通るように強制する（点ごとのマッチング）のではなく、ALIは中間的な周辺分布の「分布」に一致するように、ニューラル補間関数 $G_\phi(x_0, x_1, t)$ を学習する。この補間関数は、線形補間とニューラル補正項の組み合わせとしてパラメータ化される：
   $$G_\phi(x_0, x_1, t) = (1-t)x_0 + tx_1 + t(1-t)f_\phi(x_0, x_1, t)$$
   補間の押し出し分布が観測された中間的な周辺分布 $q_{t_i}$ と一致することを保証するために、著者らはGANに着想を得た敵対的目的関数を採用している。識別器 $D_\gamma$ は、真の周辺分布 $q_{t_i}$ からのサンプルと、補間関数 $G_\phi$ によって生成されたサンプルを区別する。生成器（補間関数）は、これらの分布間のイェンセン・シャノン距離を最小化するように訓練される。

   補間の一意性を確保し、極端に湾曲した経路を防ぐために、著者らは以下の正則化項を導入している：

   • 線形リファレンス (Linear Reference): 結合された端点間の直線からの偏差を罰する。
   • 区分線形リファレンス (Piecewise Linear Reference): マルコフ連鎖型最適輸送（OT）結合によって定義される区分線形パスへの回帰を行う。
   • 二階微分ノルム (Second Derivative Norm): 補間の曲率を罰することで、平滑性を強制する。

2. 条件付きフロー・マッチング (CFM) による周辺化:
   一度補間関数 $G_\phi$ が訓練されると、それは固定される。次に、これらの補間関数を生成するベクトル場を近似するために、別のニューラルネットワーク $u_\theta$ が訓練される。これは、学習された補間関数の時間微分がターゲットとなるベクトル場となる、標準的な条件付きフロー・マッチングの目的関数を用いて行われる：
   $$L_{CFM} = \mathbb{E}_{t, (x_0, x_1) \sim \pi} \| u_\theta(G_\phi(x_0, x_1, t)) - \frac{d}{dt}G_\phi(x_0, x_1, t) \|^2$$
   ここで、$q_0$ と $q_1$ の間の結合 $\pi$ は、独立であっても最適輸送（OT）プランであってもよい。

主な貢献

1. 新しい補間技術: 点ごとの割り当てではなく、分布のマッチングを用いて補間を学習する手法を導入した。これにより、緩やかな仮定の下で一意性が証明されている通り、滑らかな軌跡の学習が可能となる。
2. ノイズと複雑性への堅牢性: 本手法は、ノイズを含むデータセットや、点ごとの補間法（スプラインや区分線形など）が不安定になるような複雑で時変的な幾何学的構造を扱うように設計されている。
3. スケーラビリティと汎用性: 著者らは、ALI-CFMが高次元データや多数の周辺分布（$K$）に対してもスケール可能であり、特定の複雑なシナリオにおいて既存のベースラインを凌駕しつつ、他のシナリオでは競争力のある性能を維持することを実証している。

実験結果
著者らは、4つの異なるタスクでALI-CFMを評価している：

• 合成「結び目 (Knot)」分布: 1,200個の周辺分布を持つ結び目状の合成データセットにおいて、ALI-CFMのみが正確に時系列の幾何学的構造を捉えることができた。競合する手法（OT-MFM、三次スプライン、区分線形）は、非平滑な補間を生じさせ、高い勾配分散を示した。
• 細胞追跡 (Glioblastoma): 移動する細胞の顕微鏡データを用い、ALI-CFMはノイズを含むサブサンプリングされたデータから時空間ダイナミクスを学習することに成功した。対照的に、OT-CFMの軌跡は発散し、時間依存性のないOT-MFMは時間的依存性を捉えることができなかった。
• シングルセル・トラジェトリー推論 (scRNA-seq): Embryoid Body (EB)、Cite-seq、およびMultiomeのデータセットにおいて、ALI-CFMはアースムーバー距離（EMD）において最先端のベースライン（OT-MFM、OT-CFM）と同等の性能を示した。著者らは、敵対的学習により、これらの特定の指標において点ごとのデータに過学習した手法を厳密に上回ることが困難であるものの、本手法は非常に競争力があることを述べている。
• 腫瘍座標推論 (空間トランスクリプトミクス): 空間的に秩序立てられた組織切片を持つ乳がんデータセットにおいて、ALI-CFMは、保持された（held-out）切片における腫瘍領域の分布予測において、すべてのベースライン（OT-CFM、OT-MMFM、OT-MFMを含む）を大幅に上回った。著者らは、この成功を、マルチモーダルな分布を扱う能力と、組織切片のアライメントにおける課題に対する本手法の能力に帰している。

意義と主張
本論文は、ALI-CFMが、厳格な点ごとの割り当てから分布のマッチングへと転換することで、マルチ・マージナル・フロー・マッチングに対する概念的に新しいアプローチを提供すると主張している。この転換により、滑らかで一意な補間を学習することが可能になり、従来のメソッドが失敗するような、ノイズが多く複雑なダイナミクスをモデル化できる。著者らは、本手法が数百または数千の周辺分布とノイズの多いデータが存在する状況において特に効果的であることを強調しており、そのような領域において、他のどのフロー・マッチング・アルゴリズムも同等の性能を示せていないことを指摘している。標準的なシングルセル・ベンチマークでは既存の最先端手法と同等の性能を示す一方で、空間トランスクリプトミクスや合成的な複雑な幾何学タスクにおける優れた性能は、本手法の汎用性とスケーラビリティを裏付けている。