Flow Diverse and Efficient: Learning Momentum Flow Matching via Stochastic Velocity Field Sampling

2025-06-10

AI Summary

To address the limited generation diversity and weak multi-scale noise modeling capability of rectified flow (RF) diffusion models, this paper proposes the Discrete Momentum Flow (DMF) model. DMF is the first to incorporate momentum into the flow matching framework, explicitly modeling variable-speed trajectories in the velocity field and enhancing trajectory diversity via stochastic noise injection in velocity space. It replaces the conventional constant-velocity assumption through discrete path decomposition and a dedicated momentum flow matching loss. Evaluated on multiple benchmark datasets, DMF achieves significant improvements in both generation diversity and image fidelity while preserving efficient single-step or few-step sampling. The method enables stable, high-fidelity, and diverse sample generation without compromising inference speed or computational efficiency.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Search and Optimization: Mixed Discrete/Continuous Search

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Vertical and domain-specific search

Abstract

Recently, the rectified flow (RF) has emerged as the new state-of-the-art among flow-based diffusion models due to its high efficiency advantage in straight path sampling, especially with the amazing images generated by a series of RF models such as Flux 1.0 and SD 3.0. Although a straight-line connection between the noisy and natural data distributions is intuitive, fast, and easy to optimize, it still inevitably leads to: 1) Diversity concerns, which arise since straight-line paths only cover a fairly restricted sampling space. 2) Multi-scale noise modeling concerns, since the straight line flow only needs to optimize the constant velocity field $m v$ between the two distributions $mpi_0$ and $mpi_1$. In this work, we present Discretized-RF, a new family of rectified flow (also called momentum flow models since they refer to the previous velocity component and the random velocity component in each diffusion step), which discretizes the straight path into a series of variable velocity field sub-paths (namely ``momentum fields'') to expand the search space, especially when close to the distribution $p_ ext{noise}$. Different from the previous case where noise is directly superimposed on $m x$, we introduce noise on the velocity $m v$ of the sub-path to change its direction in order to improve the diversity and multi-scale noise modeling abilities. Experimental results on several representative datasets demonstrate that learning momentum flow matching by sampling random velocity fields will produce trajectories that are both diverse and efficient, and can consistently generate high-quality and diverse results. Code is available at https://github.com/liuruixun/momentum-fm.

Problem

• Enhancing diversity in flow-based diffusion models
• Improving multi-scale noise modeling capabilities
• Optimizing variable velocity field sub-paths

Innovation

• Discretized-RF with variable velocity sub-paths
• Noise on velocity for improved diversity
• Stochastic velocity field sampling for efficiency