技術的サマリー：FLUX – 混合専門家モデルを用いた幾何学的意識を持つ縦断フローマッチング

1. 問題定義

神経回路や分化する細胞集団など、多くの生物学的システムは、学習、刺激コンテキスト、または発達段階によって定義される潜在レジーム間を遷移しつつ、連続的な局所ダイナミクスを通じて進化します。しかし、これらのシステムは頻繁に非対の縦断スナップショットとしてのみ観測されます。広視野カルシウムイメージングや単一細胞 RNA シーケンシングなどのシナリオでは、同じ個々の細胞やニューロンを時間点間で追跡することはできません。代わりに、研究者は marginal 分布の系列 ($\mu_0, \mu_1, \dots, \mu_{T-1}$) を観測します。

この設定は、2 つの絡み合った課題を提示します：

1. 幾何学的制約: 高次元の生物学的測定値は、しばしば低次元の曲がった非線形多様体の近くに集中します。非対の端点間の標準的なユークリッド補間関数は、低密度領域を通過し、生物物理学的制約に違反し、状態空間の内在的な幾何学を無視する可能性があります。
2. レジーム発見: 基盤となる輸送メカニズム（「フロー」）は、時間とともに変化することがよくあります。例えば、神経活動は、広範な刺激・報酬期待から特定の弁別へとシフトしたり、細胞は多能性から系譜コミットメントプログラムへと遷移したりします。これらは、単なる空間的クラスターではなく、局所ダイナミクス（レジーム）の変化です。既存の手法はしばしば輸送精度に焦点を当てていますが、速度場を分解してこれらの潜在ダイナミックレジームを特定することには失敗しています。

2. 手法：FLUX

著者は、連続的な輸送と教師なしレジーム発見を共同でモデル化するフレームワークであるFLUX（Unpaired longitudinal data with miXture-of-experts に対するフローマッチング）を導入します。この手法は、3 つの連続的な段階で動作します：

A. 幾何学的意識を持つ条件付きパス

FLUX は、Metric Flow Matchingを縦断設定に拡張します。隣接する marginal 間の直線的なユークリッドパスを使用する代わりに、FLUX はプールされた観測（ラベルなしサンプルを含む）からデータ依存のメトリックを学習します。

• メトリック学習: 幾何学バックエンド（RBF または RBF-MLP デープカーネルメトリックを使用）は、多様体からの距離スコアを推定し、データ多様体からの経路を罰するメトリック $G(x)$ を定義します。
• ベンドネットワーク: ネットワーク $B$ は、端点 $x_k$ と $x_{k+1}$ 間の近似幾何学的意識パス $\gamma_k(\alpha)$ をパラメータ化します。このパスは多様体上に留まるように曲がり、局所幾何学を尊重するトレーニングポイント $z_{k,\alpha}$ と目標接線 $\dot{z}_{k,\alpha}$ を提供します。

B. レジーム切り替え速度場

FLUX は、グローバルな速度場 $v_\theta$ を**混合専門家モデル（MoE）**に分解します。

• 専門家ベクトル場: モデルは $M$ 個のスパースな専門家ベクトル場 $f_m$ を学習します。
• ルーター: ルーターネットワーク $g(t, x)$ は、時間と状態を専門家のログオッズにマッピングします。
• ルーティング機構: 訓練中は、Straight-Through Gumbel-Softmax機構が各サンプルに対してアクティブな専門家を選択し、離散的なルーティング決定の微分可能な最適化を可能にします。推論時には、レジームは決定論的なハード割り当て（argmax）によって決定されます。
• 目的: 速度場は、幾何学的意識パスの接線と一致するように訓練されます。ルーターは、スパースで時間的に一貫した割り当てを促進し、専門家の崩壊を防ぐために正則化されます。

C. 訓練パイプライン

安定性を確保するために、訓練は 3 つの段階で進行します：

1. メトリック学習: プールされたサンプル上で幾何学バックエンドを訓練します。
2. ベンドネットワーク訓練: 凍結されたメトリックを用いてベンドネットワークを訓練し、幾何学的意識補間関数を生成します。
3. 速度場訓練: 幾何学とベンドネットワークを凍結し、生成されたパスと接線を用いて速度専門家とルーターを訓練します。

3. 主要な貢献

1. 縦断 marginal への拡張: 著者は、軌跡ラベルが疎であるがラベルなし観測が利用可能な場合、教師なし多様体学習の役割を強調しながら、幾何学的意識フローマッチングを順序付けられた marginal 分布の連鎖に拡張しました。
2. 速度分解によるレジーム発見: 彼らは、スパースな専門家ルーティングが静的な観測を事後にクラスター化するのではなく、局所輸送構造から直接潜在ダイナミックレジームを発見する、レジーム切り替え速度場を導入しました。
3. 幾何学とレジームの相互作用: この研究は、特にレジームが静的な空間的分離ではなく局所ダイナミクスの変化によって定義される場合、幾何学的学習がレジーム発見に不可欠であることを実証しました。
4. 生物学的検証: このフレームワークは、2 つの異なる生物学的設定で検証されました：
   • 神経データ: 連合学習における早期/中間訓練と後期訓練の分離。これは行動の分岐と一致します。
   • 発生生物学: 胚性体における多能性と分化した細胞集団に関連する発現進化レジームの回復。

4. 結果

本論文は、Stanford Bunny（制御された多様体）、Lorenz アトラクター（レジーム切り替えダイナミクス）、広視野カルシウムイメージング（神経学習）、および胚性体分化の 4 つのデータセットで FLUX を評価しました。

• Stanford Bunny: FLUX は、ユークリッド法が多様体上に留まることができない高次元埋め込みにおいて、表面輸送を保持します。
• Lorenz アトラクター: FLUX は、カオス的および臨界下レジーム間の切り替えを識別することにより、完全なレジーム回復（ARI=1.0、NMI=1.0）を達成します。重要なのは、多様体学習なしの FLUX（ユークリッドパスを使用）は、同じ MoE アーキテクチャを持ちながら、単一の専門家に崩壊し、レジームを発見できないことです。これは、幾何学的意識パスが専門家特化に必要な曲率と速度のシグネチャを露出させることを証明しています。
• 神経データ: FLUX は、行動的弁別（CS+ 対 CS−）の出現と整合する皮質ダイナミクスの遷移を特定し、ユークリッドルーティングおよびスイッチングダイナミクスベースラインを上回ります。
• 胚性体: ここではユークリッドルーティングが良好に機能します（段階が空間的に分離可能であるため）。しかし、幾何学的学習を備えた FLUX は競争力があり、レジームが局所ダイナミクスに符号化されている場合、幾何学が最も重要であることを再確認させます。

アブレーション調査: 論文は明示的に、混合専門家ルーティングだけでは不十分であると述べています。幾何学的学習なしでは、モデルは局所輸送に適合できますが、それらのレジームが静的なクラスターではなく局所ダイナミクスに符号化されている場合、レジームを回復することはできません。

5. 意義と主張

著者は、幾何学的意識速度分解が、非対の縦断スナップショットから潜在生物学的状態遷移を発見するための一般的な戦略を提供すると主張しています。

• 科学的影響: FLUX は、輸送精度の最適化（ワッサーシュタイン距離の最小化）から、解釈可能なダイナミック構造の回復へと焦点を移します。これは、生物学的システムが学習や分化などの過程において、しばしば明確なダイナミックレジームに再編成されることを示唆しています。
• 限界: 著者は、FLUX が 3 つの異なるネットワークを訓練するため、ユークリッドフローマッチングよりも計算コストが高いことを認めています。また、専門家の数は事前に固定されており、局所ダイナミクスがノイズが多い場合やサンプリングが弱い場合、推定されたレジームは慎重に解釈されるべきであると指摘しています。
• 結論: 本論文は、FLUX が非対の生物学的スナップショットの解釈可能なダイナミックモデリングのフレームワークを提供し、幾何学的多様体学習とレジーム切り替えダイナミックシステムの間のギャップを埋めると結論付けています。