Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Transport Based Mean Flows for Generative Modeling

Elaheh Akbari, Ping He, Ahmadreza Moradipari, Yikun Bai, Soheil Kolouri

TL;DR

OT-Mean Flow addresses slow inference in flow-based generative modeling by learning a mean velocity aligned with an optimal transport coupling between the source and target distributions, i.e., the OT objective $OT(\mathbf{p},\mathbf{q})$. The method unifies conditional flow matching, minibatch OT flow matching, and mean-flow training under a single objective, enabling one-step generation. It introduces OT acceleration strategies such as Sinkhorn OT and Linear OT to scale training, and demonstrates improved inference accuracy and speed on low-dimensional controls, image generation (MNIST), and 3D point clouds and unpaired image-to-image translation. The results indicate that geometry-aware transport couplings yield straighter, more faithful trajectories than vanilla mean-flow.

Abstract

Flow-matching generative models have emerged as a powerful paradigm for continuous data generation, achieving state-of-the-art results across domains such as images, 3D shapes, and point clouds. Despite their success, these models suffer from slow inference due to the requirement of numerous sequential sampling steps. Recent work has sought to accelerate inference by reducing the number of sampling steps. In particular, Mean Flows offer a one-step generation approach that delivers substantial speedups while retaining strong generative performance. Yet, in many continuous domains, Mean Flows fail to faithfully approximate the behavior of the original multi-step flow-matching process. In this work, we address this limitation by incorporating optimal transport-based sampling strategies into the Mean Flow framework, enabling one-step generators that better preserve the fidelity and diversity of the original multi-step flow process. Experiments on controlled low-dimensional settings and on high-dimensional tasks such as image generation, image-to-image translation, and point cloud generation demonstrate that our approach achieves superior inference accuracy in one-step generative modeling.

Transport Based Mean Flows for Generative Modeling

TL;DR

OT-Mean Flow addresses slow inference in flow-based generative modeling by learning a mean velocity aligned with an optimal transport coupling between the source and target distributions, i.e., the OT objective . The method unifies conditional flow matching, minibatch OT flow matching, and mean-flow training under a single objective, enabling one-step generation. It introduces OT acceleration strategies such as Sinkhorn OT and Linear OT to scale training, and demonstrates improved inference accuracy and speed on low-dimensional controls, image generation (MNIST), and 3D point clouds and unpaired image-to-image translation. The results indicate that geometry-aware transport couplings yield straighter, more faithful trajectories than vanilla mean-flow.

Abstract

Flow-matching generative models have emerged as a powerful paradigm for continuous data generation, achieving state-of-the-art results across domains such as images, 3D shapes, and point clouds. Despite their success, these models suffer from slow inference due to the requirement of numerous sequential sampling steps. Recent work has sought to accelerate inference by reducing the number of sampling steps. In particular, Mean Flows offer a one-step generation approach that delivers substantial speedups while retaining strong generative performance. Yet, in many continuous domains, Mean Flows fail to faithfully approximate the behavior of the original multi-step flow-matching process. In this work, we address this limitation by incorporating optimal transport-based sampling strategies into the Mean Flow framework, enabling one-step generators that better preserve the fidelity and diversity of the original multi-step flow process. Experiments on controlled low-dimensional settings and on high-dimensional tasks such as image generation, image-to-image translation, and point cloud generation demonstrate that our approach achieves superior inference accuracy in one-step generative modeling.

Paper Structure

This paper contains 43 sections, 5 theorems, 88 equations, 8 figures, 5 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background and Related Work
  3. Classical Flow Matching and Conditional Flow Matching
  4. Unconditional Flow Matching.
  5. Optimal Transport and Related Flow Matching Models
  6. Optimal Transport.
  7. Mini-batch Optimal Transport Flow Matching.
  8. Other OT-based Flow Matching Models
  9. Mean-Flow Model
  10. Our Method: OT-Mean Flow
  11. OT Acceleration Methods
  12. Sinkhorn OT (Entropic Regularization).
  13. Linear Optimal Transport.
  14. Experiments
  15. Toy Example: Controlled Low-dimensional Point clouds
  16. ...and 28 more sections

Key Result

Theorem B.1

[Flow existence and uniqueness lasalle1968stabilityperko2013differentiallipman2024flow] If $v:[0,1]\times \mathbb{R}^d\to \mathbb{R}^d$ is continuously differentiable, then the ODE problem (eq:ode_u) admits a unique solution $\psi$. Furthermore, $\psi_t$ is a diffeomorphism for each $t\in[0,1]$, i.e

Figures (8)

  • Figure 1: Velocity visualization of a pair of points from the source and target distributions. The straight line denotes the average velocity from an intermediate time to $t=1$. The OT-MF trajectory is noticeably straighter compared to the vanilla Mean Flow.
  • Figure 2: Comparison of different transport-based mean flows for $N\to S-$curve. The first row is NFE=1, the second row is NFE=2. We report the 2-Wasserstein distance between the predicted distribution and the target distribution.
  • Figure 3: Single step (NFE=1) sample generation on ShapeNet Chairs and ModelNet10
  • Figure 4: Single step (NEF=1) shape interpolation on two samples from ShapeNet chairs. $\alpha\in[0,1]$ controls the interpolation of the source and target features.
  • Figure 5: Comparison of one-step mean flow method for Image-to-Image translation on Adult→Child and Man→Woman.
  • ...and 3 more figures

Theorems & Definitions (20)

  • Remark 2.1
  • Theorem B.1
  • Remark B.2
  • Theorem B.3
  • Remark B.4
  • Theorem B.5
  • Theorem B.6
  • Proposition B.7: liu2022flow
  • Remark B.8
  • Remark B.9
  • ...and 10 more