Beyond Optimal Transport: Model-Aligned Coupling for Flow Matching

TL;DR

This paper introduces Model-Aligned Coupling (MAC), a novel approach that improves flow matching by aligning couplings with the model's preferred transport directions, resulting in better quality and efficiency.

Contribution

MAC is the first method to incorporate model alignment into coupling construction for flow matching, enhancing training effectiveness and sample quality.

Findings

MAC outperforms existing methods in few-step generation tasks.

MAC achieves higher sample quality with fewer integration steps.

Extensive experiments validate the efficiency and effectiveness of MAC.

Abstract

Flow Matching (FM) is an effective framework for training a model to learn a vector field that transports samples from a source distribution to a target distribution. To train the model, early FM methods use random couplings, which often result in crossing paths and lead the model to learn non-straight trajectories that require many integration steps to generate high-quality samples. To address this, recent methods adopt Optimal Transport (OT) to construct couplings by minimizing geometric distances, which helps reduce path crossings. However, we observe that such geometry-based couplings do not necessarily align with the model's preferred trajectories, making it difficult to learn the vector field induced by these couplings, which prevents the model from learning straight trajectories. Motivated by this, we propose Model-Aligned Coupling (MAC), an effective method that matches training couplings based not only on geometric distance but also on alignment with the model's preferred transport directions based on its prediction error. To avoid the time-costly match process, MAC proposes to select the top-$k$ fraction of couplings with the lowest error for training. Extensive experiments show that MAC significantly improves generation quality and efficiency in few-step settings compared to existing methods. Project page: https://yexionglin.github.io/mac