Active Timepoint Selection for Learning Measure-Valued Trajectories

技術要約：学習値としての測度軌跡のための能動的なタイムポイント選択

問題定義
本論文は、疎で破壊的なスナップショットから連続的な確率パス $\{\mu_t\}_{t \in [0,1]}$ を推論するという課題に取り組んでいる。この問題は、シングルセル生物学などの領域において極めて重要である。なぜなら、高精度なデータ取得（例：シーケンシング）はコストが高く、かつ破壊的であるため、密な時間的サンプリングが困難だからである。目標は、離散的な時点における観測された経験的測度の有限集合から、基礎となる連続的な軌跡を復元する「分布補間（distributional interpolation）」を行うことである。

著者らは、この設定に能動学習（active learning）を適用する際の3つの主要な課題を特定している：

1. 非ユークリッド幾何学： 確率測度は、非線形なワッサースタイン空間 $\mathcal{P}_2(\mathcal{X})$ に存在する。標準的な能動学習手法は、ユークリッド距離（例：標準的なガウス過程）に依存しているため、これらは失敗する。なぜなら、測度の線形結合は有効な測度を生成せず、ナイーブな補間は質量を輸送するのではなく分割してしまうからである。
2. エピステミック不確実性の欠如： 現在の分布補間手法（例：フローマッチング、シュレディンガー・ブリッジ）は、通常、決定論的である。これらは単一の適合されたパスを返すだけであり、能動学習ポリシーを導くために不可欠なエピステミック不確実性（認識論的不確実性）を定量化するメカニズムを欠いている。
3. 非定常性： 測度値のダイナミクスは、変化する速度（例：恒常性と、急速な分岐イベントが交互に現れる現象）を示すことが多い。一様に配置された取得時間は、これらの過渡的で高速度な領域を捉えるには最適ではない。

手法
著者らは、扱いやすい確率的サロゲート（代理モデル）を構築することにより、測度の空間へと能動的実験を拡張するフレームワークを提案している。この手法は、以下の4つの主要な段階で構成される：

1. 線形化された最適輸送（LOT）による線形化：
   非ユークリッド幾何学を扱うために、フレームワークは観測された経験的測度 $\hat{\mu}_t$ を、固定された参照測度 $\sigma$（観測されたスナップショットのワッサースタイン・バリセンターとして選択される）で定義される共通の接空間 $T_\sigma \mathcal{P}_2(\mathcal{X})$ へと写像する。LOT対数写像を用いることで、各測度は変位ベクトル場 $v_\mu(x) = T_{\sigma \to \mu}(x) - x$ として表現される（ここで $T$ は最適輸送写像である）。これにより、ワッサースタイン幾何学が線形化され、測度をヒルベルト空間内のベクトルとして扱うことが可能になる。

2. 低次元表現：
   高次元の変位場は、低次元の潜在空間へと圧縮される。著者らは、再重み付けされた変位場を平坦化し、主成分分析（PCA）を適用することで、各時点 $t$ に対する $K$ 次元の表現 $c_t$ を得る。

3. ガウス過程（GP）を用いた確率的モデリング：
   潜在係数 $c_t$ の時間的進化に対して、マルチアウトプット・ガウス過程（MOGP）が配置される。モデルは行列値カーネルを持つゼロ平均事前分布を仮定している。GP事後分布は、予測平均と、決定的ながら、エピステミック不確実性を定量化する共分散行列を提供する。

• 再構成： 任意の時刻 $t$ における予測測度 $\hat{\mu}_t$ を復元するために、潜在的な平均は逆PCA投影と逆LOT写像（参照ランドマークに変位を加える）を通じてマップされる。

4. タイムワーピングによる非定常性の処理：
   生物学的プロセスの非定常な性質に対処するため、著者らは固有の時間パラメータ化を導入している。ワーピング関数 $\Phi(t)$ は、物理的な時間を、ワッサースタイン空間における軌跡の累積弧長を表す固有の時間 $\tau$ へと写像する。定常カーネルがこのワーピングされたドメイン内で適用されることで、実質的に、分布の変化の局所的な速度に適応する非定常カーネルが物理的な時間において誘導される。

5. 獲得ポリシー：
   フレームワークは、次の測定時刻 $t^*$ を選択するための獲得関数を採用している。このポリシーは、情報利得を最大化するように、特に固有時間ドメインにおける高いエピステミック不確実性を標的として、反復的に $t^*$ を選択する。2つの獲得戦略、すなわち、点ごとの不確実性（事後共分散のトレースを最大化する）と、期待積分リスク減少が提案されている。

主な貢献

• 概念的： 本論文は、確率パスの学習のための能動学習を定式化し、能動的実験をユークリッド空間から確率測度の空間へと拡張した。
• 技術的： 線形化された最適輸送とマルチアウトプット・ガウス過程を組み合わせることで、ワッサースタイン空間における扱いやすい確率的サロゲートを構築した。この組み合わせにより、従来の分布補間手法には欠けていた、確率パスに対するエピステミック不確実性の定量化が可能となった。
• 実証的： 著者らは、分岐イベントを模倣した合成データセットおよび実世界のシングルセル・トランスクリプトミクス・データセットの両方において、提案する獲得戦略が不確実性を考慮しないベースライン（一様サンプリングおよびランダムサンプリング）を上回ることを示した。

結果

• 合成データ： 非定常なダイナミクスと分岐イベントを持つ合成データセットにおいて、能動的戦略は、特に低〜中程度の予算において、一様およびランダムなベースラインを一貫して上回った。本手法は、高速度領域（分岐イベント）に測定を集中させることに成功し、平均ワッサースタイン誤差および速度加重誤差を減少させた。感度分析により、本手法は分岐イベントが高度に局在化している場合に最も効果的であることが示された。
• 実世界データ： マウス線維芽細胞からiPSCへのシングルセル・リプログラミング・データセットに適用したところ、能動的戦略は低〜中程度の予算（$B \le 12$）において、最も低い再構成誤差を達成した。この優位性は、予算が増加しサンプリングが密になった場合に減退し、制約のあるリソース環境における本手法の有用性を裏付けた。
• アブレーション研究： 固有時間ワーピングを削除すると性能が低下し、非定常モデリングの必要性が検証された。潜在次元を削減（$K=2$）すると大幅な性能低下が生じ、複雑な変動を捉えるための十分な次元数の必要性が浮き彫りになった。参照測度を（バリセンターではなく）固定すると誤差が増大し、線形化誤差を最小化することの重要性が強調された。

意義と主張
本論文は、その主要な貢献が、非ユークリッド幾何学と不確実性定量化の欠如という二重の課題を解決することにより、測度の空間における能動学習を可能にしたことにあると主張している。LOTとGPを活用することで、本フレームワークは、再構成された確率パスの不確実性を最小化する最適な測定時刻を選択するための実用的なメカニズムを提供する。

著者らは限界についても述べている。本アプローチは接空間近似に依存しており、大きな分布シフトが発生した場合に誤差が生じる可能性がある。また、計算コストは最適輸送のサブルーチンによって支配されるため、意図された範囲（スナップショットあたり最大 $\sim 10^5$ サンプル）には実用的であるが、さらなるスケーラブルな近似なしには、百万点規模のデータセットには困難である可能性がある。今後の課題として、リーマン幾何学のツール（マルチチャート）や多次元入力の探索が示唆されている。