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

Authors: Neta Shaul, Itai Gat, Marton Havasi, Daniel Severo, Anuroop Sriram, Peter Holderrieth, Brian Karrer, Yaron Lipman, Ricky T. Q. Chen
Primary area: generative models
Keywords: flow matching, discrete generative modeling

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.

Decision: accept (oral)
Reviewer scores: 8, 6, 8, 10, 6

Reviews and discussion:
<thread>
creation date:
1730370261832

thread_type:
official_review

<thread_object>
creator: reviewer_vzy2

<content>
summary:
This paper proposes extending discrete flow matching paths from a kinetic optimal perspective. Previous research in discrete flow matching has primarily focused on masking-related corruption paths. In contrast, this paper demonstrates that the design space can be broadened to include arbitrary probability paths. The authors explore two types of velocities: probability-advancing and probability-preserving. They also derive a closed-form velocity for the probability-advancing type and the ELBO for general discrete flow matching models.

strengths:
1. This paper extends the commonly used probability paths in discrete diffusion to a more general class of paths, providing a closed-form velocity and ELBO. This approach has the potential to inspire further research and improve a broad range of discrete flow matching models.

2. The proposed method is general and can recover previously proposed methods, such as Gat et al. (2024).

3. The experiments are comprehensive, covering image, text, and material generation. Results show that the proposed path achieves better performance compared to previous mixture paths.

weaknesses:
I have some concerns regarding the experimental results:

1. The results on ImageNet 256 fall short of state-of-the-art and even popular baselines, including both the autoregressive baseline and the proposed method. The network structure follows MaskGit; however, MaskGit reports an FID of 6, which is only half of the number reported in Table 3. What accounts for this discrepancy?

2. Why did the authors train the VQ-VAE for ImageNet rather than using off-the-shelf pretrained weights?

questions:
Please see the weakness part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730409400478

thread_type:
official_review

<thread_object>
creator: reviewer_kvbr

<content>
summary:
This paper takes a holistic approach on discrete generative models. It pioneers the use of arbitrary discrete probability paths. The methods are validated on text generation, inorganic material generation and image generation.

strengths:
1. This paper is very well written and succeeds in explaining concisely yet thoroughly the theory and background. It is well positioned in existing literature.

2. The authors propose holistic, novel methodologies for flow matching in discrete spaces.

3. The methods are experimentally validated on a broad set of modalities.

weaknesses:
None.

questions:
Is the use of the concept *flux* in this context novel, or already described in literature? If so, please provide a citation.

## Minor remarks
- There is a weird hyperlink at line 77 from the superscript 1 in $x^1$ to an equation in the appendix. This is also present in equation 6.
- line 174 'abusing a bit *the* previous ..'
- lines 202 - 206: it would be nice if the black triangles where vertically aligned
</content>
</thread_object>

</thread>
<thread>
creation date:
1730661408031

thread_type:
official_review

<thread_object>
creator: reviewer_s7cb

<content>
summary:
This paper proposes to broaden the design space for constructing better discrete flow generative models. To this end, the authors first decouple the probability paths from the velocity, effectively allowing to choose the velocity independently of the target probability path beyond the traditional masking method. Through a kinetic-optimal perspective, they propose closed-form velocities for a certain subset of symmetric kinetic energy, and further show that the mixture path model from previous work is kinetic-optimal when used with a specific scheduler. Additionally, a new tractable ELBO for mixture paths is introduced, which generalizes previous formulations. Experiments across modalities such as text, materials, and image generation demonstrate that the proposed framework outperforms existing methods.

strengths:
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.

weaknesses:
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.

questions:
I would be interested to know how do DFM approaches compare with FM approaches in the current state of the art. For example, with equal training time and computational resources, how does discrete flow matching compare with continuous flow matching in the task of pixel space image generation?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730675173935

thread_type:
official_review

<thread_object>
creator: reviewer_zxqn

<content>
summary:
Most existing methods rely on masking for the corruption process, which leaves the design space of discrete diffusion and flows underexplored. This paper introduces a new design space by optimizing symmetric kinetic energy, allowing the proposed velocities to apply across multiple probability paths. Thus, the proposed discrete flow matching method enables desired probability paths for domain-specific applications. The effectiveness of this approach is validated across various modalities — including text generation, material generation, and image generation — where it outperforms mask-based methods, especially in language tasks, and achieves a new state-of-the-art in crystalline material generation.

strengths:
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 paths, enhancing applications for downstream tasks.

weaknesses:
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 Markov Chains, if any. This would help clarify the theoretical foundations of the proposed method.

5. The authors are requested to provide a brief comparison of the key mathematical or conceptual differences between their kinetic optimal velocity formulation and the approach used in SEDD [1].

6. Typo in Line 511: “can be provide” --> "can provide"

7. Typo in Line 528: “We additional show generated” --> “We additionally show generated”

Overall, this paper is well-written with sound theoretical justification and experimental results supporting its main claims. The reviewer is unfamiliar with some related work (e.g. crystalline material generation) and is open to discussion with other reviewers/authors for a more comprehensive assessment.

questions:
Please see the weaknesses section above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730713178565

thread_type:
official_review

<thread_object>
creator: reviewer_zfhf

<content>
summary:
This paper broadens discrete generative modeling by introducing discrete Flow Matching models with flexible, iterative probability paths using continuous-time Markov chains (CTMC). Unlike conventional masked models, this approach enables optimized, domain-specific probability paths.
Key contributions:
A decomposition of path velocities, making generation simpler.
Closed-form velocities that allow any probability path, unifying and expanding prior methods.
A kinetic energy-based optimization of paths, leading to new, source-dependent schedulers.
A general ELBO that improves performance on mixture paths, surpassing masked models.

strengths:
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.

weaknesses:
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?

questions:
Could you provide more background on the mixture path and explain how it connects to the claim in Section 4.2?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732647408350

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We thank all the reviewers for their insightful review. We summarize the changes we did in the revision and additional experiments we ran following the reviewers' suggestions.  

**ImageNet** Following reviewer vZy2 concerns, we have investigated the gap in FID between our models and results reported by previous works [1,2].  The main reason is: we used one-dimensional positional embedding, compared to [1,2], that used two-dimensional positional embeddings. Following this, we reproduced the results from LlamaGen [1] (which is a state-of-the-art discrete autoregressive model on imagenet256). We modified Table 3 accordingly, with state-of-the-art FID values. Note: the key finding remains unchanged, metric probability path outperforms both autoregressive modeling and masked construction.


**Text generation:** After the submission we realized that our ELBO computation was not consistent with previously reported methods and would result in slight differences in reported values. For this reason, unprompted by any reviewer, we have since updated our results for the text experiments. In particular, we reimplemented SEDD and MD4, hence, we ensured that the results do not differ due to implementation or architectural differences. Figure 1 has been updated with new values; our key finding remains: the kinetic optimal scheduler lies on the Pareto front more times than the linear scheduler and thus can be a strong default scheduler for future works. Table 1 has also been updated, and shows results trained on FineWeb, a recent high-quality larger (by an order of magnitude) dataset that many works are adopting for text generation, consistent with Figure 1. We will add an autoregressive baseline to Table 1. Note that our key findings remain: our proposed kinetical optimal choices outperform SEDD and MD4, which use masked construction with linear scheduler, on most zero-shot evaluations. Furthermore, we obtain better perplexity on the training set, and the use of non-mask source distributions can further drastically reduce generative perplexity.

**Relation to SEDD [3]:** In response to Reviewer ZXqN’s request, we now include more  details regarding the relation between our approach and SEDD. To summarize, we provide three key observations:

1. **Relation between score and velocity.** The relation between SEDD’s concrete score parameterization of the time-reversal rates and our velocity rate can be written in a closed-form, akin to the relation of score and velocity of Gaussian paths in the continuous case. We show we can always convert from a SEDD parameterization to equivalent velocities in our formulation, but not vice versa.

2. **Our probability paths strictly generalize SEDD’s paths.** SEDD builds probability paths by: (i) considering rate matrices of a forward (noising) process; (ii) figuring out its marginal probabilities in closed form; and (iii) expressing the reverse-process rate, which is used for sampling, using the forward-process rate and closed-form marginal probabilities. To accomplish (ii), SEDD requires solving an ODE in closed form and consequently introduces heavy restrictions upon the noising rate matrices, considerably limiting the reverse-processes that can be considered in SEDD. In contrast, our approach to build probability paths is: (i) writing down the desired probability path that interpolates the source and target; and (ii) find generating velocities for these probability paths. In particular we show that (ii) can be solved in closed form for general probability paths that form a strict superset of SEDD’s paths. 

3. **SEDD loss is an instantiation of our ELBO.** The training loss of SEDD is a particular case of the ELBO in our formulation. This implies that SEDD and MD4, since they both use mask with a linear scheduler, differ mainly in the parameterization of the reverse-process. 

Points 1 and 2 both impose restrictions on the probability paths / corruption processes that SEDD can work with. The derivation details and justifications for these claims are provided in Appendix F.

[1] Sun, Peize, et al. "Autoregressive Model Beats Diffusion: Llama for Scalable Image Generation." arXiv preprint arXiv:2406.06525 (2024).

[2] Chang, Huiwen, et al. "Maskgit: Masked generative image transformer." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.

[3] Lou, Aaron, Chenlin Meng, and Stefano Ermon. "Discrete diffusion language modeling by estimating the ratios of the data distribution." arXiv preprint arXiv:2310.16834 (2023).
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper makes a significant contribution to discrete generative modeling by broadening the design space of flow matching methods, allowing the use of arbitrary probability paths with a strong theoretical foundation grounded in kinetic-optimal velocities. Reviewers highlighted the paper's generality, elegant theoretical formulation, and its ability to unify and expand upon prior methods, while achieving state-of-the-art performance in tasks such as text and material generation. The experimental validation across diverse modalities, along with the clear and well-written presentation, further demonstrates the real-world utility and impact of the proposed approach.