Beyond Optimal Transport: Model-Aligned Coupling for Flow Matching

• URL: http://arxiv.org/abs/2505.23346v1
• Date: Thu, 29 May 2025 11:10:41 GMT
• Title: Beyond Optimal Transport: Model-Aligned Coupling for Flow Matching
• Authors: Yexiong Lin, Yu Yao, Tongliang Liu,
• Abstract summary: 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.<n>We propose Model- Coupling Coupling (MAC), an effective method that matches training couplings based on geometric distance.<n>Experiments show that MAC significantly improves generation quality and efficiency in few-step settings compared to existing methods.
• Score: 59.97254029720877
• License: http://creativecommons.org/licenses/by/4.0/
• 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

Related papers

• NAN: A Training-Free Solution to Coefficient Estimation in Model Merging [61.36020737229637]
  We show that the optimal merging weights should scale with the amount of task-specific information encoded in each model.<n>We propose NAN, a simple yet effective method that estimates model merging coefficients via the inverse of parameter norm.<n>NAN is training-free, plug-and-play, and applicable to a wide range of merging strategies.
  arXiv  Detail & Related papers  (2025-05-22T02:46:08Z)
• Flow Models for Unbounded and Geometry-Aware Distributional Reinforcement Learning [5.628557798049032]
  We introduce a new architecture for Distributional Reinforcement Learning (DistRL) that models return distributions using normalizing flows.<n>This approach enables flexible, unbounded support for return distributions.<n>We show that our approach outperforms PDF based models while remaining competitive with quantile based methods.
  arXiv  Detail & Related papers  (2025-05-07T10:49:53Z)
• Learning Straight Flows by Learning Curved Interpolants [19.42604535211923]
  Flow matching models typically use linear interpolants to define the forward/noise addition process.<n>This, together with the independent coupling between noise and target distributions, yields a vector field which is often non-straight.<n>We propose to learn flexible (potentially curved) interpolants in order to learn straight vector fields to enable faster generation.
  arXiv  Detail & Related papers  (2025-03-26T16:54:56Z)
• Not-So-Optimal Transport Flows for 3D Point Cloud Generation [58.164908756416615]
  Learning generative models of 3D point clouds is one of the fundamental problems in 3D generative learning.<n>In this paper, we analyze the recently proposed equivariant OT flows that learn permutation invariant generative models for point-based molecular data.<n>We show that our proposed model outperforms prior diffusion- and flow-based approaches on a wide range of unconditional generation and shape completion.
  arXiv  Detail & Related papers  (2025-02-18T02:37:34Z)
• Flow Map Matching [15.520853806024943]
  Flow map matching is an algorithm that learns the two-time flow map of an underlying ordinary differential equation. We show that flow map matching leads to high-quality samples with significantly reduced sampling cost compared to diffusion or interpolant methods.
  arXiv  Detail & Related papers  (2024-06-11T17:41:26Z)
• Preference Alignment with Flow Matching [23.042382086241364]
  Preference Flow Matching (PFM) is a new framework for preference-based reinforcement learning (PbRL) It streamlines the integration of preferences into an arbitrary class of pre-trained models. We provide theoretical insights that support our method's alignment with standard PbRL objectives.
  arXiv  Detail & Related papers  (2024-05-30T08:16:22Z)
• Optimal Flow Matching: Learning Straight Trajectories in Just One Step [89.37027530300617]
  We develop and theoretically justify the novel textbf Optimal Flow Matching (OFM) approach. It allows recovering the straight OT displacement for the quadratic transport in just one FM step. The main idea of our approach is the employment of vector field for FM which are parameterized by convex functions.
  arXiv  Detail & Related papers  (2024-03-19T19:44:54Z)
• Exploring Straighter Trajectories of Flow Matching with Diffusion Guidance [66.4153984834872]
  We propose Straighter trajectories of Flow Matching (StraightFM) It straightens trajectories with the coupling strategy guided by diffusion model from entire distribution level. It generates visually appealing images with a lower FID among diffusion and traditional flow matching methods.
  arXiv  Detail & Related papers  (2023-11-28T06:19:30Z)
• Manifold Interpolating Optimal-Transport Flows for Trajectory Inference [64.94020639760026]
  We present a method called Manifold Interpolating Optimal-Transport Flow (MIOFlow) MIOFlow learns, continuous population dynamics from static snapshot samples taken at sporadic timepoints. We evaluate our method on simulated data with bifurcations and merges, as well as scRNA-seq data from embryoid body differentiation, and acute myeloid leukemia treatment.
  arXiv  Detail & Related papers  (2022-06-29T22:19:03Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.