[論文レビュー] Simulation-free Schrödinger bridges via score and flow matching

Alexander Tong, Nikolay Malkin+6 more|arXiv (Cornell University)|Jul 7, 2023

Model Reduction and Neural Networks被引用数 6

ひとこと要約

本稿では、エントロピー最適輸送を通じてスコアマッチィングとフローマッチィングを統合することにより、任意のソース分布とターゲット分布間の確率的ダイナミクスをシミュレーションフリーに学習する[SF]2Mを提案する。この手法により、高次元の細胞ダイナミクスの正確でスケーラブルなモデリングが可能となり、SDEシミュレーションを必要とせず単細胞データから既知の遺伝子調節ネットワークを回復できる。

ABSTRACT

We present simulation-free score and flow matching ([SF]$^2$M), a simulation-free objective for inferring stochastic dynamics given unpaired samples drawn from arbitrary source and target distributions. Our method generalizes both the score-matching loss used in the training of diffusion models and the recently proposed flow matching loss used in the training of continuous normalizing flows. [SF]$^2$M interprets continuous-time stochastic generative modeling as a Schrödinger bridge problem. It relies on static entropy-regularized optimal transport, or a minibatch approximation, to efficiently learn the SB without simulating the learned stochastic process. We find that [SF]$^2$M is more efficient and gives more accurate solutions to the SB problem than simulation-based methods from prior work. Finally, we apply [SF]$^2$M to the problem of learning cell dynamics from snapshot data. Notably, [SF]$^2$M is the first method to accurately model cell dynamics in high dimensions and can recover known gene regulatory networks from simulated data. Our code is available in the TorchCFM package at https://github.com/atong01/conditional-flow-matching.

研究の動機と目的

• ガウス分布ソースやシミュレーションを多く要する学習に依存する既存のスコアベースおよびフローベースモデルの限界を克服すること。
• トレーニング中に確率過程のシミュレーションを必要とせず、任意のソース分布とターゲット分布間の確率的ダイナミクスを学習できること。
• 反復的なSDEシミュレーションを要する伝統的なシュレーディンガー・ブリッジ手法のスケーラビリティと数値的非効率性を克服すること。
• スコアマッチィング（拡散モデル）とフローマッチィング（ノーマライジングフロー）を任意のソース分布に一般化する統合フレームワークの構築。
• アンパイルドスナップショットデータから、高次元の生物学的システム（例：単細胞遺伝子発現ダイナミクス）の正確なモデリングを可能とすること。

提案手法

• エントロピー最適輸送を用いてシュレーディンガー・ブリッジ問題を混合ブラウン運動ブリッジとして定式化し、静的でシミュレーションフリーな目的関数を可能にする。
• Sinkhornアルゴリズムを用いて、ソース分布とターゲット分布間の静的エントロピー最適輸送マップを効率的に計算し、動的SDEシミュレーションに置き換える。
• ニューラルネットワークをトレーニングし、連合的な確率的回帰目的関数において、ODEのドリフト（フローマッチィング）と条件付きスコア（スコアマッチィング）の両方を予測する。
• 中間点をブラウン運動ブリッジの周辺分布からサンプリングする確率的回帰目的関数を定義し、モデルをドリフト関数およびスコア関数に回帰させる。
• 神経的グラフィカルモデル（NGM）をODEドリフトに統合し、遺伝子-遺伝子相互作用ネットワークをパラメータライズし、L1正則化を用いてスパarsityを強制する。
• トレーニング時にODE統合を回避するシミュレーションフリーな目的関数を用いてモデルをトレーニングし、高次元データにおける効率的かつスケーラブルなトレーニングを可能にする。

実験結果

リサーチクエスチョン

• RQ1 シミュレーションフリーな目的関数は、任意の分布間の確率的ダイナミクスの学習において、スコアマッチィングとフローマッチィングを統合できるか？
• RQ2 エントロピー最適輸送を用いて、シュレーディンガー・ブリッジ問題の静的でシミュレーションフリーな代替手段を構築できるか？
• RQ3 [SF]2Mは、アンパイルドスナップショットデータから高次元の生物学的ダイナミクス（例：単細胞遺伝子発現）を正確にモデリングできるか？
• RQ4 [SF]2Mは、シミュレートされた単細胞データから既知の遺伝子調節ネットワークを回復でき、既存のベースラインを上回る性能を示すか？
• RQ5 ドリフトとスコアダイナミクスの連合モデリングは、個別に学習されたスコアまたはフロー・モデルと比較して、一般化性とロバスト性を向上させるか？

主な発見

• [SF]2Mは、トレーニング中にSDEシミュレーションを一切行わず、シミュレーションベースの手法よりも正確で効率的なシュレーディンガー・ブリッジ問題の解決策を達成した。
• 合成データにおいて、推論時に拡散スケジュールを変化させることで、ODEとSDEダイナミクスの間を滑らかに補間でき、確率的ダイナミクスのモデリングにおける柔軟性を示した。
• 高次元の単細胞データにおいて、[SF]2Mは、全遺伝子空間における細胞ダイナミクスの正確なモデリングを達成した最初の手法であり、シミュレートされたデータから既知の遺伝子調節ネットワークを回復した。
• 分岐的および三叉分岐の遺伝子調節ネットワークのシミュレーションデータにおいて、[SF]2MはAUC-ROCスコアがそれぞれ0.92および0.89を達成し、Spearman、Pearson、DREMI、Grangerベースラインを上回った。
• 単細胞データの最初の1000個の高可変性遺伝子において、予測分布と真値分布の間の1-Wasserstein距離が0.12にまで低下し、強い分布一致を示した。
• NGMの隣接行列にL1正則化を適用することで、真値のGRNsに近いスパースで解釈可能な遺伝子-遺伝子相互作用ネットワークが得られた。