FreeBridge: Variational Schrödinger Bridges for Cellular Transition Dynamics

技術要約：FreeBridge – 細胞遷移ダイナミクスのための変分シュレディンガー・ブリッジ

問題提起
高コンテンツ・イメージング・アッセイは、化学的および遺伝的な摂動に対する細胞の応答を定量化するために不可欠である。しかし、根本的な限界として、画像取得時に細胞は化学的に固定されているため、個々の時間的軌跡を観測することができない。その結果、摂動モデリングは、コントロール群と処置群の間の確率論的な輸送の推論へと還元されてしまうが、そこでは周辺的な終端分布のみが既知である。

近年の生成モデル（拡散モデル、フロー・マッチングなど）は、コントロールと摂動後の終端間の強力な整合性を達成しているが、終端の一致性だけでは中間的な進化を決定するには不十分である。同一の周辺分布を接続する確率過程は無数に存在する。さらなる制約がなければ、これらのモデルは、観察された単一細胞の形態学的に支持されていない潜在空間の領域を横断する可能性があり、生物学的に不自然な中間状態を招く恐けがある。既存の手法は、分布レベルの目的関数を最適化することが多く、その結果、中間的な細胞構造が未決定のままとなり、FIDやKIDといった指標を満たしながらもアーティファクトを生成してしまう可能性がある。

手法：FreeBridgeフレームワーク
著者らは、有効な細胞状態の定義と遷移メカニズムを明示的に分離した「セル・エンジン」として定式化された、制約付き確率輸送フレームワークであるFreeBridgeを提案する。この手法は、エンドポイントのみの監督下で動作し、以下の3つのコアコンポーネントで構成される：

1. 状態ユニット（明示的な単一細胞幾何学）:

   • 画像を集団レベルの混合物として扱うのではなく、FreeBridgeは、インスタンス分割された単一細胞の表現（Cellposeを使用）を原子的な状態として定義する。
   • 固定された学習済みエンコーダ $\phi$ が、これらのクロップを潜在空間 $z = \phi(x)$ へと埋め込む。このエンコーダは、遷移ルールの学習前に定義された固定の潜在幾何学を確立するため、輸送学習中は凍結される。
   • 経験的なサポート・バンク $Z_{emp}$ が、観察された細胞の潜在埋め込みから構築される。このバンクは、生物学的に妥当な形態の「多様体（マニフォールド）」を定義する。

2. 遷移ルール（シュレディンガー・ブリッジ輸送）:

   • コントロール分布 $\mu_z$ と摂動分布 $\nu_{c,z}$ の間の遷移は、制御された確率微分方程式（SDE）としてモデル化される：$dZ_t = u_t(Z_t) dt + \sigma dW_t$。
   • 目的は、時刻 $t=0$ での分布を時刻 $t=1$ でのターゲット分布へと輸送するドリフト場 $u_t$ を見つけることである。
   • 標準的な最適輸送とは異なり、この定式化は経路を正則化するためのタスク固有の状態コストを組み込んでいる。

3. 幾何学的正則化（経験的サポート制約）:

   • 確率過程が潜在空間内の「空隙」を横断することを防ぐため、FreeBridgeはサポート・コスト $V(z) = \lambda \cdot d_{emp}(z)^2$ を導入する。ここで、$d_{emp}(z)$ は経験的サポート・バンク $Z_{emp}$ への最近傍点までの距離である。
   • 全目的関数は、運動エネルギー（ドリフトの大きさ）とサポート・コストの和を最小化する：
     $$\min_u \mathbb{E} \left[ \int_0^1 \frac{1}{2\sigma^2}\|u_t(Z_t)\|^2 dt + \lambda \int_0^1 d_{emp}(Z_t)^2 dt \right]$$
   • この正則化は、ペアとなる細胞レベルの軌跡を必要とせずに、中間状態を経験的に観察された形態学的な多様体に繋ぎ止める。

主な貢献

• 定式化: 単一細胞の遷移モデリングのためのシュレディンガー・ブリッジ定式化を導入し、状態の指定と確率的輸送を明示的に分離した。
• 幾何学的接地: 非パラメトリック密度推定（経験的サポート・バンクへの最近傍距離）に基づく新しい正則化項を提案し、中間的な軌跡を生物学的に妥当な領域内に制約した。
• 統一プロトコル: インスタンス分割された単一細胞クロップを用いた複数のデータセット（BBBC021, RxRx1, JUMP）にわたる統一された評価プロトコルを確立し、手法間の公平な比較を保証した。
• 性能: 幾何学的制約が、終端の忠実度だけでなく、摂動メカニズム（作用機序：Mechanism-of-Action）のセマンティックな保持を向上させることを示した。

結果
BBBC021、RxRx1、JUMPデータセットを用いた評価において、FreeBridgeは以下の結果を示した：

• 終端の忠実度: 最先端のベースライン（PhenDiff, IMPA, CellFlux）と比較して、競争力のある、あるいは改善されたFréchet Inception Distance (FID) および Kernel Inception Distance (KID) スコアを達成した。BBBC021において、最も低いFID (15.35) およびKID (1.10) を達成した。
• 中間の一貫性: 経験的な形態学的サポートから外れた中間状態の割合を測定する指標である**サポート違反率（$R_{viol}$）**を大幅に減少させた。BBBC021において、FreeBridgeは $R_{viol}$ を0.11に減少させた（PhenDiffは0.38、CellFluxは0.31）。
• セマンティックな保持: アブレーション研究により、幾何学的正則化を取り除くと作用機序（MoA）分類精度が低下すること（BBBC021において72.08%から65.08%へ）が示され、中間進化を制約することがメカニズム的な信号の保持に役立つことが示された。
• 定性的分析: 可視化により、FreeBridgeは、テクスチャの平滑化や不完全なパターンを示すことが多いベースラインよりも、化合物特有の形態学的特徴（例：核の凝縮、細胞質の収縮）をより良く保持することが確認された。

意義と主張
本論文は、生物学的な遷移をモデル化するには、境界の整合性以上のものが必要であり、中間的な進化の明示的な幾何学的接地が必要であると主張している。確率的輸送を経験的に支持された細胞多様体に繋ぎ止めることで、FreeBridgeは生物学的に解釈可能な摂動ダイナミクスを生成するための原理的なアプローチを提供する。著者らは、この定式化が「仮想細胞」シミュレーションのための堅牢な基礎を提供し、エンドポイントに整合したモデルが現実的な中間細胞状態を捉えきれないという重要なギャップに対処するものであると断言している。本研究は、幾何学的制約がなければ、生成モデルは統計的に妥当な終端を生成しても、潜在空間内の生物学的に支持されていない領域を横断してしまう可能性があることを浮き彫りにしている。

限界
著者らは、FreeBridgeが凍結されたエンコーダに依存していることを認めている。つまり、エンコーダの事前学習に含まれていない形態学的要因は、輸送学習のみでは回復できない。また、サポート距離の計算はミニバッチサイズに対して二次的（$O(B^2)$）にスケールするため、より大きなサポート・バンクに対しては近似最近傍探索が必要になる可能性がある。