Posterior Augmented Flow Matching
Abstractの概要
本論文は、標準的なフローマッチング（FM）が各中間潜在状態に対して1つのターゲット軌道のみを対応させるため、疎な教師信号を提供し、高次元生成タスクにおいて高分散の学習信号やフロー崩壊を引き起こす可能性があると主張している。Posterior-Augmented Flow Matching（PAFM）を導入し、単一ターゲットの教師信号を、与えられた中間状態と条件に対する複数の妥当なターゲット完了の期待値で置き換える。この手法は、扱いにくい事後分布を条件付きパス尤度と条件尤度に分解し、自己正規化重要度重み付けを用いて学習中に候補ターゲットを集約する。著者らはPAFMがFM目的関数の不偏推定量であり、かつ勾配の分散を低減することを証明し、クラス条件付きImageNet-1Kおよびテキスト画像生成CC12Mベンチマークにおいて、複数のアーキテクチャとモデルスケールで評価を行っている。
新規性
主な新規性は、フローマッチングを再定式化し、各中間点を単一の終点ではなく、扱いにくい事後分布を条件付きパス尤度と条件尤度に分解することで導出された事後重み付き有効ターゲットの混合で教師付けする点にある。また、実用的な自己正規化重要度サンプリングの実装を提案し、最近傍探索、ランダム空間拡張、VAEモーメント再サンプリングを含む候補ターゲットの構築方法の柔軟性を実証している。
成果
ImageNet-1KおよびCC12Mにおいて、PAFMは標準FMに対して一貫した改善を示し、SiT-B/2（K=16）でFIDが27.57から24.88へ、SiT-XL/2（K=16）で11.14から9.85へ、CC12MのMMDiTで10.37から9.45への改善が報告されている。また、ImageNet実験において約4倍低いミニバッチ勾配分散が測定され、K=32の最近傍変種でスループット低下はわずか6.6%、メモリ増加は0.4%にとどまっている。
論文の注目点
- PAFMは、各中間潜在状態に対して一対一の教師信号を、条件付きパス尤度と条件尤度から導出された重みによる自己正規化重要度サンプリングを用いた、複数の妥当な継続軌道の事後重み付き期待値で置き換えることで、フローマッチングを一般化している。
- 理論的解析により、PAFMが元のFM目的関数の不偏推定量であり、Kishの有効サンプルサイズに関連する因子で勾配分散を低減することが証明されており、ImageNetにおいて約4倍低いミニバッチ勾配分散として経験的に確認されている。
- ImageNet-1KおよびCC12Mでの経験的評価により、SiT-B/2、SiT-XL/2、MMDiTモデル全体で一貫したFID改善（最大3.4 FID50K）が示され、代替候補選択戦略（拡張およびVAE再サンプリング）も無視できる計算オーバーヘッドで標準FMを上回っている。