Branched Schrödinger Bridge Matching

既存の生成モデルが抱える単一経路の制約を克服し、共通の起源から複数の終端分布への分岐的な進化を表現するために、複数の時間依存速度場と成長過程を学習する新しい枠組み「Branched Schrödinger Bridge Matching（BranchSBM）」を提案する論文です。

要約

この論文は、**「BranchSBM（ブランチ SBM）」**という新しい AI の仕組みを紹介しています。

一言で言うと、**「ある場所から出発した人々が、途中で分岐して複数の異なる目的地へ向かう様子を、AI が完璧にシミュレーションできる」**という画期的な技術です。

従来の AI は「出発点から目的地まで、一本の道でつなぐこと」は得意でしたが、「出発点から複数の道に分かれて、それぞれ違うゴールに行く」ような複雑な動きを再現するのは苦手でした。この新しい技術は、その壁を乗り越えました。

以下に、日常の例えを使ってわかりやすく解説します。

---

1. 従来の AI の限界：「一本道の迷路」

Imagine（想像してみてください）：
ある大きな広場（出発点）に、同じ服を着た大勢の人が集まっています。
彼らは「山を越えて、反対側の村 A と村 B のどちらかへ行く」という任務を持っています。

• **従来の AI（Flow Matching など）**は、広場から村 A へ、そして村 B へ向かう「一本の道」を作ろうとしました。
• しかし、AI は「みんなが同じ道を通るはずだ」と思い込んでいました。
• その結果、AI は「村 A へ行く人」と「村 B へ行く人」を無理やり同じ道に押し込めてしまい、**「村 A へ行くはずだった人が、間違って村 B に行ってしまう」あるいは「全員が村 A だけに行き、村 B には誰も行かない（モード崩壊）」**という失敗をしてしまいました。

これは、細胞が分裂して違う性質になる現象や、薬を投与した細胞が「治る」「悪化する」「変わらない」と分かれる現象を説明するときに、従来の AI だと「全部が同じ結果になる」という誤った予測をしてしまうことを意味します。

2. 新技術「BranchSBM」の仕組み：「分岐する川」

この論文の「BranchSBM」は、**「川が本流から分かれて、複数の支流になり、それぞれ違う湖（目的地）に流れ込む」**というイメージで動きます。

• 本流（スタート）： 最初は全員が一つの川（本流）を流れています。
• 分岐点（ブランチ）： 川の流れの中で、AI が「ここから右へ曲がるべきか、左へ曲がるべきか」を学習します。
• 支流（ゴール）： 右へ曲がった人は「村 A（治る細胞）」へ、左へ曲がった人は「村 B（悪化する細胞）」へと、自然に分かれていきます。

この AI のすごいところは、「どのタイミングで、誰が、どの道へ進むのか」を、最初から決めるのではなく、「川の流れ（成長）」と「分かれ道（速度）」を同時に計算して、最適な分岐パターンを見つけ出す点にあります。

3. 具体的な活用例：3 つのシナリオ

この技術は、以下のような現実世界の複雑な問題を解くために使われます。

① 3D 地図でのナビゲーション（LiDAR）

• 状況： 山のような複雑な地形（LiDAR データ）の上を、スタート地点から「東側の村」と「西側の村」の 2 つの目的地へ移動する必要があります。
• BranchSBM の活躍： AI は「山を登って東へ行く道」と「谷を抜けて西へ行く道」という、2 つの全く異なる最適なルートを同時に発見し、それぞれの人を正しい目的地へ導きます。従来の AI は、どちらか一方のルートしか見つけられませんでした。

② 細胞の運命決定（幹細胞の分化）

• 状況： 人間の体内で、未分化な細胞（幹細胞）が分裂し、「赤血球になる細胞」と「免疫細胞になる細胞」に分かれる過程をシミュレーションします。
• BranchSBM の活躍： 最初は同じ細胞でも、途中で「赤血球コース」と「免疫細胞コース」に分岐します。BranchSBM は、この**「細胞がいつ、どのように分かれるか」**を正確に再現し、それぞれの細胞集団がどう成長するかを予測できます。

③ 薬の反応予測（細胞への刺激）

• 状況： がん細胞に薬を投与したとき、細胞は「死滅する」「耐性を持って生き残る」「変化しない」など、複数の異なる反応を示します。
• BranchSBM の活躍： 薬を投与した瞬間、細胞集団がどのように「分岐」して、それぞれの運命（ゴール）へ向かうかを予測します。これにより、**「この薬を打つと、患者さんの細胞の何％が治り、何％が悪化するのか」**をより正確にシミュレーションできるようになります。

4. なぜこれが重要なのか？

これまでの AI は「平均的な答え」を出すのが得意でしたが、現実の世界（特に生物学や物理現象）では、**「一つの原因から、複数の異なる結果が生まれる（分岐する）」**ことがよくあります。

BranchSBM は、この**「分岐する未来」**を、数学的に完璧に、かつ効率的に再現できる最初のフレームワークの一つです。

• 従来の AI： 「全員が同じ道を行く」と仮定する。
• BranchSBM： 「人それぞれが、最適な分かれ道を選んでいく」と学習する。

まとめ

この論文は、**「AI に『分岐する道』を教える」**という新しいアプローチを提案しました。
まるで、交通整理の達人が、大勢の人を「一本道」で無理やり送るのではなく、「分かれ道」を設けて、それぞれが最適な目的地へスムーズに到着できるように導くようなものです。

これにより、**「細胞がどうなるか」「病気がどう進むか」「薬がどう効くか」**といった、複雑で分岐する現象を、より現実的に、より正確に予測できるようになるでしょう。

技術概要

Branched Schrödinger Bridge Matching (BranchSBM) の技術的サマリー

本論文は、ICLR 2026 で発表された「Branched Schrödinger Bridge Matching (BranchSBM)」という新しい生成モデルフレームワークを提案するものです。従来の確率的輸送手法の限界を克服し、単一の初期分布から複数の異なる目標分布へ分岐する動的な軌跡を学習することを可能にします。

以下に、問題定義、手法、主要な貢献、実験結果、および意義について詳細にまとめます。

1. 背景と問題定義

既存手法の限界

生成モデルにおける重要な課題の一つは、初期分布と目標分布の間の中間軌跡を予測することです。従来のフローマッチング（Flow Matching）やシュレーディンガー・ブリッジ・マッチング（SBM）は、2 つの分布間の単一の確率的経路をモデル化することで有効に機能します。しかし、これらには以下の根本的な限界があります。

• 単一モードへの制約: 既存手法は基本的に「単一経路」を前提としており、共通の起源から複数の異なるモード（分岐）へ発展するダイナミクスを捉えることができません。
• 質量保存の仮定: 多くの手法は初期分布から目標分布への質量保存を仮定していますが、単細胞データなどでは細胞の増殖や死滅（質量の増減）が頻繁に起こります。
• モード崩壊: 分岐する経路を持つシステムにおいて、既存の SBM を適用すると、粒子が dominant な高密度モードに集中したり、低エネルギーの中間経路のみを通過したりする「モード崩壊」が発生しやすくなります。

解決すべき課題

現実世界の多くのシステム（例：遺伝子ノックアウトや薬剤投与後の細胞運命の分岐、群衆の分岐経路）では、均質な集団が複数の異なる運命（フェーズ）に分岐します。この「分岐（Branching）」ダイナミクスを、中間時刻の教師データなしに、かつ質量の増減を許容しながら正確にモデル化する手法が必要とされていました。

2. 提案手法：BranchSBM

BranchSBM は、分岐するシュレーディンガー・ブリッジを学習するための新しいフレームワークです。

2.1 理論的基盤：分岐した一般化シュレーディンガー・ブリッジ問題

著者らは、**分岐した一般化シュレーディンガー・ブリッジ（Branched Generalized Schrödinger Bridge, Branched GSB）**問題を定義しました。これは、以下の要素を組み合わせることで、単一のソースから複数のターゲットへの分岐を表現します。

• 分岐する速度場（Velocity Fields）: 各ブランチごとに独立した時間依存のドリフト場 $u_{t,k}$ を学習します。
• 成長プロセス（Growth Processes）: 各ブランチにおける質量の増減を制御する成長率 $g_{t,k}$ を学習します。
• 不均衡条件付き確率最適制御（Unbalanced CondSOC）: 分岐 GSB 問題は、複数の「不均衡条件付き確率最適制御（Unbalanced CondSOC）」問題の和として定式化されます。これにより、各ブランチの条件付き軌跡を最適化しつつ、全体の質量保存（または増減）を制御できます。

2.2 学習アルゴリズム：4 ステージのトレーニング

BranchSBM は、安定した学習とスケーラビリティを実現するために、4 つの段階に分けたトレーニングアルゴリズムを採用しています。

1. ステージ 1: 分岐神経補間関数の最適化

   • 初期点 $x_0$ と各ターゲット $x_{1,k}$ のペアに対して、状態コスト $V_t$（データ多様体上のエネルギー）を最小化する最適補間軌跡 $\phi_{t,\eta}$ を学習します。
   • これにより、各ブランチの理想的な速度場 $\dot{x}^*_{t,\eta,k}$ が得られます。

2. ステージ 2: 速度ネットワークの初期学習

   • ステージ 1 で得られた条件付き速度 $\dot{x}^*_{t,\eta,k}$ を目標として、各ブランチ $k$ に対応するフローネットワーク（ドリフト場）$u^\theta_{t,k}$ を個別に学習します（条件付きフローマッチング損失）。

3. ステージ 3: 成長ネットワークの学習

   • 速度ネットワークを固定し、分岐のタイミングと質量の分配を決定する成長ネットワーク $g^\phi_{t,k}$ を学習します。
   • 目的関数には以下の損失が含まれます：
     • エネルギー損失 ($L_{energy}$): 軌跡の全エネルギーを最小化。
     • 重み一致損失 ($L_{match}$): 最終時刻 $t=1$ における各ブランチの予測質量と真のターゲット質量の一致を強制。
     • 質量保存損失 ($L_{mass}$): 全ブランチの質量の和が時間的に一定（または所定の増減）となるように制約。

4. ステージ 4: 最終的な共同学習

   • 速度ネットワークと成長ネットワークの両方のパラメータをアンロックし、ステージ 3 の損失に加え、終端分布の再構成損失 ($L_{recons}$) を最小化して微調整を行います。

3. 主要な貢献

1. 分岐 SBM の定義とフレームワークの提案:
   初期分布 $\pi_0$ から複数の重み付き目標分布 $\{\pi_{1,k}\}$ へ分岐する最適軌跡を学習する「Branched Schrödinger Bridge Matching」を初めて導入しました。
2. 理論的定式化:
   分岐 GSB 問題を、不均衡 CondSOC 問題の和として定式化し、その最適解が存在することを理論的に証明しました。これにより、分岐ダイナミクスを数学的に厳密に扱うことが可能になりました。
3. 多様な応用での実証:
   • 3D LiDAR 表面ナビゲーション: 複雑な 3D 地形上での分岐経路学習において、単一ブランチ SBM よりも大幅に高い精度（Wasserstein 距離の低減）を達成しました。
   • 単細胞分化ダイナミクス: マウス造血細胞や膵臓β細胞の分化データにおいて、均質な前駆細胞から複数の細胞運命へ分岐する過程を、中間データなしで正確に再構成しました。
   • 細胞状態の摂動予測: 薬剤投与後の細胞集団が複数の異なる状態へ分岐する現象を、高次元の遺伝子発現データ（最大 150 次元）で高精度にモデル化しました。

4. 実験結果の要点

• LiDAR ナビゲーション: 単一ブランチ SBM は主要な経路のみを学習し、分岐を捉えられなかったのに対し、BranchSBM は山岳地形に沿って最適な分岐タイミングを決定し、両方のターゲット分布に正確に到達しました（$W_1$ 距離が 0.975 → 0.239 に改善）。
• 単細胞分化: 中間時刻（$t_1$）のデータが学習に使われていない場合でも、BranchSBM は正しい分岐軌跡を推論し、単一ブランチモデルが失敗する中間状態と終端状態の両方を高精度に再現しました。
• 摂動モデル: Clonidine および Trametinib による薬剤処理実験において、BranchSBM は単一ブランチモデルが捉えきれない高次元 PCs における複雑な分岐パターンを再現し、MMD や Wasserstein 距離において SOTA 手法（DeepRUOT, CytoBridge など）を上回る性能を示しました。

5. 意義と結論

BranchSBM は、生成モデルの分野において、**「分岐するダイナミクス」と「質量の増減」**を同時に扱うための統一された枠組みを提供します。

• 生物学的意義: 単細胞 RNA シーケンシングデータにおける細胞運命の分岐（フェイト・バイファケーション）や、薬剤応答の不均一性を理解する上で不可欠なツールとなります。中間時刻のデータが不足している現実的なシナリオでも、生物学的に意味のある軌跡を推論可能です。
• 物理的・工学的意義: 群衆シミュレーションやロボティクスにおける分岐経路計画など、物理システムにおける複雑な分岐現象のモデル化にも応用可能です。
• 技術的革新: 従来の SBM が抱えていた「モード崩壊」の問題を、分岐構造を明示的に学習することで解決し、単一のサンプルから分岐する集団ダイナミクスを推論できる点で画期的です。

総じて、BranchSBM は、単一のソースから複数のターゲットへ発展する複雑な確率過程を、エネルギー最小化の原理に基づいて効率的かつ正確に学習する強力な手法として確立されました。