The Coupling Within: Flow Matching via Distilled Normalizing Flows

TL;DR

This paper introduces Normalized Flow Matching (NFM), a novel method that distills couplings from pretrained normalizing flow models to enhance training and performance of flow-based generators.

Contribution

The paper proposes NFM, a new approach that leverages distilled couplings from pretrained NF models to improve flow training and generation quality.

Findings

NFM outperforms models trained with independent couplings.

NFM surpasses OT-based coupling methods in performance.

Distilled couplings improve both training efficiency and sample quality.

Abstract

Flow models have rapidly become the go-to method for training and deploying large-scale generators, owing their success to inference-time flexibility via adjustable integration steps. A crucial ingredient in flow training is the choice of coupling measure for sampling noise/data pairs that define the flow matching (FM) regression loss. While FM training defaults usually to independent coupling, recent works show that adaptive couplings informed by noise/data distributions (e.g., via optimal transport, OT) improve both model training and inference. We radicalize this insight by shifting the paradigm: rather than computing adaptive couplings directly, we use distilled couplings from a different, pretrained model capable of placing noise and data spaces in bijection -- a property intrinsic to normalizing flows (NF) through their maximum likelihood and invertibility requirements. Leveraging recent advances in NF image generation via auto-regressive (AR) blocks, we propose Normalized Flow Matching (NFM), a new method that distills the quasi-deterministic coupling of pretrained NF models to train student flow models. These students achieve the best of both worlds: significantly outperforming flow models trained with independent or even OT couplings, while also improving on the teacher AR-NF model.