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
Venue: ICLR 2025 Oral (2025)
Keywords: flow matching
discrete generative modeling
Primary area: generative models
TL;DR: Through the lens of kinetic optimality, we expand the design space of Discrete Flow Matching, allowing the use of any probability path and simultaneously justifying existing mixture paths.

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.

Published as a conference paper at ICLR 2025
FLOW MATCHING WITH GENERAL DISCRETE PATHS: A KINETIC-OPTIMAL PERSPECTIVE
Neta Shaul -1, Itai Gat?, Marton Havasi? Daniel Severo?, Anuroop Sriram?, Peter Holderrieth Brian Karrer?, Yaron Lipman', Ricky T. Q. Chen?
Weizmann Institute of Science Metaa FAIR, "MIT CSAIL Work done during internship at Meta FAIR
ABSTRACT
The design space of discrete-space sprce diffusion or flow generative models are signif- icantly less rell-understood than their continuous sspucc counterparts, with many works focusing only a simple masked construction. In this work, we aim to take a holistic approach to the construction of discrete generative models based on continuous-ting 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 be applied to any given probability path, completely decoupling the probability and velocity, and giving the user the freedom specify any desirable probability path based on expert knowledge specific to the data domain. Furthermore, we find that 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 con- struction 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.
INTRODUCTION
Generative models over discrete spaces have not seen as much progress on the methodology side compared to continuous-space counterparts. For the most part, applications such as large language modeling rely solely on autoregressive models (Radford et al., 2019; Bommasani et al., 2021). The simplicity of autoregressive modeling has also motivated people to use them for multimodal gen- eration, where other modalities, such as images and videos, are tokenized and modeled within an autoregressive framework (Van den Oord et al., 2016; Team, 2024; Sun et al., 2024). While obtain- ing reasonable results, they have not yet reached the performance of continuous- generative models such as denoising diffusion (Ho et al., 2020; Song et al., 2021) and Flow Matching models (Lipman al., 2022; Albergo et al., 2023) for the visual-audio domains (Rombach et al., 2022; Dai et al., 2023; Esser et al., 2024; Zhou et al., 2024), where is believed that the ability to perform iterative refinement brings significant gains (Saharia al., 2022: Zhang al., 2024).
A promising framework that brings iterative refinement to the discrete case is to consider the use of Markov chains within dynamical generative framework. Many discrete-space generative flow and diffusion models have seen success in the generation of text (Austin et al., 2021; Lou et al., 2024: Shi et al., 2024: Sahoo al., 2024; Gat et al., 2024), proteins (Campbell al., 2024), images (Austin et al., 2021; Shi et al., 2024), and even executable code (Gat et al., 2024). However, the design space of these models currently rather limited, with many recent works instead focusing solely on the case of masking as corruption process (Shi et al., 2024; Sahoo et al., 2024). The masked construction is an extension of masked pretraining (Devlin, 2018; Yang, 2019), but it does not fully embody the concept of iterative refinement as it is equivalent to learning autoregressive models for every ordering (Hoogeboom et al., 2021; Chang et al., 2022), and has been noticed actually misleading due to low- low-piccn sampling
that some of the recent reported progress (Zheng et al., 2024) rather than the explicit design choice of masking as a comuptiio process. In spite of this, the masked construction has often been found to be the best performing choice out

---

Published as a conference paper at ICLR 2025
of the limited family of corruption processes previously considered tractable (Austin et al., 2021; Campbell et al., 2024; Gat et al., 2024).
We instead take holistic view constructing discrete Flow Matching models, massively expand- ing the design space to enable arbitrary probability paths, or colloquially, corruption processes, grounded in the framework of continuous-time Markov chains (CTMC). We list our contributions:
1. Analogous to the continuous setting, we find that an infinite number of velocities can gen- erate any given probability path. In order to reduce this search space, we consider de- composition into a balbaability-dvvaccii velocity and a probabilityprresenee
velocity. 2. To explore the space of probability-advancing velocities, we motivate family of closed- form velocities that be formulated as optimizing kinetic energy. In particular, we are the first to formulate velocities that can work with any choice of probability path, completely opening up the design space of probability paths, e.g., domain-specific constructions, while recovering the velocities used by prior works for existing paths in the literature.
3. We also find that the probability path itself can also be optimized with the same kinetic energy criterion. A closed-form solution surprisingly recovers the mixture paths considered by Gat et al. (2024) but with novel source-dependent schedulers.
4. We derive the ELBO for discrete Flow Matching models in full generality. This leads to an improved ELBO for training mixture probability paths that has not been used before, and recovers the ELBO derived by Shi et al. (2024) for the masked construction. We find that with this ELBO, our kinetic-optimal mixture paths outperform the masked construction.
2 BACKGROUND: DISCRETE FL.W MATCHING
We are interested in learning generative model that approximates a data distribution ((z), where (x² 2 ..... D) E S = TD with T = [K] A (1. K} being discrete set of possible token values, and D E number of discrete variables. For brevity and without loss of generality, we consider all dimensions to have the same number of discrete values.
Probability paths. We denote by p(x) and g(x) the source and target, respectively, probability mass functions (PMFs) over the state space S. We consider probability paths (r()). E [0, 1J. to be time-dependent PMFs taking the form
(1) and Pr (x" (xf) conditional probability path which interpolates between simple PMF time t = 0 and delta PMF centered around xi at = 1. That is, we assume the boundary conditions Po(="I>)) = p(x') and Pı (1' [rt) = dzt (x*). Hence we can interpret these probability paths pe (x) in equation 1 as interpolating between factorized source distribution p(x) & II p(x²) and the data distribution (xx). A common family of probability paths used in previous works is the collection of mixture paths (Gat al., 2024), with x dependent schedulers similar to Shi et al. (2024):
(x)))= Ex, Es ((())11))(11), where (x(xx) & P1
Pr
(r [xf] = (1 - Ne (xf)))))()) ()x(')) + K. (x.)) (x*),
where k-(-) = 0 and k-)-) to satisfy the boundary conditions. Specifically, with p(x') 8mm (x') we recover the masked construction (Shi et al., 2024; Sahoo et al., 2024).
Probability velocities. As generative process, we simulate Continuous Time Markov Chain (CTMC) (X1)eefo.1) in S such that its time marginals follow a prescribed probability path,
X1 Pt-
(3)
In order to do so, we define the concept of probability velocity, also known as rate matrix. We say that a probability velocity u, generates Pt if u, characterizes a Markov process X, with marginal Pt (equation 3) for all t E [0, 1) in the following sense:
P(Xe+h+ = x X, = =) = (z)) + hug 2) + o(h),
(4)
where (h) denotes function which asymptotically smaller than h, ie., lim o/h/// 0.
Intuitively, it describes the Markov transition of X, for small step sizes h > O. We note that for equation 4 to be a valid PMF, u, must at least satisfy the Rate Conditions:
Ue u1 2) 2 0 for all # and E₂ ve(x,2) 2)
Rate Conditions
(5)
2

---

Published as a conference paper at ICLR 2025
Single-wariable--hhngg probability velocities. It is natural to consider modeling CTMC process X₁ over S by defining a u, 2) for all pairs 2, 2 E S. However, the state space is of size [T] so this is generally prohibitive for high dimensions. A remedy is to consider rates that only allow state to change in a single variable (Campbell et al., 2022), egg., in the following example we only change the variable the i-th coordinate:
(6) To model only such changes we restrict our attention to velocities of the form w/(x', 2) that de- scribe the probability rate between the state z and the state with the i-th coordinate replaced, i.e. as described in the r.h.s. in equation 6. We can express the full velocity ux(x, 2) via u/ 2) as
(2
2) - (2
D
ue(x, 2) = Σ ua
I) (x'). jih
which states the probability velocity between two states = - x is zero if they differ by more than variable and equal u,(x', 2) if they differ by exactly one variable. Plugging this velocity into equation 4, it can be shown that (Gat et al., 2024):
P(X++h x X2 = 2) = II=
(1(((x)))) 2)] + o(h)
(8)
This implies we can sample each variable Xith independently from the distribution 8g (x') hu)("' 2)) and only incur error of o(h).
+
The marginal velocity, Previous works (Campbell et al., 2024; Gat et al., 2024) have shown that constructing generating velocity for pı(z) be achieved by considering only the conditional probability paths in equation That is, assume we have conditional velocities ut (x", =* which velocities in the state space T, that generate the conditional paths (x(1(1z) in equation 1. Then a marginal velocity wi(x", z) that generates pr (x) takes the form:
2x(x²=)) ΣxŞET ul(xr', (1/pr(
where (11z)) is the posterior probability of the i-th token taking the value 21, i.e.,
(x*1)) = dan (2²*) pr(slar p.(2)
(10)
Parameterizing the factorized posterior 11 Piee is an approach taken by prior works (Austin et al., 2021: Campbell et al., 2022). To train, simple option is the cross-entropy objective:
(0) = EE- 2010))
[- Σl_ [og])(((((())
(11)
We use this training loss for general probability paths as is generally applicable. However, for the case of mixture paths (equation 2) it is possible to derive a tractable ELBO as the marginal can be written closed form without a summation as in equation We cover this later in Section
SAMPLE GENERATION THROUGH THE FACTORIZED POSTERIOR
The most direct approach to sample from this model is to use the marginal velocity uz(x', 22, e.g., with a first-order sampling scheme defined by removing the o(h) term in equation 8, Le., given X we advance time with step size h by sampling X according
X ~ 6x;(,) + (.,((,,)),
(12)
for each [D), where of is computed with equation 9. However, for general discrete paths this sampling procedure intractable for large discrete spaces as computing (x*,,)) with equation for all x' E T has a computational complexity of !T12.
Alternatively, we propose a more efficient sampling scheme by noticing that
_. (x') + hut (x+, 22)
Σ
EJET (x*) + hug(a* =*|=[)] Pije (xf,
(13)
which leads to a sampling process that avoids computing the full marginal velocity: given the current state X,, sample X, from the factorized posterior, then sample X That is, for each E [D].
3

---

Published as a conference paper at ICLR 2025
1) Sample Xi ~ Pije (1);; and
2) Sample Xi+ 6x; (.) +
This sampling procedure still results in X, with the same time marginals while avoiding the compu- tational cost of the summation in equation 9. To enable the use of any step size h, we use slightly modified step 2; see Appendix for more details and pseudocode Algorithm 1.
KINETIC OPTIMAL VELOCITIES AND PROBABILITY PATHS
We first decouple the design space of probability paths and their generating velocities, providing the means to effectively explore this large design space. This section covers two contributions: (i) propose family of kinetic optimal (KO) velocities that generates any given probability path, and (iz) we solve for kinetic optimal probability paths, recovering a special case of mixture paths. The first contribution enables us to work with general discrete probability paths. The second con- tribution justifies the choice of mixture probability paths used by Gat et al. (2024) but offers novel af dependent schedulers. For both, we center our designs based on optimizing a discrete notion of kinetic energy (Peyré et al., 2019).
Notation. As the discussion in this section applies to arbitrary probability paths and discrete state spaces, we will use a simplified notation, where our state space is now T and for states we use 2; 2 E TT abusing a bit the previous notation (where x², 2² E T). Furthermore, we will denote by Pc(z) and (z z) arbitrary probability path and velocity field in T, respectively.
Continuity Equation. Given a probability path pe(r), the entire collection of velocities (I, 2) generating Pr (z) are the solutions to the Continuity Equation (a.k.a. the Kolmogorov forward equa- tion) that also satisfy the Rate Conditions. is useful to formulate the Continuity Equation through the flux Ji, that is
(x) + div = O, Vx ETE
with je 2) = ux (2).
(14)
Intuitively, the flux je(r,, quantifies the amount of probability mass per unit of time moving from state z to state x. The divergence operator then measures the total outgoing flux minus the total incoming flux, which in the discrete case takes the form
div..(ie) = E=faJe(E, Exje
(15)
Velocity from flux. Given a flux je satisfying the Continuity Equation (equation 14) we can get velocity from the flux by defining for æ +
u 2) je(z, 2) pp())) if pf
> 0, else ue uz 2)
(16)
and the case x = 2 is uniquely set by the Rate Conditions (5), ue (2, 2) = - -unee ue( z). The
velocity defined in this way will satisfy the Continuity Equation and the Rate Conditions if the flux satisfies the following conditions:
fx 2) 2 0, for x x
Non-negatyvity
(17)
(=(2)= = 2x
Safe Flux Condition
(18)
Intuitively, the Safe Flux Condition ensures no flux is leaving a zero probability state 2.
Proposition 4.1. Given non-negative safe flux j: that satisfies the Continuity Equation, the velocity defined in equation 16 satisfies the Rate Conditions and generates the Pt probability path.
Kinetic optimality. Motivated by the approach employed in the continuous case of minimizing the kinetic energy for the conditional velocities (Lipman al., 2022; Shaul et al., 2023), we take similar approach for finding velocities for the discrete case. The standard convex formulation of the kinetic energy adapted the discrete case (Peyré et al., 2019):
minp..de
So Extr vi(222)) je(x, d
Kinetic Energy
(19a)
pe(3)
s.t.
div (iz) = (x)
EIET
Continuity Equation
(19b)
(1 2) 2 0,
ZET
Non-negative flux
(19c)
Pa = pl =
Boundary conditions
(19d)
4

---

Published as a conference paper at ICLR 2025
where 2x, O is some problem-dependeen weighting; higher weight implies smaller flux from 2 - x, i.e., the higher this value the smaller the velocity uz 2). The optimality criterion (equation 19a) is the kinetic energy, equivalentyy Pe(2) ue =) (z). The benefit of formulat- ing terms of the flux (instead of the velocity) is that the problem becomes convex in its unknowns (rt, and in particular the Continuity Equation constraint in (19b) is linear. Lastly, in case of
wx(xs,s) oc the energy in equation 19a is defined to be 0 if je 2) = O, and oo je(x, 2) O.
(2)
Therefore, to ensures the solution jł is safe (equation 18) we ask that wy satisfies:
(c(2) 0 =
(e. 0c Pi(2)
Safe Weight Condition
(20)
and that problem 19 is feasible, i.e., it has a finite energy solution. Although problem 19 is convex, solving it for general WN requires numerical approximation. Since we want to solve it for con- ditional probability paths with different #1 T. i.e., q(x) 6z, (a:), this can be computationally challenging. Instead, we will explore cases of wi where problem (19) is solvable in closed-form. We start with assuming Pt is mowwn/given, and find the kinetic optimal velocity u,, then afterwards we discuss optimizing the pt as well.
4.1 KINETIC OPTIMAL VELOCITY
Assuming Pt > is fixed in (19), our goal is to find the kinetic optimal solution jł, and consequently obtaining a velocity in via (16). One observation make is that (19) can be efficiently solved when symmetric, i.e., when Pt(2] (x.
(xr_2)
As we prove Appendix B, (19) can be efficiently solved
p=(2)
via the following linear relaxation:
Σ
p=(=)
Va E T
(21)
(2.3)
where : R is the unknown function over the state space. The linear equation in (21) is of Laplacian form, and many properties (including elosed-form solutions) are known in many cases (Vishnoi, 2012). The solution fo to (21) is unique up global constant and using ft we construct the kinetic optimal flux,
(x 2)
po(=) wi()s) [f(()) - f((=)]+
(22)
where [s], max {s, 0} is the ReLU operator. This provides solution to (19) with fixed and positive Pt. Consequently, using (16) we get the kinetic optimal velocity. We have shown that a certain family of kinetic optimal velocities can be computed by solving linear system (21) for arbitrary probability paths Pa over state-space T. Next we will further instantiate this family and provide some closed form solution for jž and uz.
Closed-form U,. We will consider the case where (x, (3) and T, : T + R2 is a design choice of our method. To ensure wy is safe (20) we require that pz(2) = 0 implies Tz(2) = 0. can be checked with substitution
pe(s)
The solution ft to
fe(x))
(((x) E.er (3) T2(x)
(23)
One choice is T. (x)) mm(x))=00)-- that leads to the Kinetic Optimal flux ji (xx = 1 [0)p((x) (x)( for x x
(24)
which upon converting to velocity via (16) recovers the velocity proposed in Campbell et al. (2024)
for positive paths, Pt > O. Note however, that the above flux is not safe (does not satisfy equation 18) the flux jē (xr, 2) for some general x is not necessarily small, showing potential
and if pz(2) numerical issue. Campbell et al. (2024) formulate a limit case for general pt that also requires adding an extra assumption on Pt (that Pc(x) 0 = (e(x) = 0), which does not hold even for common probability paths that are typically used, such as the masked mixture path with linear schedulers.
Alternatively, we propose a more numerically stable choice. Consider (x)x) = Pe())) Le..
wo(2, 2) (zp.
(25)
This results in f((x) (x(x))//(p) and the kinetic optimal flux in this case is:
5

---

Published as a conference paper at ICLR 2025
(z (p((2))(())) (x))+ for x +
(26)
Note that in contrast to before, this flux is safe (satisfies equation 18) and therefore works for general Pt. Furthermore, (26) exhibits stable limiting behavior for continuously differentiable ps: when
(=(2) + 0, so too will ji, 2) + O.
We note that for uniform and mask source distributions with the mixture path (2), the velocity considered by Campbell et al. (2024) and our velocity resulting from (26) coincide. However, for mixture paths (2) and general discrete paths, they generally do not coincide. Additionally, the choice of velocity in (26) also recovers the velocities used by Gat et al. (2024) for mixture probability paths. See Appendix C.1 for detailed derivations. Finally, we discuss broader family of closed-form velocities involving different choices of T; in Appendix C.3, which we find can significantly boost performance at low-cost sampling regimes.
Metric-induced (z)).. The velocity resulting from (26) can be applied to any user-defined Pc. We propose metric-induced conditional probability paths of the form
(((2)(1)) softmax (-3,d(z,z1))
(27)
where B : 00 12 is monotonic scheduler with Bu 0, 81 and d T x T + Rz0 such that d(x, 21) = 0 + x = I1 iineerpreted loosely as metric over discrete values. If we apply the flux in (26) for the paths in (27) and simplify, obtain the velocity:
" (z, 2| = (t (x) [ log Po - 0 log (e = pe ( (d(2, d(.r,=1))+
(28)
(29)
This velocity has the property that we only move from state = to state x if x is closer than 2 to i... d(x, x1) < d(z, hence resulting in flow that only moves closer to TT.
4.2 KINETIC OPTIMAL PROBABILITY PATHS
Interestingly, for the weighting choice that we have motivated for the numerically stable velocity (25), is also possible solve for the kinetic optimal probability path př. As we show in Ap- pendix B, in this case, the problem (19) can be formulated equivalently as
minler J. E Edlla)22 dt
Kinetic Energy
(30a)
s.t. [x)(x)= 11
11 E [0, ao(z) Vp(x), (1(x) Vg(z)
Hypersphere constraints
(30b)
Boundary conditions
(30c)
where (z) (p) Problem (30) is the kinetic energy of a curve over the hypersphere connect-
ing vD and 7q. The optimal solution thus corresponds to the geodesic curve on the hypersphere,
ax(x)
sin(1-812
p(x)
sin x sin 17
g(z),
where Ω = arccos E. vp(2)a(3))
(31)
52
and consequently the optimal probability path and velocity for (30) are
pi (x) (2),
ut (a, 2) (x))) [a log at (ar) a log of (2())]
(32)
In the particular case of conditional probability paths (अx) = 82, (z). we get that the optimal solution recovers the mixture path (equation 2) with specific 2, dependent scheduler:
Ke (x1) 1 sin ( (1-1)((a1)
where ((z)) arccos /p(x1).
(33)
sin (1(1)))
This justifies the mixture paths (2) as kinetic optimal, and furthermore, it naturally utilizes an ø1- dependent scheduler for general source distributions p when p(x1) > 0.

---

Published as a conference paper at ICLR 2025
5 PROBABILITY-PRESERVVN VELOCITIES
While we have found a particular flux J;, the space of fluxes for a given Pt is much larger, and in this section we show how to explore it further. We first observe that since the Continuity Equation (14) is a linear equation, any flux ji satisfying this equation can be written as a sum of two fluxes:
je = + JE. where div=((r'' = O,
(34)
where jī is particular solution to the Continuity Equation and j is solution to the homogenous version of the equation, i.e., divergence-free. We call the velocity resulting from je a probability- preserving, or corrector, velocity sampling with this velocity has pa steady-state. For simplic- ity, we mainly consider the special case of symmetric flux. Symmetry is a sufficient condition for being divergence-free is evident from (15). A natural choice for symmetric flux is consider symmetrization of (22) taking the form
ve(a) jit (1r, 2) (=(
[f(())) f((=)) and uf 2) (1, 2)/pe(2),
(35)
for any function ft- For convenience, the kinetic energy (19), e.g. the same as in (26). In contrast to the kinetic optimal velocity, which results in a unidirectional flow in the sense that samples will only move from lower fe(+) to higher fe(-), the symmetric flux in (35) results in a bidirectional flow that allows equal movement between any two states with non-equal fc(.). Hence Jr acts as corrector to redirect samples back previous states in way that leaves p. invariant.
will simply re-use the same fe that comes from optimizing
ELBO FOR DISCRETE FLOW MATCHING
We show in Appendix D that we can produce a continuous-time ELBO bound on the likelihood log (x1 for any conditional probability path and conditional probability velocity in terms of the marginal ut (x+1, 2) and conditional uf(x,, 2'1x+) as follows
log (x1) So Encep EP
(u 21) u/(xt, 2 ml1=1)
(36)
En uf(s x szlaa) log
uf(y" uf(s"=)ik5) dt
Evaluating this ELBO is difficult for the same reason as sampling in Section 3, for large discrete spaces T computing (9) for x² E T has a computational complexity of [T However, for mix- ture paths (2), our conditional velocity resulting from (26) is used to obtain a closed-fore expression for the marginal velocity (see Appendix C.2), yielding a tractable ELBO for mixture paths:
log (())) 2 Jo Bacom(le))) Eln
(
Σ ((r')pf
(37)
-1- (=t 1o
log log pfje (at)(((( de,
[=] This ELBO has not been used previously, e.g., Campbell et al. (2022) had to resort to a doubly stochastic estimator. Specifically for the masked construction, we recover the ELBO used by Zheng et al. (2024) for 2 sr-iineepndeen schedulers and used by Shi et al. (2024) for of dj-dependent schedulers; see Appendix D.I.
where A(x) = (xr
7 RELATED WORK
Generative modeling through marginalization. Denoising diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2021) construct generative models by reversing a nois- ing process. The Flow Matching framework (Lipman et al., 2022; Albergo et al., 2023; Liu et al., 2022) shares similar traits but instead constructs generative models through marginalization of conditional Markov processes, allowing a larger design space of probability paths. These types of models can be trained at scale both efficiently and stably relative to other frameworks, and thus have seen massive success in the large-scale generation of images (Rombach et al., 2022; Esser et al., 2024), videos (Singer et al., 2022), and audio (Ie et al., 2024; Vyas al., 2023).
7

---

Published as a conference paper at ICLR 2025
WidTestt-13
LAMBADA
Penn Treebank
FineWWeb Edu
OpenWeeslext
102
103
102
23
03
92
93
93
Leeer
4.0
93
90
90
90
92
8
80
80
80
-4
73
75
75
75
200
EU80
EL0O
ELBO
ELBO
ELBO
Figure 1: Generative perplexity vs. ELBO of kinetic optimal (KO) and linear sched-
ulers of FineWeb-Edu models. The ELBO is evaluated: WikiText-103. LAMBADA, Penn
and OpenWehTexs Bold highlights the Pareto front.
Table 1:
Zero-shot unconditional perplexity bound as equation 36 of Fineweb- -du models, more details are in Appendix E.1. denotes our reimplementation of the method.
METHOD
LAKBADA
WIKITEXY2L
PTBL
WIKITEXTI03_
18WL
OPENWeSTEXIL
Pisewes Eou
SEDD" (mask)
2024)
$58.57
<42.84
133.99
<42.8
$114,17
$36.55
$19.41
MD4" 2024)
<43.08
$157.00
<43.02
1122.55
$35.57
$18.99
DFM Linsar (A, 1034)
$60.59
<44.17
<44.20
1442221
DFM Kinztic Oytinell (msok)
3:5855
341.80
3148.66
241.83
2123.83
318.71
Optimal UB, 1024)
558.41
242.19
144700
242.34
2115.51
Continuous-time time Markov Chains. These aforementioned frameworks have been adapted to the discrete domain by making use of Continuous-time Markov Chains (CTMC) as the choice of gener- ative process (Campbell et al., 2022; 2024; Lou et al., 2024; Gat et al., 2024). Many works discuss both uniform noise and mask construction (Campbell et al., 2024; Lou et al., 2024); however, more recent works have focused more and more the simple masked construction where each discrete variable is randomly replaced with a dummy or mask token (Sahoo et al., 2024; Shi et al., 2024) as it often performs favorably compared to uniform noise. However, the simple masked construc- tion leads to a generative model that is equivalent to any-order autoregressive modeling under mild assumptions (Hoogeboom et al., 2021: Zheng et al., 2024).
Any-order autoregressive modeling. While autoregressive models prespecify a fixed ordering, any-order autoregressive models learn conditional probabilities for every ordering. Training is often carried out by randomly masking out parts of the data sample (Devlin, 2018; Yang, 2019). Some works have focused on architectural choices: Germain al. (2015) randomly masks out weights to induce randomized ordering, while Pannatier et al. (2024) uses a fixed causal attention architecture but randomly permutes the input ordering, with the end goal of learning all combinations of condi- tional distributions so that generation of the variables can be done in any order. The ordering itself is often optimized further by the use heuristic scoring functions (Chang et al., 2022; Ziv et al., 2024).
8 EXPERIMENTS
We evaluate Discrete Flow Matching (DFM) on multiple modalities: text, crystalline material, and image generation. Our main goal to show that Discrete Flow Matching can outperform autore- gressive models, and within the class of Discrete Flow Matching, explore new additions such as the kinetic optimal and the metric-induced constructions. In text, we mainly explore the kinetic optimal probability paths (equation 33) with different source distributions, these have access to the closed-form ELBO (37). In material generation, we find that enabling permutation invariance for DFM easily outperforms autoregressive models at de novo generation, achieving state-of state-of-the-an art results. Furthermore, in domains where a natural metric exists, we demonstrate our method's 's abil- ity to inject inductive bias into the velocity and probability path using equations 27 and 29, We show that our large design space enables competitive results even with non-mask probability paths, showcasing the capabilities of our expanded design space.
8.1 TEXT GENERATION
We explore our method on the task of text generation. We use the kinetic optimal probability path as in equation 33, which only has one hyper-parameter, the source distribution p(x). For source distri- bution, we compute the statistics of tokens appearances in the training data Pstans (x*) and construct family of source distributions:
p(a)= IP=))(
p(x'), p(x*)= softmax log Panas (2:))
(38)
8

---

Published as a conference paper at ICLR 2025
13
10
(pe
Method
FDD
D3PM (Austin al., 2021)
7.34
02
04
125
226
512
1054
CTDD (Niscnoff ct al., 2024)
7.86
LLDR- (Camphell et al., 2022)
3.74
DFM w/ mask et al., 2024)
3.63
DEN cluse (36 40 2004)
DUM mutti 10
DFM w/ metric (Ours)
3.43
1036 2046
255
212
1024
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 "i significantly affects the low NFE regime while adding the cebbability-prservinn component stabilizes the high NFE regime. (right) Comparison of FID values for discrete generative models.
Table 2: De novo material generation. Our primary metric, Stability Rate, is the fraction of materials with energies below the convex hull formed by stable materials, following Miller al. (2024).
Validity (%)
Coverage (
Property L
Stubility (%)
Stractural
Composition
Recall
Precision
wdist (p)
wilsst (N,r
CDVAE (Xe et al., 2021)
5000
100.00
86,70
99.15
99.49
0.688
0.278
1.57
DiPCSP Olio 2023)
2023)
1000
100.00
83.25
99.71
99.76
0.350
0.125
5.06
FlowMM (Miller
0.083
CrysalLLM (70B) (Gruver al., 2024)
99.6
95.4
85.8
98.9
0.81
0.44
5.28
Autoregressive
86.43
89.33
63.31
99.74
0.088
0.030
1.99
Pemm invoriant DFM Mask w/ Cubic
250
94.40
84.40
98.25
99.40
0.244
0,144
6.90
Perm invariant Mask Kinetic Optimal
90.11
0.151
recovers the data statistics, 930 yield a uniform distribution on all tokens. Also, 90 -> co yields a uniform distribution on the set of least probable tokens in the data, which behaves similarly to a mask source distribution.
where 80
For this experiment, we used linear and kinetic optimal schedulers with mask, p(x) 5m (zr). and 8o E (-0.5, 0.0, 0.5, 1,2, 4, 64, 256, 1024) source distributions. The models are trained on the FineWeb-Edu (Lozhkov et al., 2024) data set. Table compares the evidence lower bound (ELBO), as in equation 37, of our trained models with previous works. See Appendix E. for experimental setup. We find that the kinetic optimal scheduler yields the best results most of the evaluation sets. Notably, to the best of our knowledge, this is the first time a non-mask source distribution obtains comparable results and sometimes outperforms the mask source distribution. Figure 1 presents a view of the generative perplexity (as measured by GPT2-large) vs. the ELBO of each model. Generative perplexity represents the likelihood determined by external model, whereas the ELBO indicates the likelihood as assessed by the evaluated model. We see that models trained using the kinetic optimal scheduler achieve better tradeoffs than those trained with the linear scheduler, as they more frequently appear on the Pareto front.
8.2 CRYSTALLINE MATERIAL GENERATION
To showcase the flexibility of our approach, we use discrete Flow Matching to generate crystals. We train on inorganic materials from the MP-20 dataset, a subset of the Materials Project database (Jain al., 2013). Crystalline materials are represented using combination of continuous and discrete variables, which we tokenize using the same method as Gruver et al. (2024), which fine-tunes 70B LlaMa-2 autoregressive model (Touvron al., 2023). In contrast, we are the first perform crystal generation with a purely discrete non-autoregressive model.
An important distinction that since discrete Flow Matching directly predicts the factorized pos- terior (10), we can easily impose permutation invariance of the atoms, which should significantly reduce the complexity of the learning problem. This is as opposed to prior works on using au- toregressive models for material generation (Flam-Shepherd & Aspuru-Guzik, 2023; Gruver et al., 2024) which must impose an unnatural ordering on the variables. We show results in Table 2 where achieve state-of-the-an results using discrete Flow Matching, in particular, with kinetic optimal scheduler (33). We believe non -on-autoregressvve generation is key ingredient in performing well due to the ability impose structure such as permutation invariance. Compared to continuous-space

---

Published as a conference paper at ICLR 2025
DFM metrie LlamaGen
Figure 3: Generated samples for ImageNet 256x256, with the same class label per column. (top) Autoregressive LlamaGen model (Sun et al., 2024). (bottom) Discrete Flow Matching with metric- induced probability path (27).
models such as FlowMM (Miller et al., 2024) and DiffCSP (Jiao al., 2023), large per- formance gain in terms of our main metric, stability rate (2 38% relative improvement), from using discrete generative models due to the discrete nature of crystal generation.
8.3 PIXEL SPACE IMAGE GENERATION
We first consider the case of image generation in pixel space. Here, T= 255} and we have access to a natural choice of metric, by embedding T on the interval (-1, [] C and using the Euclidean distance d y) Ir y) in (27), is typically done for continuous-space image generative models. We use the CIFAR-10 dataset (Krizhevsky et al., 2009) for these experiments. can improve upon the masked construction while also retaining performance at low number of function evaluations (NFE). Generated samples are shown in Figure and in Appendix G. We find that optimizing the velocity after training can provide significant gains: the choice of probabblity--ddvnnee samples while the adding the probability prreeerrityy component (Section 5) improves at high NFE.
Results are shown in Figure 2, where
velocity (Appendix C.3) affects the low NFE
8.4 DISCRETE LATENT IMAGE GENERATION
Table 3: Face-blurred ImageNet- 256 with the
Llama-B B architecture (111M parameters). denotes our reimplementation.
We also explore the use of discrete Flow Match- ing as generative model within a discrete latent space learned by vector quantized variational au- toencoder (VQVAE; Van Den Oord al. (2017)). We use images from face-blurred ImageNet (Deng al., 2009; Chrabaszez al., 2017) at 256x256 resolution. For training the VQVAE model, follow the setup in Sun et al. (2024) and use 16x downsampling produce latent space of dimen-
Method
NFE
FID
LlamaGen (AR) (Sun et al., 2024)
256
5.46
LlamaGen (AR)* (Sun al., 2024)
256
4.81
DFM Mask" (Gat et al., 2024)
100
5.72
DFM Metric (Ours)
100
4.50
sion 16x 16 with codebook size of [ 2 16384. As our choice of d(., we use the same metric that used to train the VQVAE model, which is d(æ, y) ||~// s/ll=ll]- We show quantitative results in Table 3, where we find that discrete Flow Matching model with the metric probability path outperforms the autoregressive approach, while the masked construction lags be- hind. In addition, we show generated samples in Figure 3 and Figure 8, along with a visualization of the metric probability path Figure 4 and ablation studies on NFE and CFG scale in Appendix G.
CONCLUSION
We have opened up the design space of discrete Flow Matching models based on the the continuous- time Markov chain generative process. particular, we propose a kinetic optimal point of view for constructing velocities given prescribed probability paths. This leads to, for the first time, allow- ing arbitrary probability paths to be used. Furthermore, we justify mixture paths with particular schedulers being kinetic optimal solutions, and showcase for the first time, competitive results for non-mask source distributions. Our method naturally encapsulates existing approaches, and we showcase the flexibility of approach to designing discrete Flow Matching models across multi- ple application domains, ranging from text generation, to materials and image generation, where we significant gains over autoregressive models.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Fijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303. 08797, 2023.
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg, Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34, 2021.
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportu- nities and risks of foundation models. arXiv preprint arXiv:2 072588 2021.
Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. A continuous time framework for discrete denoising models. Advances Neural Information Processing Systems, 18-28266--22799 2022.
Andrew Camphell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design. arXiv preprint urXiv:2402.04499, 2024.
Huiwen Chang, Han Zhang. Lu Jiang, Ce Liu, and William Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 11315 1325, 2022.
Ciprian Chelba, Tomas Mikolov, Mike Schuster. Qi Ge, Thorsten Brants, Phillipp Koehn, and Tony Robinson. One billion word benchmark for measuring progress in statistical language modeling, 2014. URL https: / arxiv. org/ abs /1312. 3005.
Patryk Chrabaszez, Iya Loshchilov, and Frank Hutter. downsampled variant of imagenet as alternative to the cifar datasets. arXiv preprint /707: 08889, 2017.
Xiaoliang Dai, Hou, Chih- Yao Ma, Sam Tsai, Jialiang Wang, Rui Wang, Peizhao Zhang, Simon Vandenhende, Xiaofang Wang, Abhimanyu Dubey, et al. Emu: Enhancing image generation models using photogenic needles in a haystack. arXiv preprint arXiv:2309. 15807, 2023.
Daniel W Davies, Keith Butler, Adam Jackson, Jonathan M Skelton, Kazuki Morita, and Aron Walsh. Smact: Semiconducting materials by analogy and chemical theory. Journal of Open Source Software, 4(38):1361, 2019.
Bowen Deng, Peichen Zhong, KyuJung Jun, Janosh Riebesell, Kevin Han, Christopher J. Bartel, and Gerbrand Ceder. Chgnet as a pretrained universal neural network potential for charge- informed atomistic modelling. Nature Machine Intelligence, pp. 1-11, 2023. doi: 10. 1031/ $42256-003-000116-3
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hier- archical image database. In 2009 IEEE Conference Computer Vision and Pattern Recognition, pp. 248-255, 2009, doi: 10. LLL.L///VVR...00.550000
Jacob Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04880 2018.
Prafulla Dhariwal and Alex Nichol. Diffusion models beat gans on image synthesis. arXiv preprint 00Xi::2005.05233, 2021.
Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid linear units for neural network function approximation in reinforcement learning. Neural networks, 107:3-11, 2018.
Patrick Esser, Robin Rombach, and Bjom Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recogni- fion, pp. 12873-12883 2021.
11

---

Published as a conference paper at ICLR 2025
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederie Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first International Conference Machine Learning, 2024.
Daniel Flam-Shepherr and Alán Aspuru-Guzik. Language models can generate molecules, materi- als, and protein binding sites directly in three dimensions as xyz, cif, and pdb files. arXiv preprint DrXii:2305.05708, 2023.
Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky Q. Chen, Gabriel Synnaeve, Yossi Adi, and Yaron Lipman. Discrete flow matching, 2024. URL http::////x...o///////oo/ 1 15595.
Mathieu Germain, Karol Gregor, lain Murray, and Hugo Larochelle. Made: Masked autoencoder for distribution estimation. In International conference machine learning, pp. 881-889. PMLR, 2015.
Aaron Gokaslan and Vanya Cohen. TpenWebbeextCozpus. 2019.
corpus. httpp//.w..o...aaat//////////ata io/
Nate Gruver, Anuroop Sriram, Andrea Madotto, Andrew Gordon Wilson, C Lawrence Zitnick, and Zachary Ulissi. Fine-tuned language models generate stable inorganic materials as text. arXiv preprint 24XX::24002.0379., 2024.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-688 6851, 2020.
Emiel Hoogeboom, Alexey A Gritsenko, Jasmijn Bastings, Ben Poole, Rianne den Berg, and Tim Salimans. Autoregressive diffusion models. arXiv preprini arXiv:2110.02037, 2021.
Anubhav Jain, Shyue Ping Ong, Geoffroy Hautier, Wei Chen, William Davidson Richards, Stephen Dacek, Shreyas Cholia, Dan Gunter, David Skinner, Gerbrand Ceder, et al. Commentary: The ma- terials project: A materials genome approach to accelerating materials innovation. APL materials, 1(1), 2013.
Rui Jiao, Wenbing Huang, Peijia Lin, Jiaqi Han, Pin Chen, Yutong Lu, and Yang Liu. Crystal structure prediction by joint equivariant diffusion. arXiv preprint arXiv:2309.0444 2023.
Walter Kohn and Lu Jeu Sham. Self-consistent equations including exchange and correlation effects. Physical review, 140(4A):A1133, 1965.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Matthew Le, Apoorv Vyas, Bowen Shi, Brian Karrer, Leda Sari, Rashel Moritz, Mary Williamson, Vimal Manohar, Yossi Adi, Jay Mahadeokar, et al. Voicebox: Text-guided multilingual universal speech generation at scale. Advances neural information processing systems, 36, 2024.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0300 2022.
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion language modeling by estimating the ratios of the data distribution. International Conference Machine Learning, 2024.
Anton Lozhkov, Loubna Ben Allal, Leandro von Werra, and Thomas Wolf. Fineweb-edu, May 2024.
URL https::////// htttps://hugginaee http://huggnnggggcccccc co/ co/dataset s
Mitchell P. Marcus, Beatrice Santorini, and Mary Ann Marcinkiewicz. Building large annotated corpus of English: The Penn Treebank. Computational Linguistics, 19(2):313-335 1993. URL
https: /aps:///acaathhee wow.laanthoooooooooooooo. 1000y.300/933320....
12

---

Published as a conference paper at ICLR 2025
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016. URL https: arxiv. orgg aa////0s1110000 07843.
Benjamin Kurt Miller, Ricky T. Q. Chen, Anuroop Sriram, and Brandon M Wood. FlowMM: Gen- erating materials with riemannian flow matching. In Forty-first International Conference Ma- chine Learning, 2024. URL https: / cppnnevieeeee net/ forumoormmm.. -i
Hunter Nisonoff, Junhao Xiong, Stephan Allenspach, and Jennifer Listgarten. Unlocking guidance for discrete space diffusion and flow models. arXiv preprint arXiv:2406.01157 1572, 2024.
Shyue Ping Ong, William Davidson Richards, Anubhav Jain, Geoffroy Hautier, Michael Kocher, Shreyas Cholia, Dan Gunter, Vincent L Chevrier, Kristin A Persson, and Gerbrand Ceder. Python materials genomics (pymatgen): A robust, open-source python library for materials analysis. Computational Materials Science, 68:314-319, 2013.
Arnaud Pannatier, Evann Courdier, and François Fleuret. a-GPTs: new approach to autoregressive models, 2024.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring broad discourse context, 2016. URL https: aarxivvvo..///////////aa org/ abs/1606 06031.
William Peebles and Saining Xic. Scalable diffusion models with transformers. arXiv preprint arXiv:2212.09748, 2022.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4195-4205, 2023.
2023.
Gabriel Peyré, Marco Cuturi, al. Computational optimal transport: With applications data science. Foundations and Trends@ in Machine Learning, (15-66)35556077, 2019
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models unsupervised multitask learners. OpenAl blog. 1(8):9, 2019.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer: High- resolution image synthesis with latent diffusion models. Proceedings of the IEEE/CVF confer- ence on computer vision and pultern recognition, pp. 10684-10695, 2022.
Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David J Fleet, and Mohammad Norouzi. Image super-resolution via iterative refinement. JEEE transactions on pattern anal- ysis and machine intelligence, 4514444773 2022.
Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin T Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. arXiv preprint arXiv:2406.005524 2024.
Neta Shaul, Ricky T. Q. Chen, Maximilian Nickel, Matthew Le, and Yaron Lipman. On kinetic op- timal probability paths for generative models. International Conference on Machine Learning, pp. 30883-30907. PMLR, 2023.
Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, and Michalis K Titsias. Simplified and gener- alized masked diffusion for discrete data. arXiv preprint S.Kiv:224060432222 2024.
Uriel Singer, Adam Polyak, Thomas Hayes, Xi Yin, Jie An, Songyang Zhang, Qiyuan Hu, Harry Yang, Oron Ashual, Oran Gafni, et al. Make-a-video: Text-to-video generation without text-video data. arXiv preprint arXiv:2209.14770
2022.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learn- ing, pp. 2265- PMLR, 2015.
Yang Song, Jascha Sohl-Dicksed Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. International Conference Learning Representations, 2021.
13

---

Published as a conference paper at ICLR 2025
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. Neurocomputing, 568:127063, 2024.
Peize Sun, Yi Jiang, Shoufa Chen, Shilong Zhang, Bingyue Peng, Ping Luo, and Zehuan Yuan. Autoregressive model beats diffusion: Llama for scalable image generation. arXiv preprint arXiv:2406.06523, 2024.
Chameleon Team. Chameleon: Mixed-modal early-fusion foundation models. arXiv preprint h0t.:2405.09818. 2024.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2307.092288 2023.
Aaron Van den Oord, Nal Kalchbrenner, Lasse Espeholt, Oriol Vinyals, Alex Graves, et al. Con- ditional image generation with pixelenn decoders. Advances in neural information processing systems, 29, 2016.
Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
Nisheeth Vishnoi. Lx=b. laplacian solvers and their algorithmic applications. Foundations and Trends in Theoretical Computer Science. 8, 01 2012. doi: 10. 10.1561/04000000544
Apoorv Vyas, Bowen Shi, Matthew Le, Andros Tjandra, Yi-Chiao Wu, Baishan Guo, Jiemin Zhang, Xinyue Zhang, Robert Adkins, William Ngan, et al. Audiobox: Unified audio generation with natural language prompts. arXiv preprint arXiv:2312.15821, 2023.
Logan Ward, Ankit Agrawal, Alok Choudhary, and Christopher Wolverton. A general-purpose machine learning framework for predicting properties of inorganic materials. npj Computational Materials, 2(1):1-7, 2016.
Tian Xie, Xiang Fu, Octavian- Eugen Ganea, Regina Barzilay, and Tommi S Jaakkola. Crystal diffusion variational autoencoder for periodic material generation. In International Conference on Learning Representations, 2021.
Zhilin Yang. Xlnet: Generalized autoregressive pretraining for language understanding. arXiv
preprint arXiv: 190.008.77. 2019.
Yizhe Zhang, Jiatao Gu, Zhuofeng Wu, Shuangfei Zhai, Joshua Susskind, and Navdeep Jaitly. Plan- ner: generating diversified paragraph via latent language diffusion model. Advances in Neural Information Processing Systems, 36, 2024.
Kaiwen Zheng, Yongxin Chen, Hanzi Mao, Ming- Yu Liu, Jun Zhu, and Qinsheng Zhang Masked diffusion models are secretly time-agnostic masked models and exploit inaccurate categorical sampling. arXiv preprint arXiv:2409.02908, 2024.
Chunting Zhou, Lili Yu, Arun Babu, Kushal Tirumala, Michihiro Yasunaga, Leonid Shamis, Jacob Kahn, Xuezhe Ma, Luke Zettlemoyer, and Omer Levy. Transfusion: Predict the next token and diffuse images with one multi-modal model. arXiv preprint arXiv:2408. 1039, 2024.
Nils ER Zimmermann and Anubhav Jain. Local structure order parameters and site fingerprints for quantification of coordination environment and crystal structure similarity. RSC advances, 10 (10):6063-6081, 2020.
Alon Ziv, Itai Gat, Gael Le Lan, Tal Remez, Felix Kreuk, Alexandre Defossez, Jade Copet, Gabriel Synnaeve, and Yossi Adi. Masked audio generation using a single non-autoregressive transformer. In ICLR, 2024.
14

---

Published as a conference paper at ICLR 2025
A ALWAYS-VALID SAMPLING SCHEME
The second step of the sampling scheme defined in Section requires the condition / < (=
to be a valid PMF, allowing only small step sizes. To avoid this constraint on ht we use an alternative first-order sampling scheme. We replace step 2 from Section 3 with
2) Sample X h+h -=Mxxiixx)
() (1 -- ol (IXI) (1 0x;()))
where x(x;)xi)= = u (xI,
Interpreting this expression, e (x is the probability the state does not change. If we do not change state, then we sample from (x;(.). If we do change state, then we sample from
xx(xx((()) (1 - -xx())), which a normalized distribution over states not equal to XZ- X(XIXX
This is still first-order sampling scheme, i.e. it is o(h) error from P(X + Xt). However, unlike the simple Euler procedure, this alternative is always a valid PMF for any step size h.
Algorithm 1 Euler Solver
Require: model 8, To. h t -+0
X, 35
while < 1 do
for i=O.....D {in parallel) do
Xi 0, (+)x))
A xXIx))
jump
U(0,1] 1]
if 7 jump
I 1 -h then
X ~ x.(xxxixi) (1 xx)
end if end for
+ h
end while
return X,
B SYMMETRIZED KINETIC OPTIMIZATION PROBLEM
What makes the optimization of the kinetic energy problem (19) hard for general weighting w; is the non-negative flux constraint (19c) which is non-lincar. Nevertheless, for symmetric weighting defined as
wr 2)
wx(2,z) (x
V2,2ET,
(39)
(z(2)
we are able to overcome this constraint and obtain a close-form solution of the kinetic optimal flux jt. Specifically, prove that for symmetric weighting (39) and any probability path Pt, the optimization of the flux reduced to the relaxed kinetic optimal problem.
ming. So' Exts we(z.z) pe(x) jo(x, 2)
Kinetic Energy
(40a)
s.t. div,(i))= -p((),
WI TT
Continuity Equation
(40b)
which is similar to the original problem but without the non-negativity constraint (19c).
Lemma B.I. Consider a probability path pe and assume a symmetric weighting we (39), then if j. is a solution to the relaxed kinetic optimal problem (40) then
(x, = (x, je(2,x))
(41)
solution to the kinefic optimal problem (19).
15

---

Published as a conference paper at ICLR 2025
Proof of lemma B.I. Every flux can be written the of symmetric matrix and anti- symmetric matrix,
je 2) = St (x, 2) + jum (x, 2); 5e (x, je(x, + +,((Z,,)) (x, 2) je(2, (e(zza)
(42)
where
(43)
The divergence of any symmetric flux vanishes (15), i.e.,
div,(J; = 0.
(44)
Hence writing the kinetic optimal problem (19) for the flux in terms of the symmetric and anti- symmetric flux components gives,
min yo Stm
So Eafz (x, (px(2)
[1c (I,, 2)² + st 2){] dt
Kinetic Energy
(45a)
s.t. div ()t = (x)),
E
Continuity Equation
(45b)
, r 2) b (xr, =)
Ve F
Non-negative flux
(45c)
Notice that for the symmetric term "m we are left with optimization of a quadratic function that is constrained only by equation 45c which bounds it below by a non-negative value. Hence for any Jt the optimal symmetric term is
Jt (x, (x, = 1 (c,
VT,
(46)
So the solution for kinetic optimal problem is
(x, 2) = antii.* =x) + Jt
(xr, anti (xr, =)
St anti, (2,
(47)
where Jr... (z. 2) is an optimal anti ayti-symmetric term. Finally, notice that the optimization prob- lem for the anti-symmetric term ie., equations 45a and 45b, is exactly our relaxed kinetic optimal problem (40).
The relaxed kinetic optimal problem (40) is actually easy solver and can be reduced linear system of equations, leading to our main proposition.
Proposition B.2 (Kinetic-optimal relaxation). Consider a probability path Pt > and assume symmetric weighting we (39). Let ft be a solution to equation 21, which is unique to a constant. Then, jt equation 22 is the unique solution to the Kinetic Optimality problem equation 19.
Proof of proposition B.2. By lemma B.1 the kinetic optimal flux is
Ji (x. 2) = [Jrelax (x. 2) - jr (z, x)]
(48)
where jiclax is the solution to the relaxed kinetic optimal problem (40). Since the continuity equation constraint (40b) is linear, we can solve the relaxed problem using Lagrange multipliers.
Denoting the Lagrange multipliers by (x)), T, the solution to the relaxed kinetic optimal problem (40) is
Pr(2) (2,(2) xx))),
gelax (z,
(49)
2w, 2)
where the Lagrange multipliers are determined by the linear system of equations,
P2(2)
(Ar(z) - Ax())) (x)) Va T..
(50)
W/ (x,
VET
Denoting the solution to this linear system of equation fx(),, the optimal flux is,
p((2)
P((x)
(=,,)) = (x, z)
(f((z) - -((()))
(2, (f((z) - -((())
(51)
(z) [f((2) f+(x))++
(52)
w,(r, 2)
where we have used the symmetric weighting assumption in the second equality.
n
16

---

Published as a conference paper at ICLR 2025
Proposition B.3 (Kinetic Optimal paths.). For Pt 0 and the choice of (2. p((z) the
solution to the Kinetic Optimality problem in equation 19 is equivalent problem 30.
Proof. According to equation 22 the optimal flux in this case takes the form (2, (x(x)
(x)
Plugging this in problem 19 we get the energy
((( pi(2)
(x)) 2 (((x) (((z)
pe(xr)
1 E ((()))((())
(i(()
pc(z)
(p(())((2))
2
(e(a)
Pz(z)
[i((a)
Σ pela)) (2((x) pt(x))
Σ
p((x)
(pe(x)
pr(x))
Σ
Vpr(x)
2
Σ
dt vP, (x)
where in the previous to last equality we used the fact that P((x)= = d left with the following optimization problem:
de Pi()) = O. We are
minp.
E. ( (TVP,))))
(53a)
pr(x) >0
(53b)
Ex P.
(53c)
P1 =
(53d)
Making the change of variables (x) (x),, we get the form in equation 30, desired.
The case of y 6..- Given that q 8z, we get that the Kinetic Optimal solution equation 32 takes form a mixture path (equation 2) with the scheduler specified in equation 33. Specifically show that the probability path is
( (x)))= sin? (1 t)Ω
sin Ω p(x) + (1 (1- sin? (1 t)52 sin Ω
8z, (x),
(54)
and the velocity is
2Ω (x 2/211 tan (1 t) (5z, (x) (x) (z()).
(55)
We start by substituting the q 6z, in equation 31,
or (xx sin(1 t)
sin 1s
p(x)
sin Ω 5,, (x).
(56)
sin Ω
17

---

Published as a conference paper at ICLR 2025
Hence the probability path is
sin 1)52
sin 19?
2
pe (1(
Vp(x)=
+ sin S2
Szy (x)
(57)
sin
(1 1152
(2D(zz)) sin(1 - t) sin
sin 19
p(x))+
8eg (x)
(58)
sin 2 52
sin 52
sin 2 5
sin 2 1)02
(2cos - 1)2 1S2
sin? (1) (1- - 1)92)
p(x) +
8zn (x)
sin Ω
sin Ω
sin 5
(59)
sin - 1)22 p(x)+
2200s - 1)] (S (1 -01-090 1192
(60)
sin sin²10 5
sin2 Ω
(sin 52 cos(1 1) cos Sinin1 1)52)²
+
8z, (1)
(61)
sin Ω
sin sin - t) p(x)+
cos? sin (1 - t) sin l2 cos (1 - t)
8Rg (x)
(62)
sin Ω
sin? Ω
sin? Ω
sin - 1)92
-(1 - sin" (2) sin (1 1)22
sin Ω p(x)+
(63)
sin Ω
sin 12(1 sin 2 (1 +)52) 8z, (x)
(64)
sin Ω
sin 2 (1- t1-
p(x)+ (1-1111111 sin 2 (1 +).2
82, (z), sin² 8
(65)
sin 2
where in the third equality we used Ω = arccos v //p(2) ). Substituting
(x1) = sin ? (1 t)s2
(66)
sin 22
in equation 75 yields the desired velocity as in equation 55.
C CLOSED-FORM KINETIC OPTIMAL VELOCITIES
C.I KINETIC OPTIMAL VELOCITIES FOR MIXTURE PATHI
We examine the velocities from the kinetic optimal fluxes given in (24) and (26) under mixture paths (2). We show (24) and (26) produce the same velocity for the uniform mixture for which p(x) = 1/[T], and different velocities for non-uniform mixtures. We also demonstrate our kinetic optimal velocity from (26) the velocity proposed by Gat al. (2024) for any mixture path.
Positive mixture paths using (24): For (24), we only consider mixture paths for which Po(x)x;) > 0 for all E T, including uniform p(æ) 1/|T| as special case. For these mix- ture paths, where we recall that x * z, we have
2 (x, 2|œ1) (a)p((x((1)) - @cpe(=|or))++
(x1)) 18z, (x) 8z, (2) p(z) - p(x)]+
(67)
(T(p(()(11)
We now examine the uniform and arbitrary p(x) > O cases separately.
Uniform mixture using (24): For uniform, have for
u (x, a|æ1)= (x(x1) [8 (x) 8z, (2) + p(2) - -(t))-(())]_
IT (2(a))
(x1)) ) (o,, (x) 82, (=)]+
(68)
i (2)(())
18

---

Published as a conference paper at ICLR 2025
This is only positive if x 21 and 21. In which case, we have
u1 2 # x1|21) = (x(x11) [1 0
1 - (x(
(x1)
(69)
- (-re(xz)
So in total we have
(x1 ug (1,2/21) 1 Kig (311
(82s (x) 02()))
(70)
Arbitrary p(r) > 0 using (24): For a non- non-uniform positive p(x) we do not arrive at the same velocity as the uniform mixture. Consider # # I1, and 2 of I1, then
" + 2 = + x()t) (e(x1) [p(2) p(x)] (=((1))))
(71)
This is not zero if p(2) p(z) for any pair of and æ, proving this is different velocity in general.
Arbitrary mixture paths using (26): and (=(2(21) > 0, we have
Substituting in the mixture path, where recall that
ui (v, 2|211) = pr(2(a1)
- (x
(((()(21) pe(2)a)) (6x, (=) - p(=))
= (
Oen (x) - p(x)
(72)
We consider several cases. First, # 21 and z z1, then the term in brackets is negative and hence 2r = O. Second if x # 21 and 2 * 21, we have
-p(x) pe(2))1))()))
ui x1, # x (m xx))
Pc (=(2211
-((x) - (x(2x))))))) (p))))()))))))) Ky (p(())
= (x1)
0.
(73)
Our final case, and + æı, gives
uT 2 # a11111) = ( 1 - p((p)) + ((1 - (x(x11))((a) + Ne (x))) p(2)
(1 Ky (x))))(2)
(x11
[(1 - k((=)((x)))(111111 p((1))(())))))))))))))(()))))p + (1 me (z) (x=))=()) + Kr (x1)4
(x(
(x1)
(74)
1- (x
So in total for any mixture path we have
(x(I))
" (I, 211,) 1- (x1))
(5,, (x) (x)
5z(x)),
(75)
recovering the velocity proposed in Gat al. (2024).
C.2 MARGINAL VELOCITY IN CLOSED-FORE FORM FOR MIXTURE PATHS
As shown in Appendix C. the kinetic optimal flux given by (26) results in kinetic optimal velocity (75) for mixture paths. To derive the marginal velocity, we insert (75) into (9) follows
19

---

Published as a conference paper at ICLR 2025
2)) Σ
ziET
E ((()))
Ko (z+z) (5 (xx)) (z'))
EiET
= 1 - (x)) Pijee
(x(xx)
- 6x) Σ (x(x1) Pip(oolle..
1 - Re( 1 p
(76)
riET
C.3 POWER VELOCITY FOR GENERAL PATHS
We begin by defining a single parameter family of kinetic optimal velocities. For every a the flux as in equation 22 for (x) Pl (z) is
1
(x)
(z 2) (x)) - fe(:),,, f((2) EVET pt (s) (z)
(77)
Further simplifying ji (x,,
pf(2)
(x)
(z = pi(x))
EPET Pf - pc(2)
(78)
EXET PJ
An interesting case of the flux above taking the limit a + oo, where
p8 (x)
EXET Pi (s) a-toc 8arg max, (s))
(79)
and the flux
(x, 2) = (2)) max, max (p) - pe (e(2)6arr max, (o.(a))
(80)
Indeed the above flux satisfy the Continuity Equation and the Rate Conditions in Indeed the above flux satisfy the Continuity Equation and the Rate Conditions as in equation 17. Note that it can also be seen that
37 (x,,))
O.
(81)
p=(#)=00
D EVIDENCE LOWER BOUND (ELBO) FOR CTMC
Let 0 <
tx 1 be uniform discretization of the interval [0. 1] with h. 4-1-44 K. Also let (1+1 (x ( = 5, (x= ( + lu,, ='(=f) be the Euler discretization of the variational process, and let P=+1)1(='1==") = 8,, (x*) + hui 2) be the Euler discretization of the learned process, with both starting at the same source distribution qo (o<((z)) p(x'). We also assume the model p(af (xt.x) = dzg (a1). The discrete-time ELBO then
log (o 2 Exo. mox [log p(11 (oa + log Po:k (Fu:k) log qu:k (#a:k [x1)]
(82)
E (11m))
log 8.gg (xf)
DEL (gk+11j (x) bx of ll (1
(xxx)
k=
(83)
Σ DKLlaale (px()<")))))))) ppt))l((x'))
(84)
20

---

Published as a conference paper at ICLR 2025
Each term in the summation:
DxL (9a (1' of (1
(85)
Σ (011n(1112, log (=11m
(86)
(2+1)((0(12))
Σ (x))(x") + hu 21121) log (20)x)) + hug 21121)
(87)
)11((t)) (x),
huf (2', 2|01)
[ua(2', =(aa) log x1(1),)'|=))
[1+ huz(z' 2 [-)) log
h
(88)
+ heul (z), x)
2((x'
Taylor series expansion around h = 0:
log(1 + hu,) = hu, + o(h)
(89)
So we can simplify
DeL ( & (x' [z, M1 (z))
(90)
[ut(x' aaaa) log uf(x', 2 2²101)
= [1 + ++ = (ali) (Nuf( =111) huf((" , +
E
o(h)
uf(x', 2)
x xt
(91)
- (uilet =((('12) + Σ
"4(a", 21122) ui(x', 2)
[ut(a*, of 21a11)
log
+ o(h)
(92)
Taking limit as K oo, hence + O, and asserting that g(xx/=²) dzz (akk) in this continuous-time limit, we obtain the ELBO:
log log))(z)) 2
(93)
, at (ot) log ut(at, vz))
Exempeliles) 1=1
Σ
ux(r),0))) ut (x), r;/=1) +
Σ 24
us (a",
dt
un (z xz (2))
(94)
D.1 ELBO FOR MASKED MODELS
The masked probability path is in equation with source distribution (x*) = 8m (2f). Assuming
the model is such that (=(r(=))r) = døg (zf) if xt unmasked (i.e. x = #f), our ELBO as in equation 37 further simplifies to
1o1) 2 /. (xsppe())))
E 6m dm((t)
-
Σ 1 Nie ( pije Pije(1"(11))
Re((")
(95)
Me(xt) 1- m-((af
(1+ log pije ( dr.
(96)
This simplified expression recovers the ELBO for masked mixture path as proposed by Shi et al. (2024).
21

---

Published as a conference paper at ICLR 2025
E EXPERIMENTAL DETAILS
E.1 TEXT GENERATION
Data. Our model are on trained OpenWebText (Gokaslan & Cohen, 2019) and Webe Edu (Lozhkov et al., 2024). For evaluation we use the split of five dataset Radford et al. (2019): WikiText-103, WikiTexx-2 Merity et al. (2016), LAMBADA Paperno et al. (2016), PrnnTreebank (PTB) Marcus al. (1993), One Billion Words (IBW) Chelba al. (2014). Additionally, we ex- tract 512 samples of length 1024 tokens of GPT2 Tokenizer from FineWeb-FFl Web-Edu, we do not see on training (our models do not complete an epoch this dataset.
Models. All of our text generation models uses DiT transformers architecture Peebles & Xie (2022) with 12 layers, 12 attention heads, and hidden dimension of 768 (150m parameters). For optimization we use constant learning rate of 34-4 with 2500 warmup steps, Adam optimizer with 81 = 0.9 and B2 = 0.999, and weight decay of 0.03. We also use dropout rate of 0.02, and we train for 2004 iterations with batch size 512.
ELBO for training. All text model are trained using our ELBO for mixture path as in equation 37. To avoid exploding terms in the loss, sample in [0, 1 - 1e 1e=
ELBO for evaluation. We want to evaluate the ELBO as in equation 37 for trained models with the mixture path as in equation We note that each choice of scheduler Ky (xf)) will results in different conditional probability path and hence different different ELBO. However for every token independent scheduler Ky (x()) = Ke we can change the integration variable from t to k,
l)g)((x1) 2 (/" dt dtE (- (mx)
Re (x(f)
Σ, ke(D²)
Σ
(e(zz) pi (x/(xx))
ne( Pilc(0
(97)
yi
+ (1 - 829 (x+)) k((x))
+ log logpl 1og (x)1)2))
(98)
I-me(ar)
{ dt dt
h4
Σ
(99)
Kit
+ (1 (x/)) (1 (1 log p (
(100)
dR
E Expxpe (11x1)
pije. (x/( pug (xi)
(101)
J0
+ (1 5uz (at)) (logpic. (aflzz))
(102)
where ly is the inverse of K4. For token dependent schedulers we only use the Kinetic Optimal scheduler as in equation 33,
sin sin? (=1) ne(xf)=
where ((xf) = arcoos / p(aft)
(103)
sin? ((())
22

---

Published as a conference paper at ICLR 2025
Note that 52 10, |. depending on v pp(z)) take ! = I and evaluate the integral,
log (21) 2 S/" dlB deRerppe m1)
(x
Re((")
Kig (xi²) pi (1/(20)) Σ
(e((") (1
(104)
1=1
+ (1 - dag (ztt)
{x(wa) [
(1 log (1f (21(111))
(105)
J. dx (2= 1 4)
N
ke(aa))
p. (2r
(106)
Re, (Q 7
ke(ar))
Σ kr(v²)
Kig (br)) p/e (ef (xr)+
(107)
Re(a)) [x1] (1 log pig.i (a (aflzz))
+ (1 - dmg (zt))
(108)
where Kr (Ω 4) is the Kinetic Optimal scheduler with Ω T,. and be is the inverse of Ne (1 = 7). Now that we have a more fair estimator all the schedulers we us, for each 21 we
discretize E [0, 1 - le to 1024 using
le
Kij = (j
j = 0, 1023, ~ U[0
(109)
1024
E.2 INORGANIC MATERIAL GENERATION
Material representation. A crystal is represented by a parallelepiped in 3D space with periodic boundary conditions, as previous works (Miller et al., 2024; Xie et al., 2021). The model input is variable-length sequences with length 6 + 4 where is the number of atoms in the unit cell. The first 3 tokens represent the lengths the sides of the parallelepiped, while the next represent the angles between the sides. Every atom is comprised of 4 tokens: discrete atom type and 3 continuous numbers representing the atom position inside the parallelopipe in cartesian coordinates. The coordinates are represented relative the side lengths of the parallelopiped, and are therefore restricted to the interval [0, 1] (known as fractional coordinates).
While lengths, angles, and fractional coordinates all continuous quantities, we discretize them uniformly to generate tokens, following the same tokenization method from Gruver et al. (2024). lengths (in Å) are truncated to one decimal place, angles (in degrees) are represented as integers, and fractional coordinates are truncated to two decimal places. The token set for these attributes can be created by the following python code:
tokens_ lens [f" .1E)" for i in range (500)
tokens_angles = [str (x) for x in range (180)
tokens_frac [f"0. (1:02d) for i in range (1 10000 + ["1.00"]
Tokens for atoms taken from Pymatgen (Ong al., 2013) like
from pymatgen. core. pabldictaaiil import Element
Lokens_atom [Element. from_2 (z) name for z in range (1, 95)
The overall vocabulary is composed of all previously mentioned sub-vocabularies, plus special tokens: beggining-of-senteen (BOS), masking, and padding, totalling 500 + 180 + 101 + +4+33- 878.
Model implementation. All of our models listed in Table 2, namely, DFM, Kinetic Optimal DFM (KO-DFM), and Autoregressive (AR), use modified version the Diffusion Transformer (DiT) (Peebles & Xie, 2023) implementation from Lou et al. (2024).
Two sequences that differ only in a permutation of their atoms, along with their fractional coordi- nates, represent the same crystal. For DFM and KO-DFM, we modified DiT account for this
23

---

Published as a conference paper at ICLR 2025
invariance by transforming the input before applying the attention mechanism. We flatten each quadruple of embeddings representing an atom (i.e., atom type plus fractional coordinates) and apply a linear layer with a SiLU (Elfwing et al., 2018) activation to create single representation for the atom. This brings the sequence length from 6 + 4 a to 6 + a. Positional embeddings are then added, where the same positional embedding is added all a output embeddings of the previous step, which establishes the invariance. After the attention mechanism, 4 independent linear layers are applied to each of the outputs, increasing the sequence length from 6 back to 6 + before computing the logits.
For the AR model, we replaced rotary embeddings (Su al., 2024) with sinusoidal positional encod- ings. Note that permutation invariance cannot be enforced in the same way as DFM and KO-DFM, as the model generates tokens auto-regressively. The AR model performs conditional generation by generating an embedding for the number of atoms E {0, ..... 0max 1}, where Omax 20 for the MP-20 dataset in Table 2. The embedding then passed to the same conditioning mechanism (adaLN) present in the original DiT architecture (Peebles & Xie, 2023).
Training and sampling. Hyperparameter values used during training are listed in Table 4. DFM and KO-DFM use the same values.
Param.
Hidden dim.
Attn. Blocks
Attn. Heads
Dropout
Batch Size
Learn, rate
AR
288
16
16
0.1
1024
1e-3
DFM, KO-DFM
256
16
16
0.1
1024
1e-3
Table 4: Hyperparameters used to train the DiT models for material generation.
The hidden dimension of KO-DFM and DFM was lowered to roughly match the same number of parameters as the AR model and FlowMM (Miller et al., 2024) (around 25 million), due to the additional layers required to ensure permutation invariance. Models are trained to predict the next token by minimizing the cross-entropy loss (equation 11).
During sampling. the softmax temperature was fixed to 0.7 for DFM and KO-DFM, and 1.0 for the AR model. Both DFM and KO-DFM have noise distribution equal to a delta function on the all masked sequence (as in Gat et al. (2024)). DFM uses the convex linear scheduler (K₂ = t), while KO-DFM uses the proposed kinetic-optimal scheduler (33).
Our primary metric for material generation is based on thermodynamic sta- bility, key indicator of the synthesizability of materials. Thermodynamic stability is measured by comparing the energy of a material to a database of previously known materials with the same ele- ments. Formally, we define Energy above Hull (Eehall) as the distance in energy landscape between the generated material and a convex hull of energies constructed from these reference database of materials. Stable materials have Ehult convex hull. Following Miller et al. (2024), we define our Stability Rate metric as the percentage of generated materials that are stable, i.e. Ehhul < 0 and n-ary 2. where n-ary of material the number of unique elements in it.
Evaluation metrics
that is the energy of the new material is below the
To compute the energies, we follow the methodology from Miller al. (2024): we first perform structure relaxations using the CHGNet model (Deng et al., 2023), followed by density functional theory (DFT) (Kohn & Sham, 1965) calculations. We generated 10,000 materials to compute the stability rate.
Due to the high computational cost of performing these energy calculations, Xie et al. (2021) pro- posed number of proxy metrics, which we also include for completeness:
1. Structural Validity: Percentage of generated materials where all pairwise interatomic dis- tances are greather than 0.5 A.
2. Compositional Validity: Percentage of generated materials that are determined to be charge- neutral using the SMACT heuristic system Davies al. (2019).
3. Coverage Precision & Recall: Precision and Recall metrics computed by comparing 10000 generated structures to the MP-20 test set. Precision is the percentage of generated struc- tures that are close to some test structure, while recall is the percentage of test structures
24

---

Published as a conference paper at ICLR 2025
which are close to generated structure. Closeness is evaluated using structural and compositional fingerprints (Zimmermann & Jain, 2020; Ward et al., 2016).
4. Wasserstein Distances of Property Distributions: Wasserstein distances between the dis- tribution of computed properties between the test set and the generated materials. We compute these distances for two properties: density (p), and number of unique atoms (Nai)
We emphasize that most of these proxy metrics have become saturated and are not very good at distinguishing state-of-the-art models.
E.3 IMAGE GENERATION CIFARIO
Models All our CIFARIO models use the U-Net architecture as in Dhariwal Nichol (2021), with channels 96,, depth 5, channels multiple [3,4,4], heads channels 64, and attention resolution 16. Additionally, we make two changes to the architecture as done in Gat et al. (2024): (i) We replace the first layer with embedding table of size 256 x 96, and stack the channel features such that the input to the U-Net is of shape 288 x 32 x 32. (ii) We enlarge the size of the final layer output tensor of shape 3 32 32 256. Overall parameters count of 113M. For optimization we use dropout rate of 0.3, and Adam optimizer with 31 = 0.9 and 32 = 0.999, a learning rate of le-4. We trained with an effective batch size pf 512 for approximately 300K iterations.
The conditional path. For our metric induced probability path (27) on pixel space we have natural choice of metric. We embed T = {O,..., } in the interval [-1, C IR using the map 2 emb(x:= 255 x 1 and with the Ip distance,
d(x, 41 (emb(x)) emb(
where Ip a Plyper-paramesd For the Br scheduler we use,
Be = (1.))
where a and e are
We find that best results are achieved with lp = 3, a 5, and = 1. For the other baselines in Figure 2 we follow (Gat et al., 2024).
E.4 IMAGE GENERATION FACE-BLURRED IMAGENET256x255
Our ImageNet256 experiments are conducted on the f'ace-burree variant of the ImageNet bench- mark dataset scaled to 256x256 pixels. We first train a tokenizer model (encoder, quantizer and decoder) that maps the images discrete latent representation and back. Then, we train latent generative model to generate latent representations conditional on the image class.
Tokenizer details. The tokenizer is realized as a VQVAE. Our architecture matches that of VQ- GAN (Esser et al., 2021). It applies 16x downscaling to the image with vocabulary size of 16384. The VQVAE is trained with the VQGAN loss for 40 epochs with a batch size of 128. We optimize using Adam with learning rate le 4, 81 0.9, and 33 0.95. We apply exponential moving average to the VQVAE weights with decay rate of 0.999. After the training is complete, our VQ- VAE model reached rFID value of 2.20, which matches the rFID reported by Sun et (2024) on con-face-blurred ImageNer256.
The baseline
The baseline with masked source distribution uses the cubic scheduler 13
The metric path.
Our metric-induced probability path uses the euclidean distance of the token VQVAE embeddings as the distance function with Ip being a free parameter:
d(z, (1) = Jemb(x) emb(x:))
(110)
Furthermore, we parameterize
25

---

Published as a conference paper at ICLR 2025
t 1 1-1
Be
(111)
with and being free parameters.
These three parameters are costly search, because each configuration requires a separate model to train. We tune these parameters visually by plotting the samples along the conditional path and
looking for configurations that make use of the whole time interval [0,1]. We settled on 0.9, 3 and Ip 4 (see Figure 4)
0.0
0.125 t = 0.25 t 0.375 -= 0.5 == 0.625 t = 0.75 t 0.875 t == 1.0
The conditional path for 0.9, c 3 and Ip 4. This path is advantageous because the the path smoothly interpolates from noise to image while utilizing the whole interval E [0, 1].
Figure 4:
Latent Generative model details. Our generative model uses the Llama architecture that is also used by the LlamaGen model (Sun al., 2024). Our comparisons are done on the Llama-B architec- ture variant with 11 111M parameters. For training hyperparameters, we used the exact configuration proposed in Sun al. (2024): batch size of 256, learning rate of le-4 with 2500 warmup steps, weight decay of 0,05, Adam optimizer with 31 0.9 and 32 = 0.95, gradient norm of 1.0 and class drop probability of 0.1. We used the same ten-crop data augmentation for training that (Sun al., 2024) used.
Following the guidance of (Sun al., 2024), the autoregressive and masked models were trained for 300 epochs. We found that the metric path model benefited from further training, so we trained this variant for 600 epochs.
The DFM models required minor architecture adjustments:
The masked configuration uses non-causal attention.
The metric path configuration uses non-causal attention and also prepend a time em- bedding token (sinusoidal embedding) before the class label token to enable the model to learn the time dependency.
Evaluation. We report the FID of 50,000 generated images w.r.t. the training set. Note that our LlamaGen reproduction obtains a lower FID value then reported Sun et al. (2024) (4.81 vs 5.46). This difference is due to us using the face-blurred variant of ImageNet. While Sun et al. (2024) compares against the pre-computed statistics of non-face-blurred ImageNet, we compile the statistics of face-blurred ImageNet, including training data augmentations.
Ablations.
We show ablations for CFG scale (Table 5) and NFE (Table 6).
F RELATION TO SEDD (Lou ET AL., 2024)
In this section we explain the relation between our method and SEDD (Lou et al., 2024). focus on three main points:
26

---

Published as a conference paper at ICLR 2025
LlamaCien, FID:
5.77
5.18
4.91
4.81
5,02
5.26
5.63
6.11
6.56
7.12
7.70
DEM musked path, NFE=100. FID:
17.78
12.85
9.53
745
6.28
5.78
5.76
6.03
6.56
7.20
7.97
DFM metric puih, NFE=100 FID:
8.58
5.99
4,87
4.50
4.82
5.47
Table 5:
Ablation of CFG scale for LeamaGen and DFM models. The missing cells were not evalu-
ated because they are far from the optima.
NFE
50
100
150
200
250
DFM masked path, CFG=1.6, FID:
5.73
5.72
5.74
5.71
5.82
DFM metric path, CFG=1.3, FID:
4.78
4.50
4.69
4.87
4.98
Table 6: Ablation of NFE for the DFM models.
1. Generality of probability paths. SEDD starting point is a diffusion matrix (i', and re- quires closed-form conditional probability pa (x*((z)) (zf) path solving the Kolmogorov equa- tion (linear ODE) with this rate matrix. This entails solving a (general) [T] dimensional ODE which can be hard to do in closed form. Therefore SEDD resorts to rates of the form Qi(x', z²) "'Q'(x' 2²). In contrast, our method offers a closed form generating rates (velocity) for every conditional probability path, see equation 16 and 26.
2. Score-velocity conversion. The concrete score function is particular way to parameterize probability velocity which is given by
(x 2) Qi 2 (=)s (2)...
(112)
3. Loss. The training loss SEDD can seen as instance our ELBO (36) when using the concrete score parameterization.
Probability velocity vs. concrete score. Using our notation, the noising process of SEDD taking distribution P1 at time 1, some simple distribution Po at time transition probability
defined by the
P(X, = x) + +QQ(2, + o(h))
(113)
where Qe E T S is called diffusion matrix and it satisfies the rate conditions as in equation 5. The reverse process, taking the distribution po at time = to the distribution pl at given by the diffusion matrix,
(z(x)
Qr(2, 2) Q1(2, x)
(114)
p+(2)
where the marginal is determined by the noising process (113) and the distribution at the boundary / 1. The transition probability the reverse process
P(X++h
X, =) = (x(x) hQ, + o(h).
(115)
To make the process tractable, the noising diffusion matrix that chosen only allows transitions from states E S to e S that differ by single token as in equation 6,
Qx(x,2)= Of II (x²), j43
(116)
where Q' R| and satisfy the rate conditions (5). In this case the diffusion matrix of the reverse process is,
Qc(2,2) = Q2(2, x) pr(x)
(117)
(z(2)
D
Sas (2²) Pc(x)
Q2(2" x")
II
(118)
Pz(2)
D Σ
Qi(2 ) 2)
II
82s (x)),
(119)
jzi
27

---

Published as a conference paper at ICLR 2025
where 2x)),))
2) is called the concrete score function and it is defined
<((x', Pr(= ;..., 2'_1,==________ x² Pe ,.... 2¹ 2-, 2 D)
2²+1 pD)
(120)
Considering the boundary condition time 1 to be the data distribution, Pl = q, since in our notation the velocity of reverse process is ue (z, = (z, 2) we have that (comparing equation 119 and equation 7)
u 2)) ='s="
(121)
In the next paragraph we show that for the boundary condition pt S2., the time marginal of the noising process is factorized.
(=(x = II px (x)).
(122)
i=1
In this case the conversion from concrete score to the probability velocity is,
(1* [x1x) Q2(2", x Pr (x)(()))
(123)
(=())
Considering equation 9, we see that the relation between the concrete score and the probability velocity in equation 121 holds only if (x', 2²) is independent of æı.
The conditional probability path. The conditional probability path is the marginal of the nois- ing process when taking P1 = 8.p- Hence, the relation between the diffusion matrix Qt and the conditional probability path is given by an ODE,
d di (1-((1(21)
(1- 2)
(124)
D
Σ 3
x²)
II م2y (x' (cc))pi((())) jØ
(125)
ZES ==1
One can check that indeed the factorized conditional probability path, ie, pe(x(x1) I Pr (~*lzf) is the (unique) solution to the above ODE in case that
de -h (x'())=
- "''_((",'')mn1eee 2EET
(126)
The ODE in equation 126 is still too hard solve in the general case, and some extra assumptions are in order if we hope to solve this equation in analytically. SEDD suggests the standard extra assumption that
Qi 2²) o, (x',
(127)
where [0, + R. and Q' is constant in time. In this case the solution equation 126 is
(=<(<((=))) = exp [[," o.dss
('ddds)) @² @ (a),
(128)
The assumption in (127) significantly restricts the space of conditional probability paths.
In contrast, our point of view is arguably simpler: We start with arbitrary conditional px(x* (xf)
and develop closed-form expression for its generating velocity using equations (16) and (26).
For example, the generating process using metric path in equation 27 should be comparable to the reverse process given by some diffusion matrix,
(= pr(2
Q1(="
Q1(a", ='|xt) (1( =1x=1) ya(a (1)
a1) d(a xxf))
(129)
assuming the diffusion matrix (x 2²) is restricted to Equation (127) we have that
Q' p=(="lat)
3, (0(((=.f)))) xxf) - d(z", xf)]+
(130)
Tt
on leading to contradiction since the L.H.S is constant in time.
28

---

Published as a conference paper at ICLR 2025
SEDD training loss. We derive the ELBO train loss for concrete score function suggested in Lou et al. (2024) from our ELBO (36). To instantiate our ELBO we need to consider two reverse processes. The first correspond to the noising process (113) with the boundary condition p1 6xy.
U (xx,, 1=1)) 02 x²) (e(x)
(131)
(x(2*
The second correspond the noising process (113) with the boundary condition pi q (i.e., data distribution),
u1((",, ,2) = Q'Q xt )s 22)
(132)
Now substitute the velocities in the ELBO (36),
log pı (11) 2 L Fereeeet (on)) Σ Σ
u/ =( x)))
(133)
0'
ui(2", 2xx)
+ nr x xx(=a) log
(134)
u/ll" 2t (x1)
//' Berxpe(-lor))
D
Σ
Σ
0 Q vn)
(v01l0)))
(135)
(x(
1=1 D
Pt (v (xr
(Po(rila))
de
+
log
(136)
Pe (1:/)1)))
1/"0 (psppe(-lan))
Σ Σ
Q (u/e n²) - (o0))
(137)
Pt (v' (v'(
(po(erlaf) (of
+
log (s((w((zz))) - (e((ttp))
de,
pa (x1
(138)
where g(b) = s(log()) - 1).
29

---

Published as a conference paper at ICLR 2025
G ADDITIONAL TABLES AND FIGURES
NFT=-6
NFF==128
Default Velocity
NFT=64
NFE=128
Optimized Velocity
Figure 5: CIFARIO Samples for 64 and 128 NFE, default velocities optimized velocities. The default velocity we use is the velocity resulting from (26). The optimized velocity searches over velocity (35) with varying weights.
(26) (80), and also searches over the probability-preserve For each 8 x table, same seed was used to generate the images.
30

---

Published as a conference paper at ICLR 2025
NFE=256
NPE=512
Default Velocity
NFE-256
NFE-512
Optimized Velocity
Figure 6: CIFARIO Samples for 256 and 512 NFE, default velocities vs. optimized velocities. The default velocity (26) or (80), and also searches over the probability-preserinng velocity (35) with varying weights. For each 8 x 8 table, same seed was used to generate the images.
is the velocity resulting from (26). The optimized velocity searches over
31

---

Published as a conference paper at ICLR 2025
NFT=64
NFE=128
Default Velocity
NFE=256
NPE=512
Optimized Velocity
Figure 7: CIFARIO samples generated from our model using the velocity from Campbell et al. (2024), which does not work for general probability paths such as our metric-induced paths. This is the same Pije model was used generate samples for Figure and Figure 6.
32

---

Published as a conference paper at ICLR 2025
LlamaGen (Sun et al., 2024)
Discrete Flow Matching Mask
Metric
Discrete Flow Matching
Figure 8: Non-curated generated samples for ImageNet256x256.
33


Reviews:
[{"review_id": "vy0z4sthbw", "reviewer": "Reviewer_zfHf", "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.\nKey contributions:\nA decomposition of path velocities, making generation simpler.\nClosed-form velocities that allow any probability path, unifying and expanding prior methods.\nA kinetic energy-based optimization of paths, leading to new, source-dependent schedulers.\nA general ELBO that improves performance on mixture paths, surpassing masked models.", "questions": "Could you provide more background on the mixture path and explain how it connects to the claim in Section 4.2?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "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?\n2. 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ojYp1V3XHq", "reviewer": "Reviewer_ZXqN", "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.", "questions": "Please see the weaknesses section above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed method, Discrete Flow Matching (DFM), is theoretically well-justified through the use of kinetic optimal velocities and probability paths.\n\n2. 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.\n\n3. 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\"\n\n2. 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.\n\n3. What is $\\bar{i}$ in Line 117? Define this when it is first introduced in the paper.\n\n4. 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.\n\n5. 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].\n\n6. Typo in Line 511: “can be provide” --> \"can provide\"\n\n7. Typo in Line 528: “We additional show generated” --> “We additionally show generated”\n\nOverall, 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NeU3OnzO6b", "reviewer": "Reviewer_S7CB", "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.", "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?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 4, "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. \n\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JnGBi2tz3y", "reviewer": "Reviewer_kVBR", "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.", "questions": "Is the use of the concept *flux* in this context novel, or already described in literature? If so, please provide a citation.\n\n## Minor remarks\n- 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.\n- line 174 'abusing a bit *the* previous ..'\n- lines 202 - 206: it would be nice if the black triangles where vertically aligned", "limitations": "", "rating": 10, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 4, "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.\n\n2. The authors propose holistic, novel methodologies for flow matching in discrete spaces.\n\n3. The methods are experimentally validated on a broad set of modalities.", "weaknesses": "None.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NlvjJQi5pH", "reviewer": "Reviewer_vZy2", "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.", "questions": "Please see the weakness part.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\n2. The proposed method is general and can recover previously proposed methods, such as Gat et al. (2024).\n\n3. 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:\n\n1. 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?\n\n2. Why did the authors train the VQ-VAE for ImageNet rather than using off-the-shelf pretrained weights?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)