[論文レビュー] Improving and generalizing flow-based generative models with minibatch optimal transport

Alexander Tong, Kilian Fatras+6 more|arXiv (Cornell University)|Feb 1, 2023

Generative Adversarial Networks and Image Synthesis被引用数 27

ひとこと要約

一般化された conditional flow matching (CFM) for continuous normalizing flows と a minibatch optimal transport variant (OT-CFM) が、dynamic OT および Schrödinger bridges を近似し、シミュレーション不要で、さまざまなデータタスクに対してより速い学習と推論を実現します。

ABSTRACT

Continuous normalizing flows (CNFs) are an attractive generative modeling technique, but they have been held back by limitations in their simulation-based maximum likelihood training. We introduce the generalized conditional flow matching (CFM) technique, a family of simulation-free training objectives for CNFs. CFM features a stable regression objective like that used to train the stochastic flow in diffusion models but enjoys the efficient inference of deterministic flow models. In contrast to both diffusion models and prior CNF training algorithms, CFM does not require the source distribution to be Gaussian or require evaluation of its density. A variant of our objective is optimal transport CFM (OT-CFM), which creates simpler flows that are more stable to train and lead to faster inference, as evaluated in our experiments. Furthermore, we show that when the true OT plan is available, our OT-CFM method approximates dynamic OT. Training CNFs with CFM improves results on a variety of conditional and unconditional generation tasks, such as inferring single cell dynamics, unsupervised image translation, and Schrödinger bridge inference.

研究の動機と目的

• Develop a simulation-free training objective for continuous normalizing flows that generalizes prior flow matching methods.
• Introduce an optimal-transport-informed variant (OT-CFM) to approximate dynamic OT and yield straighter, faster flows.
• Show that OT-CFM improves training efficiency and inference speed across conditional and unconditional generation tasks.
• Provide practical training techniques and open-source software to unify flow-based model training.

提案手法

• Propose generalized conditional flow matching (CFM) that regresses a neural vector field vθ(t, x) to the marginal transport field ut(x) using samples from conditional paths p t(x|z) and conditioning z.
• Define ut(x) as the marginal field derived from conditioned paths via ut(x) = Ez [ ut(x|z) p t(x|z) / p t(x) ].
• Show equivalence LFM ≈ LCFM under broad conditions, enabling regression without density evaluation of the source.
• Introduce OT-CFM where z is drawn from a minibatch optimal transport plan π(x0, x1), linking to dynamic OT in a simulation-free way.
• Extend to Schrödinger bridge CFM (SB-CFM) using entropy-regularized OT, matching SB probability flow with a Brownian reference process.
• Provide minibatch OT approximation to enable training on large datasets, with theory connecting minibatch OT to dynamic OT when batch size grows.

実験結果

リサーチクエスチョン

• RQ1 Can a generalized CFM framework accommodate arbitrary source distributions and transport maps beyond Gaussian assumptions?
• RQ2 Does conditioning on minibatch OT maps (OT-CFM) effectively approximate dynamic OT and Schrödinger bridge flows in a simulation-free manner?
• RQ3 Do OT-CFM and SB-CFM provide faster training and inference while maintaining or improving model fidelity across diverse tasks (single-cell dynamics, image generation, translation)?
• RQ4 What practical training strategies and implementations (e.g., minibatch OT, solver choices) yield robust, reproducible flow-based modeling results?

主な発見

• OT-CFM yields faster training and inference compared to baseline flow matching methods, and requires fewer function evaluations for similar sample quality.
• Minibatch OT maps enable learning dynamic OT without exact full-batch OT, with performance close to exact dynamic OT and improved efficiency.
• SB-CFM can reproduce Schrödinger bridge flows by using entropy-regularized OT, connecting to SB probability flows with a Brownian reference.
• OT-CFM achieves better or comparable 2-Wasserstein distances and NPE on low-dimensional benchmarks and shows improved FID/Evaluations on CIFAR-10 under revised training settings.
• In single-cell interpolation tasks, OT-CFM outperforms several baselines on average across datasets in left-out timepoint error metrics.