Flow Matching with General Discrete Paths: A Kinetic-Optimal Perspective
Neta Shaul, Itai Gat, Marton Havasi, Daniel Severo, Anuroop Sriram,, Peter Holderrieth, Brian Karrer, Yaron Lipman, Ricky T. Q. Chen
TL;DR
This paper introduces a flexible framework for discrete generative models using continuous-time Markov chains, allowing arbitrary probability paths and optimizing kinetic energy, leading to improved performance across various data modalities.
Contribution
It proposes velocity formulas for any probability path, decoupling probability and velocity, and introduces mixture paths that optimize kinetic energy in discrete models.
Findings
Outperforms masked construction in text generation.
Utilizes domain-specific probability paths in image generation.
Empirically validated across text, material, and image modalities.
Abstract
The design space of discrete-space diffusion or flow generative models are significantly less well-understood than their continuous-space counterparts, with many works focusing only on a simple masked construction. In this work, we aim to take a holistic approach to the construction of discrete generative models based on continuous-time Markov chains, and for the first time, allow the use of arbitrary discrete probability paths, or colloquially, corruption processes. Through the lens of optimizing the symmetric kinetic energy, we propose velocity formulas that can be applied to any given probability path, completely decoupling the probability and velocity, and giving the user the freedom to specify any desirable probability path based on expert knowledge specific to the data domain. Furthermore, we find that a special construction of mixture probability paths optimizes the symmetric…
Peer Reviews
Decision·ICLR 2025 Oral
I generally like this paper; it establishes a new framework for discrete flow matching and provides a method for constructing kinetic-optimal probability paths. The experiments across three downstream applications effectively demonstrate its real-world utility.
1. Is the construction in Section 4.2 unique, and is it general enough? 2. How does the proposed method relate to discrete rectified flow, and how do their equilibrium probability paths differ when using the same source and target domains?
1. The proposed method, Discrete Flow Matching (DFM), is theoretically well-justified through the use of kinetic optimal velocities and probability paths. 2. Extensive evaluations on text, crystalline material, and image generation tasks demonstrate the effectiveness of the proposed velocity formulation, achieving state-of-the-art results in material generation. 3. The construction of an infinite number of velocities for any chosen probability path broadens the design space for probability pat
1. Typo in Line 50: “misleadingly actually” --> "actually misleading" 2. The velocity $u_t(x,z)$ is only defined for $x\neq z$ in Eq. (5). The authors need to explicitly define $u_t(x,z)$ for the case where $x=z$, or explain why this case does not need to be defined separately. 3. What is $\bar{i}$ in Line 117? Define this when it is first introduced in the paper. 4. Line 227: Explicitly explain the relationship between the symmetric condition mentioned and the detailed balance equation in M
The generality of the proposed design space is definitely a significant contribution of the paper. Indeed, the new framework supports finding symmetric kinetic-optimal velocities for arbitrary discrete probability paths, expanding design flexibility over traditional masking. The paper has a strong theoretical foundation, and lays the ground for a more flexible approach to discrete generative modeling, which improves over state of the art, as highlighted by the experiments.
While the kinetic energy framework (Eq. 18) allows injecting problem dependent weighting $w_t(x,z)$ in theory, solving for arbitrary weights requires numerical approximation in practice, which can be challenging as mentioned by the authors. This can limit the practical potential of the design space.
Videos
No videos yet. Explain this paper in a talk, walkthrough, or lecture? Add one.
Taxonomy
TopicsFluid Dynamics and Turbulent Flows · Lattice Boltzmann Simulation Studies
MethodsDiffusion