Flow Matching with General Discrete Paths: A Kinetic-Optimal Perspective

TL;DR

The paper broadens discrete-flow modelling by grounding discrete generation in continuous-time Markov chains and allowing arbitrary probability paths beyond masking. It develops kinetic-energy–based, closed-form velocity formulas that can generate any prescribed path, and introduces kinetic-optimal probability paths that recover mixture paths with domain-specific schedulers. A probability-preserving component and an ELBO for mixture paths enable stable training and clear optimization targets, with results demonstrating improvements over masking in text and competitive performance in materials and images. The work thus expands the design space for discrete diffusion-like models, enabling domain-informed paths and velocities to enhance iterative refinement in discrete data domains.

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 kinetic energy for the discrete case. We empirically validate the usefulness of this new design space across multiple modalities: text generation, inorganic material generation, and image generation. We find that we can outperform the mask construction even in text with kinetic-optimal mixture paths, while we can make use of domain-specific constructions of the probability path over the visual domain.

Figures (8)

• Figure 1: Generative perplexity vs. ELBO of kinetic optimal (KO) and linear schedulers of FineWeb-Edu models. The ELBO is evaluated: WikiText-103, LAMBADA, Penn TreeBank,FineWeb-Edu, and OpenWebText. Bold highlights the Pareto front.
• Figure 2: (left) Increasing the design space of discrete probability paths and velocities allows us to perform better than prior works, while significantly boosting performance at the low NFE regime. (middle) We find that the choice of kinetic optimal $u_t^\star$ significantly affects the low NFE regime while adding the probability-preserving component $u_t^\perp$ stabilizes the high NFE regime. (right) Comparison of FID values for discrete generative models.
• Figure 3: Generated samples for ImageNet 256$\times$256, with the same class label per column. (top) Autoregressive LlamaGen model sun2024autoregressive. (bottom) Discrete Flow Matching with metric-induced probability path (\ref{['eq:metric_prob_path']}).
• Figure 4: The conditional path for $a=0.9$, $c=3$ and $\mathrm{lp}=4$. This path is advantageous because the the path smoothly interpolates from noise to image while utilizing the whole interval $t\in [0,1]$.
• Figure 5: CIFAR10 Samples for 64 and 128 NFE, default velocities vs. optimized velocities. The default velocity we use is the velocity resulting from (\ref{['eq:Neta_flux']}). The optimized velocity searches over (\ref{['eq:Neta_flux']}) or (\ref{['e:power_inf_flux']}), and also searches over the probability-preserving velocity (\ref{['e:symmetric_flux']}) with varying weights. For each $8\times8$ table, same seed was used to generate the images.

Theorems & Definitions (5)

• Proposition 4.1
• Proposition B.1: Kinetic-optimal relaxation
• proof
• Proposition B.2: Kinetic Optimal paths.
• proof