技術要約：Wasserstein 残差：集団力学からの勾配流の学習

問題提起

本論文は、疎な観測スナップショットの集合から集団力学（population dynamics）を再構成するという逆問題を扱う。計算生物学や群衆力学などの多くの科学分野において、集団の進化は、エネルギー汎関数 $F$ の最急降下によって駆動される確率分布 $\rho_t$ の曲線である**ワッサースタイン勾配流（Wasserstein Gradient Flow, WGF）**としてモデル化される。目標は、観測された周辺分布 $\{q_t\}$ （離散的な観測時刻 $t \in \mathcal{T}_{obs}$ において）に適合する基礎となるエネルギー汎関数 $F$ を復元することである。

既存の主要なアプローチは、エネルギーと最適輸送（OT）ペナルティを最小化する一連の近接ステップを定義する**ジョルダン・キンダーライラー・オットー（JKO）**スキームに依存している。しかし、JKOベースの手法には主に2つの制限がある：

1. 時間離散化に対する柔軟性の欠如： 連続するスナップショット間でコストの高い最適輸送問題を解く必要がある。
2. ギャップへの敏感さ： 大きな時間的ギャップがある場合に苦戦する（例：粒子が正弦波状の谷に沿って移動する場合など、曲がった軌跡を捉えることができず、スナップショット間の線形な弦補間（linear chord interpolation）に依存してしまう）。

手法：ワッサースタイン残差フレームワーク

著者らは、JKOスキームや最適輸送結合を回避する残差ベースのアプローチを提案している。代わりに、勾配流を特徴付ける基本的な微分方程式を、非負の損失関数を介して直接強制する。

1. 残差定式化

本論文では、勾配流の接線（Tangent）およびエネルギー散逸等式（EDE）の定式化を利用する。具体的には、以下の速度制約に焦点を当てる：
$$v_t(x) = -\nabla_x \frac{\delta F}{\delta \rho_t}(x)$$
ここで、$v_t$ は速度場であり、$\frac{\delta F}{\delta \rho_t}$ はエネルギー汎関数の第1変分である。

著者らは、$\rho$ が $F$ の勾配流である場合に限りゼロとなる速度残差 $R_{vel}$ を定義する：
$$R_{vel}[F, (\rho, v)] = \int_0^T \int_{\mathbb{R}^D} \left\| v_t(x) + \nabla_x \frac{\delta F}{\delta \rho_t}(x) \right\|^2 \rho_t(x) \, dx \, dt$$

グローバルな目的関数は、この残差と、観測時刻におけるデータ適合ダイバージェンス $D$ を組み合わせたものである：
$$\mathcal{L}(F, \rho, v) = \lambda R_{vel}[F, \rho, v] + \sum_{t \in \mathcal{T}_{obs}} D(\rho_t, q_t)$$
このフレームワークは、**パス・ファインディング（Path-Finding）やアクション・マッチング（Action Matching）**といった既存の手法を、単一の残差原理のインスタンス化として統一する。

2. 「スティッチング（Stitching）」アルゴリズム

この目的関数を最小化するために、著者らはスティッチングという、シミュレーションフリーで粒子ベースの手法を導入している。その特徴は以下の通りである：

• 学習可能な軌跡： 曲線 $\rho$ が暗黙的なJKO手法とは異なり、スティッチングは曲線 $\rho_\theta$ を、$N$ 個の学習可能な粒子軌跡 $x_{t,k}^\theta$ 上の**カーネル密度推定（KDE）**として明示的にパラメータ化する。
• シミュレーションフリー： ニューラルODEやSDEの積分を必要としない。粒子の位置と速度は、直接的な学習可能パラメータである。
• 経験的近似： 計算を効率化するため、KDEは粒子中心における経験的測度によって近似される。これにより、速度残差は粒子に関する和へと簡略化される：
  \mathcal{L}_{stitch} \approx \int_0^T \sum_{k=1}^N w_k \left\| \dot{x}_{t,k} + \nabla_x \frac{\delta F_\theta}{\delta \rho_\theta_t}(x_{t,k}) \right\|^2 dt + \sum_{t \in \mathcal{T}_{obs}} D(\rho_\theta_t, q_t)
• 関数パラメータ化： エネルギー汎関数 $F_\theta$ は、ポテンシャル項 $V_\theta$、エントロピー項（対数密度）、および相互作用項 $W_\theta$ の総和としてパラメータ化され、これらはすべてニューラルネットワークによって表現される。

主な貢献

1. 残差フレームワーク： 著者らは、ワッサースタイン勾配流の学習を、非負の残差項を最小化することとして再構成し、パス・ファインディングやアクション・マッチングを包含する統一的な視点を提供した。
2. スティッチング・アルゴリズム： 軌跡曲線 $\rho$ を第一級の学習変数として扱う、新しいシミュレーションフリーの粒子法を導入した。これにより、大きなスナップショット間のギャップに対処し、JKOベースの弦予測器が見落とすような曲がった軌跡を捉えることが可能になる。
3. 最先端の性能： 本手法は、特に疎なデータや非対（unpaired）データが存在するシナリオにおいて、優れた軌跡推論ベンチマーク結果を達成している。

実験結果

著者らは、3つの主要なベンチマークでスティッチングを評価している。

1. 例証的例（連続フロー vs. JKO）

合成された「うねった谷（wavy valley）」ポテンシャルにおいて、粒子は正弦波状の軌跡を辿る。

• 結果： JKOベースの手法（特にJKOnet*）は、疎なスナップショット間（$t=10$ から $t=20$）の曲率を捉えることができず、直線的な補間を生じさせる。スティッチングは、曲がった軌跡を正しく復元し、基礎となるポテンシャルの復元において大幅に優れた結果（$R^2 = 0.62$ 対 $0.51$）をもたらした。

2. 合成ポテンシャル復元

本手法は、対（paired）（相関のある軌跡）および非対（unpaired）（無相関のスナップショット）のレジームにおける15個の2次元ポテンシャルに対してテストされた。

• 結果： スティッチングは両方のレジームにおいて一貫した性能を維持する。対照的に、JKOベースの手法は、連続するスナップショット間のOT結合に依存しているため、非対の設定では著しく低下するか崩壊する。スティッチングのエラーは、レジーム間での差異が約2倍以内に収まるのに対し、JKO手法は非対の場合に有益な勾配場を復元できないことが多い。

3. 単一細胞軌跡推論（Embryoid Body データセット）

ヒト胚性幹細胞の分化データ（27日間にわたる5つのスナップショット）に適用された。

• 結果： スティッチングは、Neural SDEs、TrajectoryNet、および様々なJKOベースの手法を含むベースラインと比較して、最高の性能（最小の $W_1$ および $W_2$ 距離）を達成した。これは、フルデータ学習および「leave-two-out」の時間的補間プロトコルの両方において、競合手法を凌駕している。

4. 相互作用力学の復元

McKean–Vlasov 系における環境ポテンシャル ($V$) と相互作用カーネル ($W$) を分離するために本手法が使用された。

• 結果： スティッチングは、拘束ポテンシャルと引力的相互作用カーネルの両方を正常に復元し、高いパターン $R^2$ スコア（$V$ で $0.84$、$W$で$0.81$）を達成した。また、各ステップでOT問題を解く必要があるJKOベースの手法よりも、1エポックあたりの計算速度が大幅に速い。

5. 非勾配流

著者らは、勾配項を一般的なベクトル場に置き換えることで、残差フレームワークを勾配流を超えて拡張でき、カイラル（非保存的）な力学を追跡できることを示している。

意義と主張

本論文は、ワッサースタイン残差の視点が、集団力学を学習するための、JKOスキームよりも柔軟で堅牢な代替案を提供すると主張している。最適輸送結合の制約からエネルギー汎関数の学習を切り離すことにより、スティッチングは以下を可能にする：

• 疎なデータへの堅牢性： 線形補間に固有の「曲率の喪失」なしに、広く離れた観測から力学を学習する能力。
• 非対データの互換性： 生物学的イメージングで一般的な、時間経過に伴って粒子の同一性が追跡されないデータセット（非対スナップショット）から学習する能力。
• シミュレーションフリーの効率性： 学習中にODEやOT問題を解く計算オーバーヘッドを回避する。

著者らは、本研究を既存のパラダイムの統一および実用的な進歩として位置づけている。同時に、相互作用項における $O(N^2)$ の計算量や、$N \to \infty$ における理論的な収束保証の欠如といった限界についても認めている。