COT-FM: クラスタ単位最適輸送フローマッチング

要旨

本論文では、Flow Matching（FM）における確率経路を再構築し、より高速かつ信頼性の高い生成を実現する汎用フレームワークCOT-FMを提案する。FMモデルは、ランダムまたはバッチ単位のカップリングにより曲線的な軌道を生成することが多く、これが離散化誤差を増大させサンプル品質を低下させる要因となっている。COT-FMは、目標サンプルをクラスタリングし、各クラスタに事前学習済みFMモデルを逆伝播させて得た専用のソース分布を割り当てることでこの問題を解決する。この分割統治戦略により、モデル構造を変更することなく、より正確な局所輸送と著しく直線的なベクトル場を実現する。プラグアンドプレイ方式のCOT-FMは、2次元データセット、画像生成ベンチマーク、ロボット把持タスクにおいて、一貫してサンプリングを加速し生成品質を向上させる。

English

We introduce COT-FM, a general framework that reshapes the probability path in Flow Matching (FM) to achieve faster and more reliable generation. FM models often produce curved trajectories due to random or batchwise couplings, which increase discretization error and reduce sample quality. COT-FM fixes this by clustering target samples and assigning each cluster a dedicated source distribution obtained by reversing pretrained FM models. This divide-and-conquer strategy yields more accurate local transport and significantly straighter vector fields, all without changing the model architecture. As a plug-and-play approach, COT-FM consistently accelerates sampling and improves generation quality across 2D datasets, image generation benchmarks, and robotic manipulation tasks.

COT-FM: クラスタ単位最適輸送フローマッチング

COT-FM: Cluster-wise Optimal Transport Flow Matching

要旨

Support