One-Sided Quantile Coupling for Flow Matching

技術要約：Flow Matchingのための片側分位量結合（One-Sided Quantile Coupling）

問題提起
Flow Matching (FM) は、単純なソース分布（通常は標準ガウス分布）からターゲットとなるデータ分布へと確率パスを回帰することにより、連続時間生成モデルを学習する手法である。この際、結合（カップリング）の質、すなわちソースとターゲットのサンプルのペアリングが、学習軌道の幾何学的形状に決定的な影響を与える。標準的なFMは独立したランダムなペアリングを使用するが、これはしばしば長く、かつ矛盾したパスをもたらし、ターゲット速度場の分散を増大させ、学習効率を阻害する。

Optimal Transport (OT) に着想を得た構造化された結合は、これらのパスを直線化し、生成品質を向上させる可能性があるが、既存の手法（例：OT-CFM）はミニバッチ内の割り当て問題を解くことに依存している。これらのアプローチは、ペアごとのコスト行列を構築し、割り当て問題（例：ハンガリー法やSinkhorn反復）を解く必要があり、バッチサイズに対して少なくとも二次的（正確な解法では三次的）な計算コストを要する。これはスケーラビリティを制限し、不安定性を引き起こす可能性がある。本論文は、バッチごとの輸送問題を解くことなく、構造化された結合による幾何学的な利点を得られるかという問いに取り組んでいる。

手法：分位量結合フローマッチング (QC-FM)
著者らは、軽量な片側結合スキームである Quantile Coupling Flow Matching (QC-FM) を提案する。データとノイズの両方のバッチをサンプリングしてそれらを再ペアリングする従来のメソッドとは異なり、QC-FMはデータバッチのみをサンプリングし、対応するソースサンプルを直接構築する。

1. 片側構築:

   • データバッチ $\mathcal{B}_x$ が与えられると、手法は行列 $U$ によって定義される $k$ 個のランダムな直交方向（スライス）上にデータ点を投影する。
   • 各方向について、バッチ内でのデータ点の順位（ランク）を決定する。これらの経験的なランクは、逆標準正規累積分布関数を用いてガウス分位量へとマッピングされる。
   • ソース潜在コード $e^{(i)}$ は、これらの投影された分位量制約を満たすように構築される。具体的には、ソースの投影が、データランクから導出されたガウス分位量と一致するように設計される。
   • 残りの次元（スライス部分空間の直交補空間）は、条件付きガウス分布からサンプリングすることで埋められる。これにより、制約のない方向においてソースの確率性が維持される。

2. ガウス残差ハイブリッド拡張:
   完全なQC結合（すべてのソースポイントがランクを通じて構築される場合）は、$k \ge 2$ の場合にスライスコード間の依存関係により標準的な独立同一分布（i.i.d.）ガウス事前分布から逸脱するため、著者らは事前分布を保持しつつQCの利点を維持するための2つのハイブリッド戦略を導入している。

   • QC-FM-Mixture: バッチの「アンカー」サブセット（例：20%）に対してQC構築を適用し、残りの「残差」ソースサンプルを、残りのデータサンプルに対して独立したランダムペアリング（Baselineアプローチ）を用いてペアリングする。これにより、構造化された結合と非構造化された結合の混合体を作成する。
   • QC-FM-Adjacency: アンカー・サブセットにQCを適用するが、残りのソースサンプルはアンカーによって誘導される局所幾何学に基づいてペアリングされる。残りのデータ点は、（構造化された部分空間における）最も近いアンカーに割り当てられ、残りのソースサンプルはこれらのグループ内でペアリングされる。これにより、ターゲットを複製することなく、構造化されたバイアスを伝播させる。

主な貢献

• QC-FM アルゴリズム: ランダムなスライスに沿った単調な一次元輸送を課す、新しいバッチローカルな構造化結合である。これはペアごとのコスト行列や割り当てソルバーを不要にし、支配的な計算コストを $O(B^2)$ や $O(B^3)$ から $O(kB \log B)$ （ソート）へと削減する。
• 理論的分析:
  • 分散減少: 固定されたフレームに対して、QC結合は選択されたすべてのスライスにおいて既約な回帰分散を排除し、理想的なフローがその方向において正確に直線となることを証明している。
  • 事前分布の保持: $k=1$ の場合、ソースの周辺分布は正確にガウス分布であることを定量化している。$k \ge 2$ の場合、ガウス事前分布からの逸脱はスライスコードのコピュラに限定され、この逸脱による輸送コストは有界である。
• ハイブリッド学習スキーム: MixtureおよびAdjacency変体の導入により、幾何学的バイアスの注入と、Baselineからの明示的なガウス信号の保持とのバランスを実現している。

実験結果
手法は、訓練予算とアーキテクチャを一致させた条件下で、CIFAR-10、CelebA-64、FFHQ-64、および ImageNet-64 で評価された。

• 性能: QC-FM-Mixtureは、4つのデータセットすべてにおいて、独立結合を行うBaselineおよび正確なハンガリー法を用いたOT-CFMを一貫して上回った。Baselineに対する改善幅は、FID-50KにおいてCIFAR-10の6.3%からFFHQの12.9%であった。
• 効率性: 計算コストの観点では、QC-FMはOTベースの手法よりも大幅に高速である。バッチサイズ2048において、QC-FM-Mixtureはハンガリー法によるOTソルバーよりも800倍以上速く、Sinkhornソルバーよりも35倍速かった。
• トイ・ダイアグノスティクス: 合成データセット（Checkerboardおよび8-Gaussians）において、QC-FMの各バリアントは、Baselineと比較してほぼ直線的な学習軌道（経路長比 $\approx 1$）と、著しく低いターゲット速度分散を生み出し、結合の幾何学的利点を裏付けた。

意義と主張
本論文は、投影されたランク構造を保持することが、ミニバッチ輸送問題を解くことなく、Flow Matchingの結合に有用な幾何学的バイアスを注入するための、シンプルでスケーラブルかつ効果的な方法であると主張している。

• スケーラビリティ: QC-FMは、非構造化なBaselineと、計算コストの高いバッチワイズ輸送手法との間の、魅力的な中間領域を提供している。これは計算量的に軽量であり、投影数（$k$）によって制御可能である。
• 実用性: 本手法は推論のための学習済みコンポーネントを必要とせず、生成は依然として標準的なガウス分布から開始される。結合は訓練中にオンザフライで構築される。
• 限界: 著者らは、QC-FMは構造化結合の「バッチローカルな代用物」であり、グローバルなOptimal Transportの近似ではないことを謙虚に述べている。分位量はミニバッチ内で推定されており、理論はグローバルな軌道の直線性や特定のFID改善を保証するのではなく、スライスごとの分散を制御するものである。スライス数（$k$）およびアンカー比率（$p$）の選択は、現在はヒューリスティックおよびデータセット固有のチューニングに依存している。

結論として、本研究は、ランク保存を通じてノイズとデータのペアリングを構造化することにより、片側分位量結合がFlow Matchingにおけるサンプル品質と学習効率を向上させ得ることを示しており、計算量的に困難なOptimal Transport手法に代わる実用的な選択肢を提示している。