技術的概要：カウントデータ向けフローマッチング（count-FM）

問題定義

単一細胞 RNA シーケンシング（scRNA-seq）や神経スパイクトレイン解析などの分野で広く見られる高次元のカウントデータは、生成モデルにとって固有の課題を呈します。これらのデータタイプは、離散性、疎性、そして複雑な相関構造によって特徴づけられます。既存の深層生成アプローチはこの領域において困難に直面しています：

• カテゴリカル状態アプローチ： カウントをカテゴリカル変数（例えばトークン）として扱う手法は、出力次元が最大カウント範囲に比例してスケーリングする必要があります。これによりパラメータ効率が低下し、隣接するカウントが互いに関係のない状態として扱われ、カウント空間の固有の幾何学が無視されます。
• 連続変換： デquantization や潜在表現を通じてカウントを連続空間に写像するアプローチは、離散確率質量を連続密度に置き換えるため、潜在的な離散構造を曖昧にする可能性があります。
• 既存カウントモデルの限界： 一部のカウント固有のジャンプモデル（例：Poisson-JUMP）は存在しますが、それらは任意のカウント分布間の輸送を設計していないか、現代のフローベース手法のような効率的な学習および生成特性を欠いていることが多いです。

本論文は、カウント空間で直接動作し、データの離散幾何学を尊重するとともに、カテゴリカル状態モデルのパラメータオーバーヘッドなしに任意のソース分布とターゲット分布間の効率的な輸送を可能にする生成フレームワークの必要性に対処します。

手法：count-FM

著者らは、局所単位ジャンプを伴う連続時間出生・死亡過程に基づき、カウントデータに特化したフローマッチングフレームワークであるcount-FMを提案します。

1. 連続時間出生・死亡過程

モデルは、ソース分布 $p_0$ とターゲット分布 $p_1$ の間の遷移を、カウントベクトルの空間 $\mathbb{N}_0^d$ 上の連続時間マルコフジャンプ過程（CTMC）として定義します。この過程は、各座標 $i$ における局所単位出生（$+1$）と死亡（$-1$）を通じて進化します：

• 出生率： $\lambda_{t,i}(x)$
• 死亡率： $\mu_{t,i}(x)$
• ダイナミクス： 微小区間 $h$ において、状態 $x$ から $x + e_i$ へ遷移する確率は $h \lambda_{t,i}(x)$ であり、$x - e_i$ へ遷移する確率は $h \mu_{t,i}(x)$ です。カウントがゼロのときは死亡は許可されません。

このパラメータ化は非常に効率的です。モデルは、最大カウントに比例してスケーリングするすべての可能なカウントレベルにわたる完全なカテゴリカル分布を予測する代わりに、カウント範囲に関わらず、次元あたり 2 つのレート（$2d$ 出力）のみを予測します。

2. 条件付き二項ブリッジと学習

モデルを学習するために、著者らは固定された端点 $(x_0, x_1)$ の間の条件付きブリッジを定義します。各座標 $i$ について、ブリッジは二項分布を用いて構築されます：
$$X_t^{(i)} = x_0^{(i)} + \text{sgn}(x_1^{(i)} - x_0^{(i)}) B_t^{(i)}, \quad B_t^{(i)} \sim \text{Binomial}(|x_1^{(i)} - x_0^{(i)}|, t)$$
このブリッジは、条件付き平均が $x_0$ から $x_1$ へ線形に移動しつつ、整数値の中間状態を維持することを保証します。

この条件付きブリッジに対して質量保存（コルモゴロフ前方方程式）を課すことで、著者らは条件付き出生率と死亡率の閉形式式を導出します：
$$\lambda_{t,i}(x | x_0, x_1) = \frac{(x_1^{(i)} - x)^+}{1-t}, \quad \mu_{t,i}(x | x_0, x_1) = \frac{(x - x_1^{(i)})^+}{1-t}$$
ここで $(\cdot)^+$ は正の整流を表します。

3. 学習目的関数

モデルは、条件付きブリッジ過程とモデルが学習した過程との間の経路空間におけるカルバック・ライブラー（KL）ダイバージェンスを最小化することで学習されます。これは局所レートマッチング目的関数に分解されます：
$$\mathcal{L}_{train}(\theta) = \mathbb{E} \left[ \sum_{i=1}^d \ell(\lambda_{t,i}, \lambda_{\theta,i}) + \sum_{i=1}^d \ell(\mu_{t,i}, \mu_{\theta,i}) \right]$$
ここで $\ell(u, v) = v - u \log v$ は一般化された KL 項です。これにより、識別可能な条件付きレートを学習することで、シミュレーションなしで周辺遷移レートを学習することが可能になります。

4. サンプリングとカップリング

• サンプリング： サンプルは、第一-order 局所ジャンプ離散化を用いて、$t=\epsilon_t$ から $t=1-\epsilon_t$ まで学習された出生・死亡過程を前方にシミュレーションすることで生成されます。
• 端点カップリング： 本論文は、ソースとターゲットのサンプルをペアリングするための 2 つのカップリング戦略を探求します：
  • 独立カップリング： ランダムなペアリング（$p_0(x_0)p_1(x_1)$）。
  • 最適輸送（OT）カップリング： 対称ポアソンコストを用いたミニバッチ OT ペアリング。これは、より低い曲率を持つ遷移を誘発する傾向があり、解釈可能性とサンプリング効率を向上させます。

5. 条件付き生成

条件付きタスクでは、モデルは共変量 $y$ を組み込み、Classifier-Free Guidance (CFG) を利用します。学習中は、条件付け変数がランダムにドロップされます。サンプリング時には、条件付きレートと無条件レートがガイダンススケール $w$ を通じて結合され、条件の強さを制御します。

主要な貢献

1. count-FM フレームワーク： 局所出生・死亡ダイナミクスに基づくカウントデータ向けフローマッチングフレームワーク。中間状態がカウント値のまま保持されるカウント空間内での直接輸送を可能にします。
2. 実行可能な学習スキーム： 条件付き二項ブリッジを介して周辺遷移レートを学習する手法。フローマッチングの効率性を維持しつつ、カテゴリカル状態モデルの複雑さを回避します。
3. パラメータ効率性： 完全なカテゴリカル分布ではなく、局所出生・死亡レート（$2d$ 出力）のみを予測する定式化。これにより、大きなカウント範囲を持つ高次元データに適しています。

実験結果

1. シミュレーション（2 次元ガンマ・ポアソン混合）

• 性能： count-FM および count-FM-OT は、D3PM、discrete-FM、tauLDR、SEDD、Poisson-JUMP などのベースラインと比較して、最良のサンプル品質（最低の 2-Wasserstein 距離および MMD$^2$）を達成しました。
• 効率性： count-FM は、カテゴリカル状態ベースライン（14,000〜67,000 以上）と比較して、はるかに少ない学習可能パラメータ（2,372）でこれらの結果を達成しました。
• 幾何学： 中間状態はカウント空間内で滑らかに進化しました。OT カップリング版は、視覚的により直線的な遷移経路をたどり、少ない関数評価（NFEs）でターゲット品質に到達しました。

2. 単一細胞 RNA-seq（歯状回）

• 無条件生成： 2,930 個の細胞と 13,913 個の遺伝子からなるデータセットにおいて、count-FM はサンプル品質（W2 および MMD$^2$）において潜在変数ベースライン（scVI）および潜在拡散モデル（scDiffusion、scLDM）を上回り、かつ競合他社（約 29M〜39M）と比較して少ないパラメータ（約 9.8M）を使用しました。
• 輸送（P12 から P35）： このモデルは、出生後 12 日目から 35 日目までの発達的輸送を成功裡にモデル化しました。系統制限付き OT カップリングを使用することで、輸送経路は既知の発達構造（例、神経芽細胞から顆粒細胞への進行）を尊重し、カウント空間内に留まる解釈可能な中間分布を生成しました。

3. 神経スパイクトレイン（海馬/内側嗅覚野）

• 条件付き生成： このモデルは、線形化された位置と走行方向に条件付けられたスパイクカウントの生成に適用されました。
• 比較： ガイダンススケール $w=1$ の count-FM は、決定論的 MLP 回帰モデルおよびポアソン MLP ベースラインを上回りました。
• 相関構造： 集団相関を過小評価したポアソン MLP と異なり、count-FM はニューロン間の依存構造をよりよく保持しました。
• ガイダンスのトレードオフ： $w=1$ は空間チューニングと分布較正の間の最良のバランスを提供しました。$w=2$ はプレイスフィールドを鋭くしましたが、振幅較正を歪めました。

意義と主張

本論文は、count-FM がカウント空間内での直接生成および輸送を可能にすることで、高次元カウントデータのモデル化に効果的なフレームワークを提供すると主張しています。その主な意義は以下の点にあります：

• 幾何学的忠実性： 局所単位ジャンプを介した輸送をモデル化することで、この手法はカウントデータの離散幾何学を尊重し、連続近似の曖昧化効果を回避します。
• スケーラビリティ： パラメータ効率性により、カテゴリカル状態手法が計算的に実行不可能となる高次元アプリケーション（例、数千の遺伝子）において実用可能です。
• 解釈可能性： 輸送経路の中間状態は有効なカウントサンプルであり、発生生物学などの応用において意味のある遷移分布の検査を可能にします。
• 汎用性： このフレームワークは、無条件生成、任意の分布間の輸送、制御可能なガイダンスを伴う条件付き生成をサポートします。

著者らは限界を認め、条件付き二項ブリッジは比較的単純であり、すべての妥当な中間経路を探索するわけではないこと、また現在の第一-order 離散化はより正確なシミュレーションスキームによって改善され得ることを指摘しています。しかし、シミュレーションおよび生物医学データセット全体にわたる実証結果は、count-FM が競合するアプローチよりも高いモデル効率で高品質なサンプルを達成することを示しています。