Masked Diffusion Models as Energy Minimization
Authors: Sitong Chen
Shen Nie
Jiacheng Sun
Zijin Feng
Zhenguo Li
Ji-Rong Wen
Chongxuan Li
Venue: NeurIPS 2025 poster ()
Keywords: Masked Diffusion Models
Optimal Transport
Energy Minimization
Kinetic Energy
Few-Step Sampling
Primary area: deep_learning
TL;DR: We present a theoretical framework that interprets masked diffusion models (MDMs) as solutions to energy minimization problems and an efficient post-training schedule tuning method without model modification.

Abstract: We present a systematic theoretical framework that interprets masked diffusion models (MDMs) as solutions to energy minimization problems in discrete optimal transport. Specifically, we prove that three distinct energy formulations—kinetic, conditional kinetic, and geodesic energy—are mathematically equivalent under the structure of MDMs, and that MDMs minimize all three when the mask schedule satisfies a closed-form optimality condition. This unification not only clarifies the theoretical foundations of MDMs, but also motivates practical improvements in sampling. By parameterizing interpolation schedules via Beta distributions, we reduce the schedule design space to a tractable 2D search, enabling efficient post-training tuning without model modification. Experiments on synthetic and real-world benchmarks demonstrate that our energy-inspired schedules outperform hand-crafted baselines, particularly in low-step sampling settings.

Masked Diffusion Models as Energy Minimization
Sitong Chen 1.2.3 Shen Nie .Jacheng Sun',
Zijin Feng', Zhenguo Li Ji-Rong W.. Chongxuan L2
Gaoling School of Artificial Intelligence, Renmin University of China
2 Beijing Key Laboratory of Research on Large Models and Intelligent Governance
3 Engineering Research Center Wext-Generation Intelligent Search and Recommendation, MOE Huawei Noah's Ark Lab
{chensitong0809, www.eeieeennneia chongxuan Srrtuuannn@uruuu Cruc. edu. cn; (sunjiachengl, fini Li Zhenguo) Chuawei... com
Abstract
We present systematic theoretical framework that interprets masked diffusion models (MDMs) as solutions to energy minimization problems in discrete opti- mal transport. Specifically, netic, conditional kinetic, and geodesic energyare mathematically equivalent un- der the structure of MDMs, and that MDMs minimize all three when the mask schedule satisfies closed-form optimality condition. This unification not only clarifies the theoretical foundations of MDMs, but also motivates practical im- provements in sampling. By parameterizing interpolation schedules via Beta dis- tributions, reduce the schedule design space to tractable 2D search, enabling efficient post-training tuning without model modification. Experiments on syn- thetic and real-world benchmarks demonstrate that our energy energy-inspired schedules outperform hand-crafted baselines, particularly in low-step sampling settings.
prove that three distinct energy formulationski-
Introduction
Masked diffusion models (MDMs) HL tive models for discrete data. By reversing stochastic masking process, MDMs iteratively generate sequences through a series of unmasking steps, guided by learned denoising functions. This sim- ple yet flexible architecture has shown promising empirical performance across text generation [E]] protein generation [39, [0], and image generation [115. 23, A7]
6X have emerged as a powerful class of genera-
Despite their empirical success, the underlying principles that govern the sampling efficiency of MDMsparticularly in few-step regimesremain poorly understood. Existing works typically adopt manually designed mask schedules (e.g., linear or sine) without theoretical justification. In contrast, in continuous domains [27 22, [5)] and discrete flow matching [[7]], recent work has drawn deep con- nections between diffusion models and optimal transport, motivating questions such as: Can MDMs be similarly understood through a similar lens? Can we characterize optimal sampling schedules in principled way? How do these schedules relate the geometry of the underlying probability space?
This paper answers these questions by establishing theoretical framework that interprets MDMs as minimizing energy functionals over discrete probability flows (DPFs). We prove that three natu- ral formulations of transport costkinetic energy, conditional kinetic energy, and geodesic energyare equivalent under the structure of MDMs. More importantly, we show that MDMs minimize these en- ergies when the mask schedule is coupled to geometric interpolation schedule (also the weight function in the energeis) % via a simple closed-form relation: af = sin' (En))) This result unifies
Coorreppnndence to Chongxuan Li.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

seemingly disparate formulations and reveals that MDMs not only follow geodesics on the probabil- ity simplex, but also implicitly optimize sampling rate matrices despite its structural constraints.
Building on this insight, we propose an efficient parameterization of schedule functions via the cumulative distribution function (CDF) of Beta distributions. This reparameterization reduces the high-dimensional schedule design problem to 2-dimensional scalar search, enabling task-adaptive tuning with minimal overhead. We validate our theory through extensive experiments on both syn- thetic and large-scale real-world benchmarks, including language, code, and mathematical reasoning tasks. Our results demonstrate that aeery--inspinn schedules outperform commonly used manually designed schedules in few-step sampling settings for certain tasks.
In summary, contributions are:
We establish a theoretical framework that interprets MDM as optimal transport processes, and prove the equivalence of three distinct energy formulations.
We derive a closed-fom condition for energy opttmal mask schedules, showing that MDMs minimize sampling cost under this condition.
We introduce Beta-CDF parameterization that enables efficient and task-adaptive sched- ule tuning in a 2-dimensional space.
We empirically validate our theory across synthetic and real-world benchmarks, demon- strating consistent improvements in few-step sampling performance.
2 Background
2.1 Discrete Probability Flows and Masked Diffusion Models
Discrete Probability Flows (DPFs). DPFs [I6. 14, 32] define continuous time transformations over structured distributions in finite state spaces. In text generation tasks, a state (2), .... z") E DD" typically denotes a token sequence of fixed length n drawn from a vocabulary D of size d. Formally, DPFs introduce a family of parameterized distributions (p+(=)) that interpolate between tractable base distribution P=()) (e.g., uniform distribution) and a target distribution pi (2).
DPF is governed by a time-dependent transition rate matrix (Q.) D which specifies transition probabilities between states. However, the reverse implication does not hold: different rate matrices can induce the same DPF [32]. We will omit the range of the time index E [0, 1] for brevity.
Masked Diffusion Models (MDMs). MDMs [L 8] represent a subclass of DPFs built upon absorb- ing Markov chains. These models capture two temporal processes: a masking process (operating backward in time, / = + 0) and its reverse, the unmasking process (1 = 0 + 1) as follows.
(,(,,))) - ot
{or
z x ( x)),
Masking process:
otherwise
Pc(x)
a (c(a) p=())
Unmasking process: Q((2,z) Q0(2,2 2)
(x # 2), (2)
(2())
0
otherwise
where transitions in the masking process involve single-token masking operations: from x to
where a token replaced by special mask token [M], i... = (z'...., [M]..., zn) + x
(z',..,, x' # [M].... )). The unmasking process reverses this transition.
This absorbing structure yields a closed-for conditional probability flow, realized independent per-token interpolation between the data and mask states:
Peji (21) (1 ar) (a') + an (z'),
where the mask schedule a, [0, 1] is smooth, strictly increasing function satisfying on- = 0 and On = It determines the progression of masking over time. The schedule relates to the rate a, via:
0: = exp (-f,- as ds
We adopt the MDM terminology rather than the velocity functions used in the flow matching literature.

---

as shown in existing work 1R. 38). Existing methods model the likelihood ratio p=(e) in Eq. @aalso pa(2]
known as the concrete score [T9, 26] or its equivalent formulations [B88 [36, H optimize evidence lower bounds (ELBO) of the log-likelihood [4] In particular, the ELBO is invariant to the choice
of Or [12, 38, 36 by change of variables. We provide proofs for Eq. (+) in Appendix C.ll and the
invariance of ELBO in Appendix C2 for completeness.
2.2 Kinetic and Conditional Kinetic Energy
provides principled framework for quantifying the transport cost of probability flows, forming an optimization objective that yield improved sampling trajectories.
Kinetic energy [2.
M
37]
Definition 2.1 (Weighted kinetic energy). Given a weight function "tt and DPF pt governed by rate matrix Qs the weighted kinetic energy is defined as
E 1 Qc(2,z)² 2:2Ez
En (pt, Qii 7x) Er p(())
(5)
where ft denotes the temporal derivative of "t--
The function "i is smooth, strictly increasing schedule satisfying the boundary conditions To 0 and 71 = 1. Different choices of % implement various temporal weighting schemes for the DPF.
This generalizes the unweighted kinetic energy [B7]. where Yt t, and facilitates our theoretical analysis.
The quadratic dependence on the transition rates Qı(z, in Eq. (5) mirrors the classical velocity- squared form of kinetic energy. As Qt controls sampling behavior, minimizing & seeks minimal- cost sampling paths between distributions [24 22 25]. central goal in efficient generative modeling
and few-step sampling [2, IEL [77].
However, direct computation of & is generally intractable, owing to the absence of closed-form solutions for Pt in most practical cases. To address this, we introduce a conditional surrogate.
Definition 2.2 (Weighted conditional kinetic energy). Let Pet1 denote conditional flow governed by the conditional rate matrix Qti1, which satisfies the marginal consistency condition:
Q+(2,x)= Sep(z x|21) Punl)
(6)
21
The weighted conditional kinetic energy with weight function Tt. defined as
(ip.pn,Qup131 = E8, (xppp(()))))
Σ anaeee
(e (11(2,2/21))
(7)
Definition 2 extends Definition 2 to the conditional setting, in line with analogous energy formu-
lations in the continuous domain [27]] While the two energies differ in general, they are equivalent
under the MDM framework (see Sec.
2.3 MDM as Geodesic Curve
To understand the relationship between MDM and kinetic energy from a geometric perspective, draw on the geodesic interpretation [12]]. This view emerges from key geometric embedding: for each token, the conditional distribution Ptj (.. p(i) A D-1 can be mapped isometrically onto the unit sphere via square-root parameterization:
Vt = (VPhp VPRi s (v: -
(8)
Through this embedding, 22
proved geometrically that the per-token MDM conditional flow Eq. (E)
corresponds to geodesic motion on that is, movement along the great circle connecting the
masked initial state yh and the target distribution yı. The interpolation schedule 7tt (a smooth in-
creasing function satisfying to = 0, m = 1) governs the temporal progression, and its derivative ix represents instantaneous velocity. We retain the " notation from Definition DAl as it also serves as

---

Energy equivalence
Theorem 3.1
Theorem 3.2
Ek
E
Eg
Scction 3.1
Theorem 3.6
Energy minimization Scction 3.2
Theorem 3.6
Lemma 3.5
MDMs energy minimizers under Condition 3.4
Figure Illustration the theoretical results of this paper.
the weight function for energy minimization as detailed in Sec. B. A more comprehensive and intu- itive introduction to the geodesic curve and relationship with MDM is provided Appendix B
This geometric framework unifies MDM sampling dynamics with energy minimization principles, since geodesics simultaneously minimize both path length and energy functionals [I]] To formalize this link, we define the weighted geodesic energy, which MDMs minimize under suitable scheduling.
Definition 2.3 (Weighted geodesic energy). Given a weight function " and conditional flow Pij1 governed by conditional rate matrix Qij the weighted geodesic energy is defined as:
4
Ep(pi1171)= E E (mm(a))snn((( RPPep (2* x1)
dan (2'12))
where dan (=1121) = / (2)))) is derived from the embedding in Eq. (8)
The term S² in Eg again reflects the velocity-squared form of kinetic energy, capturing geometric transport costs. Although & lacks explicit dependence on rate matrix, we demonstrate in Sec. that it is equivalent to both & and & under MDM structure.
Finally, since mask schedule a, determines both the rate matrix and the probability flow in MDM, we reparameterize the energy functionals as (ot, ye), Ec(at,7e) and (y ?1s) for the rest of paper.
3 Main Results
3.1 Equivalence of Energies in MDMs
We first formally establish that the three energy functionals defined in Sec. Ethe kinetic energy & the conditional kinetic energy Eg, and the geodesic energy &,are mathematically equivalent under the MDM framework. This equivalence captured by the following theorems.
Theorem 3.1 (Kinetic-conditiona) equivalence in MDMs). For any weight function Yt and MDM with mask schedule 04, the marginal and conditional kinetic energies are proportional:
Ex(ar, 71) CyEa(au, 7t),
(10)
where C1 is scalar depending only on the sequence length and vocabulary size d. As resudt, the two objectives share the same minimizers:
arg min Ee 7t)) = arg min Es(or, 7(). ae
(11)
ar
As discussed in Sec. DL while & and & share similar structure, their inherent differences between marginal and conditional formulations typically lead to divergent values. However, proof in Appendix D.I reveals that by decomposing the concrete score in the rate matrix of MDM (see Eq. ( into temporal components and clean-data conditional probabilities [Σ4] and leveraging the inherent simple closed-form of MDM's conditional rate matrix (characterized in Appendix C.)) equivalence between &g and &, established, even in high-dimensional regimes.
Remarkbly, Theorem 3.1 establishes &, as a theoretically sound surrogate for Eg in MDMs. Cru- cially, while E, suffers from intractability due to the absence of closed-form (z(2), &, remains

---

computationally tractable in MDMs since both the conditional flow and rate matrix admit closed- form expressions. Furthermore, unlike Eg. & decomposes along sequence dimensions (see Ap- pendix C.)), enabling per-token analysis and creating natural bridge to &g energy inherently defined through per-token conditional probability flows. This connection leads to our following theorem.
Theorem 3.2 (Condiiional-gendesicc equivalence in MDMs). For any weight function "t. and MDM with mask schedule Ott, the conditional and geodesic energies are proportional:
Eo(a wr =
Tt),
(12)
where C2 a scalar depending only on the sequence length n.. This implies that they share the same minimizers:
arg min Eo( nh) arg min Ey(or, 7t)- a
(13)
Our proof in Appendix DW reveals that this equivalence originates from the token-wise decompo- sition of & (see Appendix C4I a property inherently enabled by MDM's per-token structure of Pej1 and Qqj This extends the DFM framework [GZi] beyond its original one-dimensional geodesic- kinetic energy correspondence. Morcover, more flexible choices are considered in our settings while energy functionals in [67]] use fixed weight function N t. Therefore, our Theorem [2.2 establishes interpretations applicable to real-world text generation with adaptive weighting schemes.
The insight that the sampling-orienter kinetic energy and the geometrically motivated geodesic ergy constitute two equivalent viewpoints suggests MDMs may simultaneously achieve optimality in both probability flow and matrix characteristics despite its structural constraints, which formally prove in Sec.
[ To exemplify these equivalences, we consider the single-token case (n 1) in Example 3.N. where all three energy functionals collapse an identical closed-form expression. See Appendix D for the proof.
Example 3.3. When n 1, the kinetic, conditional kinelic, and geodesic energies all reduce to:
E(ves 71)) = (,' 7t acl1 - o1)
dt.
(14)
3.2 MDMs with the Optimal Mask Schedule Are Energy Minimizers
Building on the energy equivalence established in Sec. B.I. we now turn to the core optimization question: Can an appropriately chosen MDM schedule simultaneously minimize the primary objec- fivekinetic energy Eg by equivalence, all three energy functionals?
To resolve this question, initiate our analysis from the geodesic perspective. As discussed in Sec. D. o₁ governs the unmasking process in MDMs, while "/ parametrizes the corresponding contin- interpolation along geodesic curves [+2]]] and also serves the weight function in the geodesic energy. This schedule dualityo; for discrete dynamics, "It for geometric flownaturally necessitates an optimal parametric relationship between them, formalized the following condition.
Condition 3.4 (Optimal scheduling condition). We say that the optimal scheduling condition be- tween the mask schedule of and the interpolation schedule "t is satisfied when
a sin ( ("):).
(15)
The monotonic bijection in the condition, namely f [0, 1] [0, 1], x x sin? (5z). creates an one-to-one schedule correspondence between and Ye.
While this relationship was geometrically established in [22)], demonstrating the relationship be- tween the interpolation schedule of the geodesic curve and the mask schedule of MDM that gener- ates the curve, its profound implications for energy minimization remain uncharacterized a gap our subsequent lemma addresses.
Lemma 3.5 (Geodesic energy minimization). Under Condition Tc. the schedule of minimizes the geodesic energy.
Our proof in Appendix D.4 demonstrates through energy- energy-theoretic analysis that Condition not only guarantees generation of minimal-length geodesic paths (thereby providing proof for existing
5

---

7/ CD 0.5))
7, CDF( 0.3))
70
Theoretical Minimum
24
Theoretical Minimum
60
0.8
22
0.8
50
20
40
0.6
18
0.6
Energy
30
16
0.4
14
0.4
20
12
0.2
0.2
10
10
0.25
0.50
0.75
0.25
0.50
0.75
Figure 2: Distinet weight functions "r shape different energy landscapes and consequently yield different optimal mask schedules of. Axes represent the beta-parameterization of ae (see Sec. (.R)) Color intensity indicates energy values from Eq. (EE). Red stars mark the theoretical minima under the optimal schedule condition.
geometric conclusions [+2]] in an alternative perspective) but also formally establishes the attainment of minimal geodesic energy. These results enrich our understanding of MDM optimality. One may further be curious about the practical implications: does discretizing the continuous trajectory into finite-step sampling process affect the attainment of optimality? We address this in Appendix D. demonstrating that the discretized trajectory remains optimal in a well-defined sense. Therefore, our subsequent discussion will continue to focus the theoretical continuous case.
This energy minimization perspective provides a crucial theoretical bridge connecting MDMs' ge- ometric properties with their sampling dynamics. Combining Lemma B.5 with our equivalence theorems Sec. 5.1. we extend the minimization result to & and & energy functionals of pri- mary practical interest due their direct connection MDM sampling efficiency, arriving the following central result under identical optimal scheduling conditions.
Theorem 3.6 (Kinetic energy minimization). Under Condition B44. the MDM schedule az simulta- neously minimizes all three energy functionals.
Theorem HG does not trivially follow from Lemma as & and E, require simultaneous optimiza- tion of probability flows and rate matrices fundamental departure from E,, 's exclusive dependence on the probability flow. Our proof Appendix D.G crucially relies on the fact that although the mask schedule a governs Pt and Qi jointly in MDMs, introducing parametric constraints, the Markov structure of MDM still intrinsically co-optimizes the probability flow and the rate matrix. This re- solves long-stann long-standing conceptual paradox in discrete diffusion: MDM's simple coupled framework in fact preserves optimal transport properties through its intrinsic design Markovian transitions.
Furthermore, Theorem B.G establishes that Condition 3.4 not only dictates optimal probability paths but also governs optimal sampling rates. This theoretical insight directly motivated our energy- inspired schedule tuning method in Sec. [ where select og that minimizes energies given fixed 7r. To empirically validate how distinct weight functions 7t shape different energy landscapes and consequently yield unique optimal mask schedules, we revisit the single-token case (n 1) Example UN Fig. quantitatively demonstrates this relationship by visualizing how varying 7t induces corresponding of schedules that minimize the energy functional defined in Eq. (14).
3.3 Ensrgy-Inspired Fast Samplers
Our energy minimization perspective introduced in Sec. + shows the importance of the weight function 7t- Especially, the term (1a)-1 in the energy functionals (See Definition 2.1 222 and 2.3 plays a central role: it downweights regions rapid temporal change 7h effectively focusing optimization on slower regions. Different choices of thus encode different emphases in the dif- fusion processfor instance, whether to spend more computational budget early (coarse structure) or late (fine detail) in the unmasking trajectory. This observation suggests that task-specific tuning

---

1.0
0.0...09)
1.0
8.03,
8.0.9.033
8(0.9,03) 0.3)
Interpolation Schedule %
0.5.5,0.5)
8(0.5,0.3)
0.8
0.8
BCL. 1)
Mask Schedule a.
B(1
B(2, 2)
B12 2Y
922,55
0.6
BC2 5)
MCS, 2)
M(5,2)
0.2
0.2
0.0
0.0
0.0
0.2
0.6
0.8
1.0
0.2
0.4
0.6
0.8
Time Step
Time Step
Figure 3: Heta-parameterized interpolation schedules and corresponding mask schedules. The left panel demonstrates beta-parameterized interpolation schedule morphologies, while right panel displays corre- sponding optimal oz schedules derived via Condition 6:4
of "t may yield performance improvements, different tasks may benefit from different temporal allocations.
Moreover, since MDM training objectives are invariant to the choice of schedule as (as discussed in Sec. I.)), schedule optimization can be performed post hocafter model trainingwithout requiring re- training. This enables lightweight adaptation of pretrained models to new distributions or generation objectives by simply modifying the sampling schedule.
However, direct optimization over the space of all possible schedules remains intractable due to its nifnite-dimensional nature. Existing approaches therefore rely on a small set of manually designed schedules, such as linear (a,, !) [8, 26, E], sine (o, sin { ()) [E6]), squared sine schedules (a, sin² (īt)) [ To bridge the gap between theoretical flexibility and practical feasibility, propose parameterizing "/r as the cumulative distribution function (CDF) of a beta distribution with two parameters (a, b):
* CDFB(a. (4),
(16)
which, via the condition in Eq. (IL)), yields a corresponding a schedule. This parametric framework generates diverse schedule topologies including convexity and inflection points through just two parameters (see Fig. ( and it is motivated by a simple observation formalized in the following proposition, proved in Appendix D.
Proposition 3.7. Linear and squared sine schedules correspond specific beta parameterizations:
04
=
= CCDF
(17)
= sin
= = #Dt=CDFBu.Dje
(18)
We begin with toy model in Fig. g to illustrate how tuning beta parameters (a, b) affects sampling quality under low-step regimes, showing that different target distributions prefer different schedules, highlighting the need for task-specific adaptation. In Section H. we further demonstrate that, cer- tain beta schedules can outperform standard hand-crafted schedules, especially when the number of sampling steps is limited, on real benchmarks.
Experiments
In this section, we demonstrate that our energy-inspiree task-specific tuning method introduced in Sec. B33 can be used to accelerate MDM sampling in practical applications such as mathematical rea- soning and code generation. Crucially, the invariance of training loss of MDMs to the choice of mask schedules (see Section 2. allows us to efficiently optimize the mask schedule for specific down- stream tasks without the computational burden of end-to-end model retraining. For details on how to identify favorabe schedule by tuning the beta parameters, please refer to Appendix D

---

Probability
0.4
0.2
0.0
Target Distribution
y; ()=CBB(((.,00055)
7 (.=CDF(B(06,00..)
Probability
0.04
0.02
0.00
Target Distribution
0 CDF(B(O.S,0.5) 31
7, CDF(B(0.6,003)) 0.3))
Figure 4: Toy experiments illustrating how different target distributions prefer different schedules. Each panel visualizes the effect of beta parameter tuning sampling quality under limited step budgets by showing target distribution and two distributions sampled by different schedules. More details of this experiment provided Appendix ETR
We evaluate our method using LLaDA 8B [B3], an open-source MDM that achieves performance comparable modern large language models such as LLaMA3 [3T]]. We select representative tasks: MBPP [7]. HumanEval [9], BBH [28], GSM8K [I[O], Hendrycks Math [[ and Minerva
Math [IS]. These benchmarks comprehensively assess the models capabilities general reason-
ing, mathematical problem solving, and code generation. Please refer to Appendix + for more
experimental details.
Fig. E systematically evaluates sampling performance across diverse reasoning benchmarks under varying step budgets. Our analysis reveals that on code generation tasks (MBPP and HumanEval), anta-parameterized schedules match the generation quality of the linear baseline with 2x step re- duction. For the Hendrycks Math mathematical reasoning task, our method achieves performance parity with the linear schedule using 4x fewer steps. There are also benchmarks on which beta- parameterized schedules exhibit comparable yet better performance, such BBH [28] and GSM8K [I(]], and we provide the results on these two benchmarks in Appendix. EX
On benchmarks where petteparamettrrized schedules demonstrate profound empirical advantages over manual baselines, we observe a systematic preference for convex interpolation schedules. As analyzed in Sec. [X this empirical bias suggests that for certain tasks such code generation, allocating computational resources to optimize early-stage sampling dynamics (coarse structure for- mation) may yield greater quality gains compared fine-graiie mme-grinned refinement phases. Further discus- sions on the task-specific schedule preferences are presented in Appendix [
While rigorous characterization of task-schedun correspondences remains open, constituting critical directions for future research, the schedule invariance property facilitates computationally efficient schedule exploration without model retraining. Practitioners thus perform task-specific schedule tuning through our framework, requiring no additional training infrastructure.
Raw data of our experiments corresponding to Fig. E are presented in Appendix E.3, and we provide additional samples Appendix E.S to offer more comprehensive understanding.
Related Work
Mask Diffusion Models. MDMs [L) 8, 54, 94, 66, [X] have established themselves prominent class of generative models for discrete data. While substantial progress been made in under- standing MDM training dynamics through theoretically equivalent objective formulations [12] 38

---

45
35
40
30
sin(e22
sin Lr2
25
(e22
009990330099
CDFG0003,009))
MBPP
20
CDFL8009,603)
20
15
10
5
128
256
16
32
Step
Step
35
sintm ()
sin (e/2
Hendrycks Math
Minerva Math
15
CDPURIO9900 9.00)
sin(s/2
10
sin (m22
0
128
256
16
32
64
128
256
Step
Step
Figure 5: Performance evaluation of energy-optimized schedules on LLaDA 8B [33]. Each panel cor- responds to a distinct benchmark. The x-axis displays sampling steps on logarithmic scale, while the axis quantifies task performance, where higher values denote superior generation quality. Results bench- marks where parameercceee schedules exhibit comparahle yet better performance provided in Ap- pendix. E
36, [ and parameterization strategies 14, 211 211 [9]. the sampling process remains relatively underexplored. Existing efforts primarily concentrate on developing advanced discrete sampling
algorithms, including Tweedie 7-sampling [S & [6)), k-Gillespie methods [+ and higher-orde
solvers [A5]]. along with discretizing time step optimization techniques [GS] and distillation-base
acceleration [33], A6). Notably, the critical mask schedules governing sampling trajectories have
received systematic investigation. Current implementations [8) 26, H. 117, 36]] typically employ manu- ally designed schedules (predominantly linear), even state-of-the thrt--rrttttttttt MDM-based large language models [[3]. This underscores the necessity for principled task-adaptive schedule optimization, which constitutes our primary contribution.
Energy and Geodesic Perspective. The connection between kinetic energy minimization [2. 1 and efficient sampling via optimal transport trajectories [24] 22, 25, [0 has been well-established in continuous settings. However, existing literature predominantly examines continuous diffusion processes, limiting direct applicability discrete domains. Recent geometric analyses [Σ reveal intrinsic links between MDM probability paths and geodesic curves under specific interpolation schedules, motivating our exploration of geodesic perspective netic principles. However, the energy perspectivewhich plays a vital role our optimal transport frameworkis absent in [22]] thus preventing it from establishing the optimality of Condition 3.4 in the context of sampling efficiency. The most relevant work [37]] introduces energy perspectives to Discrete Flow Matching (DFM), distinct discrete probability flow (DPF) variant. Crucially, their framework decouples probability flow and rate matrix optimization MDMs where a₁ jointly governs both components. Despite this architectural constraint, we demon- strate through Theorem 3.6 that MDMs inherently achieve optimal rate selection by leveraging key lemma from [37] (see Appendix I.6).. Our introduced " interpolation schedule also provides novel theoretical insights by establishing that every mask schedule optimizes corresponding en- ergy functional, thereby justifying sask-specific schedule tuning capability absent in [37]] where formulations reduce to the % special case. Furthermore, the training loss invariance mask
bridge between MDMs and ki-
approach incompatible with

---

schedules represents a unique MDM property distinguishing it from DFM and conventional DPF frameworks, enabling exclusive post- schedule optimization within our MDM paradigm.
Conclusion
We present a theoretical framework that establishes MDMs as optimal transport processes mini- mizing three distinct energy formulations, demonstrating that MDMs inherently achieve minimal sampling cost through energy-optimal mask schedules. Building upon this theoretical foundation, we develop a Beta-CDF paramcterization scheme that facilitates efficient task-adaptive schedule optimization. Comprehensive empirical validation across synthetic and real-world benchmarks con- firms our framework's 's effectiveness, showing consistent performance gains in few-step sampling scenarios.
Limitation. While our method enables practical task-specific schedule optimization, the intrinsic relationship between high-dimensional real-world tasks and their optimal schedules remains not fully understood- fundamental challenge requiring further investigation. This limitation highlights promising directions for future research in interpretable schedule-task correlation analysis.
Broader Impact. Our schedule tuning framework could benefit real-world applications by reducing computational costs. Conversely, the acceleration capability may lower synthetic content generation barriers, potentially exacerbating misinformation risks.
Acknowledgement
This work supported by the National Natural Science Foundation of China (No. 92470118); the Beijing Natural Science Foundation (No. L247030); the Beijing Nova Program (No. 20230484416); the Public Computing Cloud of Renmin University of China; the Beijing Major Science and Technol- ogy Project under Contract no. Z251 100008425002; the fund for building world-class universities (disciplines) of Renmin University of China; and the Huawei Research Fund.
References
[1] Michael David Spivak. "A comprehensive introduction differential geometry". In: (No Ti- dee) (1970).
[2] Robert J McCann. "A convexity principle for interacting gases". " In: Advances in mathematics 128.1 (1997), pp. 153-179.
[3] Daniel T Gillespie. "Approximate accelerated stochastic simulation of chemically reacting systems". In: The Journal of chemical physics 115.4 (2001), pp. 1716-1733.
[4] Jascha Sohl-Dickstein et al. "Deep unsupervised learning using nonequilibrium thermody- namics". In: International conference on machine learning. PMLR. 2015, pp. 2256-2265. [5] Nihat Ay et al. Information geometry. Vol. 64. Springer, 2017.
[6] W. E. Applied Stochastic Analysis. Graduate studies in mathematics. American Mathematical Society, 2019. ISBN: 9781470452414, URL: httts:://bone //books gowgle com books id=
2QRHYUEACAAJ,
[7] Jacob Austin al. "Program Synthesis with Large Language Models". In: CoRR 00s/2108.07732 (2021). arXiv: 2108 07732. URL: https: : ://
arxiv. org/ abs 2108
07732
[8] Jacob Austin et al. "Structured denoising diffusion models in discrete state-spaces". In: Ad- vances in Neural Information Processing Systems 34 (2021), pp. 17981-17993.
[9] Mark Chen al. "Evaluating Large Language Models Trained on Code". In: CoRR URL: https : / arxiv org abs 2107
abs/2107.03374 (2021). arXiv: 2107 03374.
03374
[10] Karl Cobbe et al. "Training Verifiers Solve Math Word Problems". In: CoRR abs/21 14114 (2021). arXiv: 2110 14168.
URL: https // arxiv. org abs 2110
14168.
10

---

[1]] Dan Hendrycks et al. "Measuring Mathematical Problem Solving With the MATH Dataset". In: Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks 1. NeurlPS Datasets and Benchmarks 2021, December 2021,
virtual. Ed. by Joaquin Vanschoren and Sai-Kit Yeung. 2021. URL: https 7
datasets - benchmarks - proceedings neurips- cc paper / 2021 / hash 7 htm [12] Diederik Kingma et al. "Variational diffusion models". In: Advances neural information
processing systems 34 (2021), pp. 21696-21707.
[13]
Cédric Villani. Topics in optimal transportation. Vol. 58. American Mathematical Soc., 2021.
[14] Andrew Campbell et al. "A continuous time framework for discrete demoising models". In: Advances in Neural Information Processing Systems 35 (2022), pp. 18266-28279.
[15] Huiwen Chang et al. "Maskgit: Masked generative image transformer". In: Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition. 2022, pp. 11315-
11 11325.
[16] Sander Dieleman et al. "Continuous diffusion for categorical data". In: arXiv preprint 0..22211../000 15089 (2022). arXiv:2211 I.
[17] Xiaochuang Han, Sachin Kumar, and Yulia Tsvetkov. "Ssd-Im: Semi-autoregressive simplex- based diffusion language model for text generation and modular control". In: ar preprint
arXiv:2210. 17432 (2022).
[18) Aitor Lewkowyez et al. "Solving quantitative reasoning problems with language models". In: Advances in Neural Information Processing Systems 35 (2022), pp- 3843-3857.
[19] Chenlin Meng et al. "Concrete score matching: Generalized score matching for discrete data".
In: Advances in Neural Information Processing Systems 35 (2022), pp. 34532-34545.
[20] Haoran Sun et al. "Score-based continuous-time discrete diffusion models". ".d In: arXiv preprint
arXiv:2211. 26750 (2022).
[21] Pan Xie et al. Vector Quantized Diffusion Model with CodeUnet for Text-to-Sign Pose Se- quences Generation. 2022. URL: /ttps://openreeveu. http:ppaarwviiii...eeeooornaaaaa net//oorum?dd-RRl33KKUU
[22] Michael Samuel Albergo and Eric Vanden Eijnden.. "Building Normalizing Flows with Stochastic Interpolants". In: The Fleventh International Conference on Learning Represen- tations. 2023. URL:: https::/openrevii hwtp://eet/eooneeeeeeeeee het.://frrrmmteeeelgelllcic.
[23] Huiwen Chang et al. "Muse: Text-to-mmage generation via masked generative transformers".
In: arXiv preprint ar arXiv:2301.00700 (2023).
[24] Yaron Lipman et al. "Flow Matching for Generative Modeling". In: The Eleventh Interna- tional Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net 2023. URL: https: /ppnnevieeeeeeeeee nnt/ or um hitiiddddppvwcPsS333
[25] Xingchao Liu, Chengyue Gong, and qiang liu. "Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow". In: The Eleventh International Conference on Learn- ing Representations. 2023. URL: https: /ttp///eeevvveeeee
[26] Aaron Lou, Chenlin Meng. and Stefano Ermon. "Discrete Diffusion Language Modeling by Estimating the Ratios of the Data Distribution". In: arXiv preprint arXiv:2310. 16834 (2023).
[27] Neta Shaul et al. "On Kinetic Optimal Probability Paths for Generative Models". In: Inter- national Conference on Machine Learning, ICML. 2023, 23-29 July 2023, Honolulu, Hawaii, USA. Fd. by Andreas Krause et al. Vol. 202. Proceedings of Machine Learning Research. PMLR, 2023, pp. 30883-30907. URL:
https / proceedings mlr press / v202
Shaul25a.. htmil
[28] Mirac Suzgun et al. "Challenging BIG-Bench Tasks and Whether Chain-of-Thought Can Solve Them". In: Findings of the Association for Computational Linguistics: ACL 2023, Toronto, Canada, July 9-14, 2023. Ed. by Anna Rogers, Jordan L. Graber,aber, and Naoaki Okazaki. Association for Computational Linguisties. 2023, pp. 13003-13051 DO1: 18653/V1/ 2023 FINDINGS ACL. 824, 2023. 8C1
URL: https : /7 doi
org/ / 10. 18653 v17
[29] Pengze Zhang et al. "Formulating Discrete Probability Flow Through Optimal Transport". In: Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December
10 16, 2023. Ed. by Alice Oh et al. 2023. URL: http / pappers nips cc paper % 5C files paper 2023 hash aee776/55aar22116522efffff7o111. -ww.ff/d---bbbrttttoooeoaa Abstract Conterence.nie htm

---

[30] Yunfei Chu et al. "Qwen2-audio technical report". In: arXiv preprint arXiv2407.10759 (2024).
[31] Abhimanyu Dubey et al. "The llama herd of models". In: arXiv preprint arXiv:2407.21783
(2024).
[32]
Itai Gat et al. "Discrete Flow Matching". In: CoRR abs/2407.15595 (2024). DOE: 10, 48.505
ARXTV 2407 15595. arXiv: 2407 15593.
URL: https: org /00 488500/a5iiii
22407 15595.
[33] Siqi Kou et al. "Cllms: Consistency large
language
models". In: arXiv preprint
arXiv:2403.0m0 00835 (2024).
[34] Aaron Lou, Chenlin Meng, and Stefano Ermon. "Discrete Diffusion Modeling by Estimat- ing the Ratios of the Data Distribution". In: Forty-first International Conference on Ma- chine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024, URL:
[35] Yong-Hyun Park et al. "Jump Your Steps: Optimizing Sampling Schedule of Discrete Dif- fusion Models". In: CoRR abs/2410.07761 (2024). DOI: 10 48550 ARXIV 2410 07761.
Xiv: 2710. URL: https: / org/ 10. 48550/ arxiv 2410. 077611
[36] Subham S. Sahoo et al. "Simple and Effective Masked Diffusion Language Models". In: Ad- vances in Neural Information Processing Systems 38: Annual Conference on Neural Infor- mation Processing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December 10
15. 2024. Ed. by Amir Globersons et al. 2024. URL.: http: nips cc paper'/ /tt00001cc11572220000000066666.0 -ttp://aaaddaaaaa---a-aaaaa Abstract
5C files / paper / Conference hlmll.
[37] Neta Shaul et al. "Flow Matching with General Discrete Paths: A Kinetic-Optimal Perspec- tive". In: CoRR abs/2412.03444 (2024). DOT: 10. 48550//RXIIT 2412. 03487. arXiv: 2412. 03487
URL: https: /ttps:://do.... org/10 10 48550/arXXi 2412. 2412.03487,
[38] Jiaxin Shi et al. "Simplified and Generalized Masked Diffusion for Discrete Data". In: arXiv preprint arXiv:2406. 04329 (2024).
[39] Xinyou Wang et al. Diffiusion Language Models Are Versatile Protein Learners. 2024. arXiv:
2402.18567 ccs LG] URL: https://arxiv: Sttts://arxi.oorg/aaaa. wtwpp///////......./////// 18567,
[40] Xinyou Wang et al. DPLM-2: A Multimodal Diffiusion Protein Language Model. 2024. arXiv:
2410.13782 CG. L URL: https://arzis: org/ aa////bs/24100 13782
[41] An Yang al. "Qwen2.5 Technical Report". In: arXiv preprint arXiv:2412. (2024).
[42] Jachyeong Jo and Sung Hwang. Continuous Diffussion Model for Language Languag Modeling. 2025.
arv: 2502. 11564 LG]] URL. /ttps://arxiv 2rg/ 11564
[43] Shen Nie et al. Large Language Diffussion Models. 2025, arXiv: 2502. .99922299222 Cc. CLJ URL:
attps://arxis http://wwwvvvor....o//o abs/ /2002. 09992.
[44] Jingyang Ou et al. "Your Absorbing Discrete Diffusion Secretly Models the Conditional Dis- tributions of Clean Data". In: ICLR. 2025. URL: https://openrewiie: net/forum?id
3MyXP8Tanm
[45] Yinuo Ren et al. "Fast Solvers for Discrete Diffusion Models: Theory and Applications of High-Order Algorithms". In: CoRR abs/2502.00034 (2025). DOI: 10. 48550 1.8550/ARXTVV 2502.
00244,
arXiv: 2502.000234 URL:
48550/ arviv 2502. 2502200224.......
[46] Chenkai Xu et al. "Show-o Turbo: Towards Accelerated Unified Multimodal Understanding
and Generation". " In: arXiv preprint 15X1v:2502.05441 (2025).
[47] Zebin You et al. "Effectivv and Efficient Masked Image Generation Models". In: arXiv preprint xixiv22003.0000 097 (2025).
[48] Yixiu Zhao et al. Informed Correctors for Discrete Diffusion Models. 2025. arXiv: 2407 css LGT. URL: https://arxiv. org/ abs///b///// abs//4///700 2407 2404.21243......
21243
12

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The contributions and scope of the paper are well summarized in the abstract and the introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations of work in Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are discussed in the paper.
The authors are encouraged to create separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The au- thors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the ap- proach. For example, facial recognition algorithm may perform poorly when image resolution is low images taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to ad- dress problems of privacy and fairnsss.
While the authors might fear that complete honesty about limitations might used by reviewers grounds for rejection, worse outcome might be that reviewers discover limitations that aren't it acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
13

---

Justification: Assumptions are provided in our theorems, and complete proofs are provided in the appendices.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theo- rems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proofs provided appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We include the detailed experimental settings in Appendix pendix BTG to ensure the reproducibility our experimental results.
DW and Ap-
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or the contribution it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, other means that are appropriate the research performed.
specific model and empirical evaluation,
While NeurIPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results way to re- produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case thors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
14

---

reproducible, they Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: All the datasets and models used this paper are open-sourced our experimental details are provided in Section ! and Appendix 1..2
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, understand that this might be possible, so No is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results, See the NeurIPS code and data submission guidelines (https:
Sol.eSummeee/oeSmmmmmcca ssi ssionnilccy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: We provide the complete experimental settings in Section B and Ap- pendix P
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined other appropri- ate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: The standard deviation results displayed in Appendix EKL
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
15

---

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should prefer- ably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out of range (c.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The details are in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that didn't t make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics httpgs/s//aarrrg......aal/l/ltad
Answer: [Yes]
Justification:
The authors have reviewed and followed the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss the broader impacts of our paper in Section
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact spe- cific groups), privacy considerations, and security considerations.
16

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g., gated release of models, providing defenses addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g.. pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper does not pose such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by re- quiring that users adhere usage guidelines or restrictions access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We have cited the original owners of assets used in the paper properly.
Guidelines:
existing assets. The authors should cite the original paper that produced the code package or dataset.
The answer NA means that the paper does
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the pack- age should be provided. For popular datasets, paperswi thcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the cense of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
17

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer:
[NA]
Justification:
This paper does not release new assets. Our evaluations are based on existing
data.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
For crowdsourcing experiments and research with human subjects, does the pa- per include the full text of instructions given to participants and screenshots, if applicable, well as details about compensation (if any)?
Question:
Answer: [NA]
Justification: This paper does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but the main contri- bution of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, cura- tion, or other labor should be paid at least the minimum wage the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: This paper does involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equiva- lent) may be required for any human subjects research. you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
18

---

16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [Yes]
Justification: We have described in detail in Section ! and Section E how we use LLMs to evaluate metrics.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our L.M policy (https: http:::///eurippp http://rrpps..ccooooeeeeeeeeaeee / for what should or should be described.
19

---

Contents
1
Introduction
2
Background
2.1
Discrete Probability Flows and Masked Diffusion Models
2
2.2
Kinetic and Conditional Kinetic Energy
3
2.3
MDM as Geodesic Curve
3
Main Results
4
3.1
Equivalence Energies in MDMs
3.2
MDMs with the Optimal Mask Schedule Are Energy Minimizers Energy-Inspired Fast Samplers
5
3.3
6
4
Experiments
5
Related Work
8
6
Conclusion
10
Appendix Detailed Notations and Definitions
21
Appendix An Intuitive Explanation of Geodesics
21
Appendix C
Proof of Auxillary Lemmas
22
C.1 Derivation of the Relationship between ag and or in MDM
22
C.2 Proof of the Invariance of Training Loss to the Mask Schedule
23
C.3 Derivation of Conditional Rate Matrix of MDM
23
C.4
Decomposition of Conditional Kinetic Energy along Sequence Dimension
24
Appendix D Proof of Main Results
25
D.1 Proof of Theorem 3.1
25
D.2 Proof of Theorem 3.2
26
D.3
Proof of Example 3.3
28
D.4 Proof of Lemma 3.5
28
D.5 Optimality the Discretized Case
29
D.6 Proof of Theorem 3.6
30
D.7
Proof of Proposition 3.7
32
Appendix
Experimental Details
32
E.1
Details of Beta Parameter Tuning
32
E.2
Standard Benchmarks and Evaluation Settings
33
E.3
Additional Results and Raw Data
................................
33
E.4
Further Discussions on Task-specific Schedule Preferences
33
E.5
Additional Samples
36
E.6
Details of Toy Sampling Experiment in Fig.
36
20

---

Appendix A Detailed Notations and Definitions
n: the sequence length.
I, 2: n--dimensional vector representing states in a model.
D: the vocabulary with size [D] = d.
x' 2f E D: the i-th token of data x; 2.
m(z): the number of mask tokens in z.
(2)( the DPF that connects simple distribution po(=) and data distribution Pı(z) g(2).
P=n(=/<1): the conditional probability flow conditioned on the data.
P1te(z1 the posterior distribution conditioned on time l.
0e: The mask schedule function.
Q:: The transition rate matrix the unmasking process of MDM at time t.
Quj1
: The conditional rate matrix of the unmasking process of MDM at time I.
T7: The transition rate matrix of the masking process of MDM at time f.
The transition rate that uniquely determines Q MDM settings.
Øz:
y: The interpolation schedule function the geodesic curve on the high-dimensional sphere. Also the weight function choose in three energy functionals for energy mini- mization.
(a, b): Beta distribution with parameters a and b.
CDF: cumulative distribution function
Appendix
An Intuitive Explanation of Geodesics
This appendix provides an intuitive introduction to geodesics and exponential maps to clarify the geometric interpretation of MDM in [[2]] Readers seeking formal mathematical definitions of these differential geometry concepts may refer to [S]] for complete technical specifications.
smooth high-dimensional surface, such D- dimensional sphere SD-1 In our scenario, the embedding Eq. (8) maps the per-token conditional probability flow Eq. (3) onto SD-1, shown in Sec. DXL On every point yo the manifold, there exists a tangent space Ty, containing all vectors starting from yo and tangent to the manifold.
Manifolds and Tangent Spaces. A manifold
Exponential Map and Geodesics. The exponential map is denoted exP (v). which maps tangent vector v E Tyo to point y/1 on the manifold. Geometrically, this represents moving from Yo along the "direction" of v at constant speed until reaching y1. This movement follows geodesic path, which is both the "shortest path" and the "straight path" between two points on manifold, generalizing straight lines in Euclidean space. For example, great circles geodesics spheres.
Inverse Exponential Map and Parameterized Geodesic Trajectory: The inverse exponential map is denoted as expe (y1). which maps manifold point y1 back to tangent vector E T This vector encodes both direction and distance from yo to y1- Therefore, given start/end points yo and yı on the manifold, geodesic trajectory parameterized by " (strictly increasing with "o 0, 71 1) can be expressed as:
exp
pl yo (v)) =Y exp (21), E [0,
(19)
This formulation implies:
exp (y1) is the tangent vector encoding the direction and scale that generates the geodesic curve from Yo to y1.
7t is the interpolation schedule and " t means constant-speed motion along the geodesic.
Recent theoretical advances [ reveal that MDM's conditional probability flow Eq. ( forms exactly the geodesic curve in spherical geometry (see Fig. 8 The interpolation schedule is uniquely determined by the mask schedule ar as 7t. arcsin you, which is equivalent with Condition [z!
21

---

No (VPu(a), (Po(b), (P)())))po()).......)))))))))))) (Po(c).....)
a 2
1/2
y (v) 1p.0)..0
Figure 6: The per-token conditional probability flow in MDM generates exactly the geodesic curve.
Appendix C Proof of Auxillary Lemmas
In this section, we provide complete proofs for the auxiliary lemmas referenced in the main text for completeness.
C.1 Derivation of the Relationship between or and o, in MDM
Lemma C.1. The mask schedule og relates to the rate de via expression
= exp (-J, a,ds
(20)
Proof. From the definition of mask schedule in Eq. (G), suffices to prove
P(r) = xi( = exp (-J." a,ds
(21)
Consider infinitesimal time intervals (1, - A/ where each token experiences masking probability o,Al oAt).. The preservation probability therefore satisfies the following product bounds:
[(1 [0-t)/Ar]+]
II K=
(1 + & P(x = x((=))
(22)
((1(1/At)
II K=
(1- (L o(At)).
(23)
Analyzing the upper bound through logarithmic transformation, we get
L(1 (U-t)/At)
[A-))/Ar]
(1 (1- (1=o o(At)) = exp
Σ k=
log(1 (1(1--0--sAA++((AA))
(24)
[(10)/AA)
exp
(-01-RArAt+ +((t))
(25)
k=
(1)
exp(
al-ndu)
(26)
exp
a,ds
(27)
22

---

where in (1) follows from Riemann sum convergence as At + 0 and (2)
applies the variable
substitution s = 1 align integration limits.
The lower bound converges identically through analogous arguments. This completes the proof.
C.2 Proof of the Invariance of Training Loss the Mask Schedule
Different equivalence expressions of the training loss of MDM has been proved invariant to the choice of ar multiple works [112, G8, 36). We adapt a proof from [36] by examining the negative evidence lower bound (NELBO) through token-level denoising components:
1 at /
- -à+
ONLLBO Ep (=1)).)mmn(())))
log(z0(22: t), z1)
(28)
-ùt
=en(e))mnnn() 1 ai
log (og( ,t), af)
(29)
Despite the apparent dependence in its parametric form, the loss exhibits fundamental invari- ance as formalized below.
Proposition C.2 (Schedule Invariance, Proof Adapted from [[]]). CNELBO is invariant to 0z func- tional form, depending only on its boundary values on = 0, 0 =
Proof. The invariance emerges through variable substitution via the chain rule. Let log(1 -0z). then
t=1 -a² (26,
ONELBO (e))
(30)
1-a a4
E Epp(er)
log(zs (2+, t), x)d(f(t))
(31)
2 (pt(t)))
(m)) (-)) (-))
(32)
Ep (z) (()))
log( 2))d
(33)
Here (1) applies the substitution f(t) log a.). (2) applies change of variable y f().. In
(3) we let &y = 2f 1(y) to (2,, (2, = (2,, f- ( The final expression contains no explicit dependence ag trajectory between its fixed endpoints, thereby establishing the claimed invari- ance.
C.3 Derivation of Conditional Rate Matrix MDM
Although explicit sampling through 12, a|m)) remains unnecessary in MDM, establishing closed-form representations proves valuable for theoretical characterization.
Lemma C.3 (Conditional Rate of MDM). The following conditional rate matrix generates MDM's unmasking process:
Qr(2, f
(z, 2) (2((2)
(1)(=11)> otherwise
4 - T1
Eayz Qe(z,
otherwise
(34)
Here x denotes single-token unmasking transitions defined in Sec. and x = x1 denotes that In can be generated from x through one or several steps of unmasking.
Applying Bayes' theorem establishes the posterior relationship:
P1)(0112) Pej (z)
(35)
(x(2)
23

---

we know that the positivity condition (1)((1112) > 0 consequently requires (2/2)) > 0 and 2 21 ensuring the rate matrix's 's well-posedness.
Proof. To verify that (1(2,z(211 generates MDM's demasking dynamies, we confirm the consis- tency condition Eq. (g) holds. We verify this through the following derivation:
Σ (((2, Pep (2 mz (11))
Σ (0(eeo
(36)
(e(2)
Pip Pil (1(()) Pap (2)21))
Σ (1, 2)
(37)
Paj1 (2
pi(2)
Peti (1
Σ (2),2)
(38)
pr(z)
pc(x)
(x,,2) 2)
Q1(z,z).
(39)
pz(2)
We subsequently derive its closed-form of (2 x/21). If x t then (at((( (1) = 0, yielding
Q((2, = 0. If - x > 21, on the other hand, we have
P (x*(a)))) Piji (2 [xt)
Q0(2, = 5t 2)
(40)
a0
thus completing the proof.
C.4 Decomposition of Conditional Kinetic Energy along Sequence Dimension
In this section, we show that under many DPF frameworks such as MDMs, the conditional kinetic energy can be decomposed along sequence dimension as following:
Eppp Qq) E, (app)))))))
(2le))) Quj (2,
(41)
(11(2)a11)
= (x1)
Σ
Quj (2) x|21)
(42)
uep (x((a)
Peli (eji (2))
Eg (x (mn)
c
Qep (2 11))
(43)
7epo)1(" (x)
uPPq (xt (ot( (1) Qun(=" x |21)
(44)
The pivotal decomposition in (1) leverages two structural properties: (i) The conditional probability flow (Bq. (G) exhibits token-wise independence, and (ii) The conditional rate matrix (Lemma (K) nullifies transitions altering multiple tokens simultaneously. These enable reduction of full-sequence transitions to single-token operations, with remaining n tokens contributing constant combinato- rial factors. This structural property persists across various DPF implementations including MDM and Discrete Flow Matching [32) B7i] validating the conditional kinetic energy as theoretically sound surrogate objective. Notably, standard kinetic energy E lacks such decomposition due to p. (2))s de- pendence on cross-token correlations.
For MDM's binary mask dynamics (of vs. [MASK]), the combinatorial constant becomes C = 2n-1;
(a(a)))) = 2 E (2)) m))
Σ
1 Q Qu1
(x1)) (45)
24

---

This decomposition permits notational relaxation where & analysis considers tokens 2, 2 E D inde- pendently of sequence context, slight abuse off notation adopted in Appendix D.2'ss equivalence proof.
Appendix D Proof of Main Results
D.1 Proof of Theorem 3.
Theorem 3.1 (Kinetic-conditional equivalence MDMs). For any weight function Ytt and MDM with mask schedule 0t, the marginal and conditional kinetic energies are proportional:
Ea(or, 1 = CiEe(au,
(46)
where C1 is a scalar depending only on the sequence length n and vocabulary size d. As a result, the two objectives share the same minimizers:
arg min Eg 7e) = arg min E,(o 7t). di
(47)
Our argument leverages a foundational decomposition from + regarding concrete score represen- tations:
Lemma D.2. for = = (=',...,=*= [M],...²)), x= (2), x² * (MI...2") we have
(r(x)) 04 1- pr UM
(48)
P((z)
where ZUM the vector consists of all unmasked tokens of We now prove the main theorem:
Proof. Let m(()) quantify the masked positions in assumed w.l.o.g. to occupy initial sequence positions. The key summation decomposes as:
m(=)
Pr(x)
Pr(r)
Σ
(49)
Pt(2)
(x(2)
[
m(s)
(50)
at
i=1
x'> [M]
(=)
Ot
Σ
(51)
04
1=1
z [M
m(2)
(52)
Olt
25

---

Substituting the rate matrix from Eq. (2), we therefore deduce
Pc(x)
2
Ex (x))
Σ
(53)
Pc(2)
p+(=)
vri (x)
E
o pi (z)
(54)
dt Pc(x)1e
0f (c))
E, Σ
(55)
it Pt(2)
@ F, Σ
m(z)
0t of
(56)
1 at 9t
at of Et 1 01 7/
Σ m(z)
(57)
at of
Ck E.
(58)
04 7/
Here define Cs Ex m(z). On the other hand, the conditional kinetic energy can also be break down as
Paj (xr g
E, E. ((p(
Σ
(59)
(1l((()1)( 2:2
Ptju Pej (2/x11)
of Paj (x)
= =R(p((z1)
Σ
(60)
2t Pit1 P211(2/az (
2,212 -a
04 of 1 0t 7t
Expp(()))
Σ
(61)
Σ
Σ
p1(11)
Eg Cxt o
(62)
aq 7
F1 2,
o,
n2"
Σ (111)
E.
(63)
02
7t
a₁ 0
-122-1pp
E,
(64)
0t
C.
E.
Qt
Tt
(65)
0t
1t
This expression of & is equivalent to the decomposition in Appendix C.4 by plugging in the explicit form of Qij1- The proportionality constant C₁ C2/C. emerges from comparing both expressions, with Ck, C. depending solely on architectural parameters and [D] = d. The minimizer equivalence follows directly from the strict positivity of scaling constants.
D.2 Proof of Theorem S
Theorem 3.2 (Conditional-geadee equivalence in MDMs). For any weight function "tt and MDM with mask schedule Ot, the conditional and geodesic energies are proportional:
Ec(au, 7)) = C2Eg(as, 7t),
(66)
where C2 is a scalar depending only the sequence length n.. This implies that they share the same minimizers:
arg min Eorr, n)) = arg min Eg 1c). (c).
(67)
26

---

Proof. The geodesic energy is inherently defined using the token-wise independent conditional flows Eq. (3), therefore it admits straightforward decomposition along sequence dimensions. In MDM case where all tokens follows the same mask schedule os, we further have
Eg(Po1171)=
Σ E (o
72
4
Del1 (2))
(68)
upep (2))
4
pn (x)) m)) 71Pel1 (z 21))
Val2 (2 x1)²
(69)
4
Cg Ee (ui).. (=*1)=)) 7xPoj1
Wel1 (= (x1)²
(70)
(2
On the other hand, through dimensional decomposition established in Appendix C.4 the conditional energy admits:
EpPii Qqpi7)) = -2n-1_ (11))
Σ artiar of
Q x |=)=)
(71)
Pepil (at (m))
-Ce Br.p (o1))P
Qu1 x1|=
(72)
Pepp (x* (x))
Where we define C. n22-1 as in Appendix D.I. Therefore, The equivalence proof reduces to analysis through notational relaxation, treating 2, E D as individual tokens.
Leveraging the rate matrix expression from Lemma CKI have
(2
(Pip. Qu1i:te E2 (2(2)))) ft
Quj (2,
(73)
2,210
(2(x2)
2
-(pp)(())
Σ
J,att
(74)
pqp (I) x1)
Oi
zm [
1 - a
0xatt
(75)
71
ar
[M
of
E (x(())
(76)
which is equivalent the expression in Appendix DJI by letting n 1. Applying the relationship de agd; deduced from Eq. (e), we get d: 0407 therefore we further have
1 of Ee(pi Qqjit) (r, % ar at)
(77)
On the other hand, applying Eq. (.) Eg in n 1 case can be expressed as
4
Σ di V (=](a)))) 2
Eppniiitt) = Ex,pr (mn)
(78)
7t
1
Σ
(=(2)12)2
= Er,pi
(79)
Pep (2(12))))
1
( (z) 8m (z))]²
-x,p((01)
(80)
7t
andz (2) (1 04 )om (z)
-epr(or))
1
Σ
di
+
Σ
(81)
7
at
(1 1-ar)
ac)
==
x=|
1
2
2
+
(82)
04
(1
01)
27

---

Since
d;
0g
d2
0
+
(83)
0+)
1t
(1-04)
the functional equivalence in scalar case is established, extended to n-dimensions through the di- mensional scaling factor C2 Ce/Cg 2"-1, The minimizer equivalence follows from strict positivity of scaling relations.
D.3 Proof of Example t.
Example 3.3. When it = 1, the kinetic, conditional kinetic, and geodesic energies all reduce to:
E(o,, (1) = (-¹ 5 acl1 -a))
6? dt.
(84)
ac)
Proof. In the case, we have
C Σ m(z) = m ((M)) 1;
(85)
C. 112"1=11 =1
(86)
C. 1.
(87)
Therefore, have C1 = C = 1 and the three energy functions share the same form Eq. (14).
D.4 Proof of Lemma 355
Lemma 3.5. Under Condition B.4 the schedule of minimizes the geodesic energy.
Since the geodesic energy (Definition (.3) is defined by summing the token-wise conditional proba- bility flow, we only need focus on the one-dimensional case. Therefore, we abuse notation slightly by regarding E D:
4
(p(ppn))) Bip (m) D.p ((l2))) (2 Dup (=|
(88)
Beep (m.) Σ
4
uqn (=/
(89)
%
-E,pp (w))
(90)
Here you1 denotes the d-dimensional embedding vector induced by the embedding Eq. (8) Therefore, the minimizing problem becomes:
/' [t)(1)
arg min s(t) (()
di
(91)
s.t. [|v())]| vt.
(92)
For baseline case "/ we construct the augmented functional with Lagrange multiplier X(t):
Clyl /. (11x(t)11 + x(0((Il(()))= 1)) dt.
(93)
The Euler-Lagrange formalism yields:
ac d Ac
0,
(94)
Oy dt Di
from which we derive the critical differential relationship:
i = Ay-
(95)
28

---

from |b|1 = 1 we have vÿ + |1 0, therefore we have
-1bi|| 2 yÿ = \lyvll =
(96)
Plugging this expression of À into Eq. (H3),, get
j = -
(97)
which corresponds to the uniform circular motion with its acceleration pointing towards the center of the sphere. Therefore, the route follows the great circle connecting yo and 3/. In MDM case where yo JTMj and !/1 represents clean data without mask token, we have yo y/ 0. Therefore, the angle between yo and yn is 2 and the curve shares the following simple form:
T 20 + y sin( T !)- y(() = yo cos(
(98)
Now generalize arbitrary 7n schedules through temporal reparameterization:
[
[1)(1 {(t)
(1(((()))1
arg min v(t)
dt arg min v(+(0))
dt
(99)
J
(tt)
dy
arg min
|| ied
(100)
v(v(())
Jn
d7
[
dy
arg min
11
(101)
v(7)
dry
Therefore in ordinary Y cases, the optimized route is the geodesic curve rescheduled by interpolation schedule y4:
π v()) = cos(
(+)) + yn sin( π 7t)- n
(102)
By squaring both sides of Eq. (I022,, further recover MDM's conditional probability flow:
Peji(t (t) (1) Po cos² (" 7) + Pı sin ( ())
(103)
a P1 + (1 az po.
(104)
where (1) leverages orthogonality Po = 0. Therefore, we proved that MDM with schedule schedule of generates the minimal-length curve well minimal-energy conditional probability path, validating and generalizing the conclusion in [22] from a energy perspective.
D.5 Optimality in the Discretized Case
Here show that under Riemann discretization, MDM trajectories still minimizes corresponding discrete energy functional that converges to the continuous formulation as the number of steps N + Recall that in Appendix D.4 proved that Eg can be equivalently expressed as:
Eg f/ 1b
dt.
Therefore, discretize Eg
Mn+1 (At)? T6++1 / A/
1n++1 3/ll2
EN
At Σ
1n+1 7/1
where At = 1/N. For simplicity, consider the special case where t. Then we have // 7n = N. We now demonstrate that the N equidistant points along the geodesic still minimize this functional. First, we note that the squared chord length on the unit sphere satisfies || yn|l² = 2(1 cos(O.)), where O, denotes the angle between yn. and yn++- This transforms the minimization
29

---

problem into finding N intermediate points along the great circle are between yo and y1 that minimize Σn 2(1 cos()
From the properties of spherical geodesics, we know that placing all intermediate points along the geodesic path enforces the constraint Σ0, 0n 2 Given that cos(x:) is convex x/2x/2). Jensen's 's inequality establishes:
2 cos(O)) 2 N 2(1 cos(
π 2N )),
with equality achieved when points are uniformly distributed along the geodesic. Any deviation from the geodesic path would further increase the cumulative angular separation Σ 0, beyond z. conse- quently raising the total energy. This proof extends to arbitrary schedules through proportional point distribution based on | 7n/
D.6 Proof of Theorem
Theorem 3.6 (Kinetic energy minimization). Under Condition B.. the MDM schedule of simulta- neously minimizes all three energy functionals.
Proof. Theorem B.I. Theorem 3.2 and Lemma [. collectively establish that oj optimizes the three functionals over all mask schedules Oa2, This conclusion forms the theoretical foundation for our data-driven schedule tuning framework presented in Section GIR!
However, &g and & are defined on both probability flows and rate matrices. Therefore, it remains be further proved that when the probability flow uniquely induced by a is fixed, the conditional rate matrix, which is also determined by az, still minimizes the energy functionals, i.e.
Qap(az) arg min (e(pp 2411:%),
(105)
Qu1
where Qui (a() refer to conditional rate matrix in MDM case derived in Appendix c under Con- dition B.4. We adapt the following key result from prior analysis [37]]of the "t I case:
Lemma D.7. The following conditional rate matrix minimize the conditional kinetic energy & when the probability flow pz()) and weight function 7 =t fixed, i.e.
arg min (P Qqii" 7/ = =) = Qij = āt Qaa
1- 1 (6x,
(c) 5x))),
(106)
where the arg min is taken over any possible Q4p that generates the fixed probability flow.
Here a defined using the conditional probability flow under Discrete Flow Matching (DFM) settings, resembling the MDM case (see Eq. (L)) by conditioning on both ends (t = 0 and (= 1)
of the flow. In the case when p=(=) = εMd (z), it coincides with the mask schedule defined in our work. The notation is also slightly abused by regarding 2, x E D justified by the decomposition of
conditional kinetic energy in Appendix [
We now demonstrate a non-trivial result: MDM under Condition c4 inherently achieves optimal conditional rate matrices for arbitrary %t. despite structural constraints.
First, consider % = where Appendix C yields the conditional rate matrix. We demonstrate that MDM achieves the optimal velocity specified by the RHS of Eq. (D.7). For 1, the conditional rate matrix simplifies to:
z [M] 1 11
= [M], 21
((1() =
a.a
oi
2 = [M x=
(1)
I-a:
de
2 = [M] M=
(107)
I-ac
otherwise
otherwise
where (1) follows from the identity d; ata, established Eq. (e), thus completing the 7 case.
30

---

For general Ye we reformulate the conditional kinetic energy using Definition 2
Ec(Qupi)) = E T8
Σ
(108)
where f and A are independent of Q. We first establish the following key lemma: Lemma D.8. If
(t) arg min E Qu
211,2,22)00111)) 2,
Σ
(109)
then we have
(n) E arg min
E, Qup
Σ 21)12,2))0(1(0)'
(110)
"t
2,E,RIEA
where the arg min is taken over any possible Qqp that generates the fixed probability flow.
Proof of the lemma follows via the substitution W, = Q1j1((). Let
(t) E arg min E, W(t)
Σ 21)
(111)
EA
then (+()) (t). Therefore, have
Σ (1)
arg min E, Qi11
(112)
tt
d
arg min E, 7t.
Σ E,
f(r,,
(113)
dl
d
dW dW jz
arg min
E.
Σ
((,,,,
(114)
dt
Wxc0
dt
d
dW /W/ )3 Σ (1)
arg min Erit
(115)
di
W. 7
dr(1)/dl
EA
d
dWzys f(2,0, 21)(
arg min E. W
Σ
(116)
dt
dy
d
[w)()) = Q
(117)
proving this lemma. Given MDM's optimality under 71 i.e
du
π sin ( (t) cos(s (5t) cos)
π tan π
(t) = (at)
2 t),
(118)
7=
ly=
it remains to prove that the conditional rate matrix in MDM in general "i cases satisfies the LHS of Eq. (LII). Since we have
(a)( = 1
di
r sin [ 7c) cos (natt = x tann cos (E))
π ( 7))..
(119)
thus (af( = (*()) and MDM's intrinsic optimization across arbitrary 7t obtained. Re- markably, this result transcends geodesic energy &g (defined solely through Ptj1). demonstrating MDM's dual optimization of both probability flows and sampling matrices despite structural con- straints.
D
31

---

D.7 Proof of Propasition 3.7
Proposition 3.7. Linear and squared cosine schedules correspond to specific beta parameteriza
tions:
at
It CDF ((0.5,05s)))) Y = 1 = CDF
(120)
02 = sin"
(121)
Proof.
The probability density function (PDF) of the Beta distribution Bia, b) is defined as:
p( a, x (1 x=)-11
(122)
Baa b)
where B(a, ((a)))))(( is the normalizing constant. For the case 0.5:
(00.5) n
B(0.5, 0.5)
(123)
((1)
The cumulative distribution function (CDF) therefore given by:
(DF
dt
(124)
TV t(1-t)
aresiu v7
1
(2 @ cos 0d0)
(125)
sin O cos 1
aresin Vt.
(126)
where step (1) employs the trigonometric substitution x sin² 0. Therefore, when 7t (DF we have
- 042 = sin s (1))
(127)
For the case a 1:
B(1,1) = (1))T(1)
1.
(128)
T(2)
Therefore the CDF simplifies to
CDF S/ 1dt t
(129)
inducing a; (5t). This completes the proof.
Appendix E Experimental Details
E.1
Details of Beta Parameter Tuning
As stated in Section + our theoretical analysis suggests that different downstream tasks may re- quire generated text to possess specific intrinsic structures, which in turn necessitates the generation process to emphasize particular temporal phases. These temporal preferences captured through different y schedules, which induce corresponding optimal a through Condition 3.4
Therefore, we hypothesize that the schedule preferences are mostly inherent to task nature rather than data specifics. To validate this hypothesis, we conducted the following experiments demon- strating that randomly chosen small subsets (50-150 instances) of test data suffice for reliable schedule selection. Specifically, we compare schedule performance between small test subsets and full evaluations in Table M
32

---

Table 1: Schedule performance between small test subsets and full evaluations under different Beta-CDF pa- rameters.
Beta-CDD Parameters
(0.5,0.5)
(1,1)
(0.9,0.3)
(0.3,0.9)
Task: GSM8K (f)
(length=128, steps=32)
Random subset 1 (^=132)
44.70
43.94
31.06
0.00
Random subset 2 (7=132)
46.97
41.67
40.91
0.00
Full test set (n=1319)
38.06
34.80
29.04
0.08
Task: Human (1) (166 steps=64)
Random subset 1 (n=82)
8.54
20.73
26.83
1.22
Random subset 2 (n=82)
18.29
24.39
30.49
2.44
Full test set (n=164)
11.59
22.56
24.39
1.83
While HumanEval evaluations exhibit greater variance to smaller test populations (n=164 total), the relative performance rankings remain mostly consistent across subsets critical indicator of our method's robustness. This empirical validation confirms that schedule preferences are very probably induced by intrinsic task attributes rather than specific data instances.
Therefore in practice, we recommend conducting initial grid searches using small random test data subsets (about 50 150 instances) across parameters a, {0.1, 0.2, .... 1.0} } to find a set of well- performed schedules. After the coarse search, we can perform finer- seer-ggaieed selection on shortlisted candidates using larger data subsets. This approach achieves comprehensive parameter space explo- ration while maintaining computational feasibility.
E.2 Standard Benchmarks and Evaluation Settings
In this section, briefly introduce the evaluation benchmarks and describe the experimental details.
Building upon established practices in LLM evaluation [30, 2TI PK
J. we evaluate performance across
key dimensions including: general ability (BBH [28]), mathematics (GSM8K [ITC) Hendrycks
MATH [[]]], Minerva MATH [I8]), and code generation (MBPP [[Z], HumanEval (9I). Evaluation
follows the conditional generation paradigm, where models produce completions given task prompts, with performance quantified through exact match or other -ppaiin evaluation metrics.
Our implementation leverages the open-source pretrained weights and evaluation toolkit from LLaDA [43],, modifying only the mask schedule that governs the iterative unmasking process. The mask schedule affects the number of tokens unmasked at each step, with certain schedules permitting zero-token unmasking during the process. Therefore, generation quality discrepancies occur even when sampling steps are set as the sequence length. All experiments can be efficiently conducted on a single NVIDIA A800 GPU.
E.3
Additional Results and Raw Data
Fig. Ø shows the result of our main experiments on benchmark BBH [28] and GSM8K [اC]], where our beta-parameteee parameterized schedules exhibit comparable yet not better performance than the linear schedule.
Tab. 247 shows the raw data all our main experiments. We highlight entries matching or exceeding the highest mean within statistical variance (=1std) in bold.
E.4 Further Discussions on Task-specific Schedule Preferences
Our methodology provides both theoretical foundation for understanding schedule preferences and a practical mechanism for task-specific optimization - overcoming the limitations of previous mne-size-fits-a approaches rather than proposing universally superior schedule.
33

---

25
CDFIDIO 0.33)
SEDDRZZ
20
sin Un2
40-
C..000330000
CD.L800..009)
GSM8K
BBH
30
COFLB009 0.99
20
10
16
128
Stop
Step
Figure 7: Performance evaluation of energy-optimizee schedules BBH [28] and GSM8K [TII] where
beta-parameterized schedules exhibit comparable yet not better performance than the linear schedule. Each panel corresponds distinct benchmark. The x-axis displays sampling steps on logarithmic scale, while the y-axis quantifies task performance, where higher values denote superior generation quality.
Table 2:
Performance evaluation of eetteearameterized schedules on MBPP [ benchmark. All experi-
ments fix generation length 256 and higher values indicate better sampling quality.
Interpolation Schedule
Mask Schedule
8
16
32
64
128
256
Manually Designed Schedules
(DH(B(05500.5) 0.5))
16.4 1.66
40.6
sin (
0,20 1 0.20
0.4 0.28
3.8 0.86
10.4 1. 1.37
19.8 1. 1.78
28 1. 2.01
(DF(B(1,
sin (3()
0.53
4.8 0.96
13.8 1.54
20.6 1.81
2.14 35.2
40.2 2.19
Repurameterizznng
Schedules
CDF(B 0.9)
0.20 1: 0.20
0.20 1: 0.20
0.20 0.20
0.20 0.20
0.20 1001:1000 0,20
1.6 0.56
CDF(B 0.3))
= 0.92
14.4 1.37
23.4
318
39.8 2.19
2.19
CDF(B(1,0.2)
5.8 1: 1.05
15.4 1.62
23.8
1.91
32 2.09
36.2 4 2.15
38.5 2.18
CDF (0(.9.
0.56
5-9-0.98
11.6 1.43
20.0 1.81
35.4 + 2.14
39.8 2.19
Table 3: Performance evaluation of beta- paraaeteraeee schedules HumanEval M benchmark. All experiments fix generation length at 256 and higher values indicate better sampling quality.
Interpolation Schedule
Mask Schedule
8
16
32
64
128
256
Manually Designed Schedules
CDF(6 0.5))
1.83 + 1.05
1.83 + 1.05
5.40 1.78
11.59 11:00+220 +:5.1
25.61 3.42
32.32 3.66
(5))
1.06
CDF(B(1, 1))
sin! (2.
4.27 + 1.58
11.59 2: 2.51
13:41 2.67
22.56 3.27
27.44 3.19
31.3 3.66
Reparameterizing
Schedules
CDF(B(O.3,00 0.9)) (DF.9) CDP(U) 0.2))
1.83 + 1.06
1.83 1.05
1.83 1.05
1.83 1.06
1.83 1.05
1.83 1.05
4,88
1.69
12.8 12.6..2.66
17.68
2.99
24.39
3.36
31.71
3.64
31.1
3.63
2.39
2.72
20.73
3.18
3.47
28
3.54
3.59
(D. 0.9))
3.05 1: 1.35
7.53
2.122
10:37 1 2.30
24.39 3.36
28.66
3.54
32.32 3.68
Table 4: Performance evaluation of beta-parameterized schedules BBH [2]] benchmark. All experi- ments generation length at 128 and higher values indicate better sampling quality.
Interpolation Schedule
Mask Schedule
16
32
64
128
Manually Designed Schedules
CDF 0.5))
12.76 0.32
15.83.4
17.32
0.34
17.91 0.35
19.29
0.35
siní f)
11.95 0.32
14.58
0.33
16.36
0.34
17.42
0.35
18.06
0.35
CDF(B(1, 1))
sin 2/
11.29 = 0.31
14.21 0.32
16.83 0.34
17.32 4: 0.35
18.6 : 0.35
Beta Reparameterizing Schedules
CDF 0.9))
3.7 4: 0.22
9.12 0.3
12.04 0.31
16.82
0.34
17.
0.35
CDF
7.71
16.13
0.34
(DF(6
12.23 = 0.31
15.87 3.37740.34
17.02 : 0.35
17.11
0.34
19.12
0.35
C 0.7)
12.24 0.32
14.91 0.33
0.35 16.66
17.6 0.35
19.28
0.35
11.67 0.311
1433 0.33
16.73
0.334
17.62
0.35
18.85
0.35
( 1.3)
10 0.29
13.79 0.32
16.36 0.366-0.33
17.88
0.35
34

---

Table 5: Performance evaluation of beta- parttaaaaaeeeeezedddeoo schedules GSM8K [T0]] benchmark. All exper- iments fix generation length 128 and higher values indicate better sampling quality.
Interpolation Schedule
Mask Schedule
8
16
64
128
Manually Designed Schedules
CDF(B(0.5,0.5)
2.43 + 0.42
18.20 + 1.26
40.26 + 1.35
48.37 + 1.38
49.20
1.38
sinl 51)
0.53 + 0.20
6,44 = 0.68
27.37 = 1.23
15.03 = 1.37
47,08 +: 1.37
(DF(
sin (5()
4.32 + 0.56
17.59 + 1.05
36.47 = 1.33
46.70 = 1.37
17.70 +: 1.35
Beta Reparameterizing Schedules
(..00) CDF 0.3)
0.61 + 0.22
0.91 = 0.18
0.99 + 0.27
0.400:040
8.42
0.76
6.14 + 0.66
15.39 + 0.99
29.12 = 1.26
39.20
- 1.34
46.70
1.37
(D) 0.8)
4.62 0.58
19.64 + 1.09
40.26 1.35
46.02
1.37
48.52 1.38
CDF 0.9)
4.78 + 0.53
19.94 + 1.10
30.54 = 1.33
45.49 + 1.37
48.14
1.38
Table 6: Performance evaluation of
schedules on Hendrycks Math []]]]] benchmark. All experiments fix generation length at 256 and higher values indicate better sampling quality.
Interpolation Schedule
Mask Schedule
8
16 Manually Designed Schedules
32
64
128
256
(DF(5(0.8,0.5)
11.3 0.41
11.58 0.44
12.78 1: 0.46
16.3 0.51
18.86 2: 0.54
20.24 0.56
sin (
11.5 + 0.44
11.5 + 0.44
11.68 +: 0.45
14.36 + 0.48
17.9 + 0.53
18.84
0.54
CDF(B(1, 1))
sin (5()
12.02 0.45
16.54 à: 0.51
18.8 0.54
19.3 0.55
20.18
0.58
20.18
0.58
Beta Reparameterizing Schedules
CDF( 0.99))
11.3 0.44
11.5 0.44
11.3 0.44
11.5++0.444
11.48
11.52
CDF 0.3))
16.32 0.51
18.76 0.54
19.5 1: 0.53
19.9 0.53
20.18 0.56
20.08 0.55
CDF
0.54
0.54
0.55
0.55
0.55
0.55
CDF CDF(9(0099.0 (B(0.900.9)
1..6
15.82
2 0.51
18.5 0.54
19.28 + 0.55
20.2 1: 0.56
20.2 1: 0.56
Table Performance evaluation of pett-parameterized schedules Minerva Math [S]] benchmark. All experiments fix generation length at 256 and higher values indicate better sampling quality.
Interpolation Schedlule
Mask Schechle
16
32
64
128
256
Manually Designed Schedules
CDF(8(0.,0.5)
0.12 + 0.05
0.50
4.70 0.30
16.54
30.10 +
sin (1)
0.04 + 0.03
0.26 + 0.07
0.92 + 0.13
8.74 + 0.39
21.70 - 0.56
27.26 0.60
CDF(8(1, I))
sin (24)
0.30 + 0.08
1.84
6.50 0.34
15.94 0.50
25.8 0.50
29:20
Beta Reparameterizing Schedules
(DF CDF (.39).
0.04
0.12
0.03
0.20
0.34
0.13
1.10 = 0.17
5.56 : 0.32
13.16 = 0.46
20.84
= 0.55
26.841 0.59
29.36 0.61
(DF (,22).
1.28 0.16
4.82 + 0.30
10.64 0.43
17.14 0.51
23.68 0.37
26.78 0.59
CDF CDF(B (.0.9.0..))
0.18 : 0.06
1.48 : 0.17
6.90 + 0.35
17.28 +: 0.51
26.0 0.59
29.5-1 0.61
Therefore, the empirical parity on certain benchmarks (see Appendix EM confirms that linear sched- ules already serve as near-optimal candidates for specific task categories. In fact, linear schedule ar = t is a special case in our framework since it corresponds to ?t- CDF(( 0.5)). a point in our parameter space.
Specifically, ?t = CDF 0.5)) has higher derivatives on both initial and final generation phases while being relatively static in the middle. Since 7 appears in the denominator of our ergy functional expressions, tasks requiring sustained refinement throughout generation (particularly middle phases) rather than the starting or ending phases might inherently favor the linear schedule.
Among the six tasks we experimented on, BBH [28] is only one that focuses on general reason- ing problems. As for GSM8K [MT]] although it is mathematics-focusel its answers use detailed natural language explanations compared to other math benchmarks. This might be the reason why sustained refinement in the middle phases are preferred. Considering this similarity, it is reason- able that GSM8K [TTO] and BBH [28] have similar preference on schedules although it is difficult to mathematically deduce the exact expression of the energy functionals.
Considering the complexity of real-world tasks, it is indeed challenging direction to systemati- cally study the relationship between tasks and their prefered schedules. A promising research di- rection involves designing intermediate benchmarks that balance practical relevance with analytical tractability- more sophisticated than our current toy examples yet simpler than real-world tasks.
35

---

Table Comparison of beta reparameterized and manually designed schedules math-reasoning task Hendrycks Math [[]] Sampling steps set as 64.
Prompt
In fluid ounce bottle of Vitamin Water, there 125 calories. How many calories would be contained in 12 fluid ounce bottle? Express your answer in decimal form.
Beta Parameterized Schedule ((,,000.
Manually Designed Schedule
0g l:
find the 12 fluid ounce bottle, we need
187.5
to multiply the number of calories the 8 fluid ounce bottle by 1.5. So, the answer to be 125 * 1.5 187.5..
Solution:
125/8 187.5
1:8 = 125.5
1:1 125.5
Therefore, the answer 187.5 calories
1:2 = 125.5
are contained in 12 fluid ounce bot- tle of Vitamin Water.input::
1:3 125.5
1:4 125.5
1:5 125.5
1:6 = 125.5
1:7 125.5
1:8 125.5
1:9 125.5
1:10 = 125.5
1:11 125.5
1:12 = 125.5
1:13 125.5
1:14 = 125.5
1:15 = 125.5
1:16 = 125.5
1:17 = 125.5
1:18 125.5
1:19 = 125.5
1:2
E,5
Additional Samples
We present comparative samples from identical tasks and sampling steps under different schedules. Fig. & illustrates how optimized schedules produce mathematically coherent outputs, while Fig. provides code generation examples demonstrating better quality.
E.6 Details of Toy Sampling Experiment in Fig. g
The implementation the toy experiment Fig. B considers target distribution defined over
sequences with 5 tokens and vocabulary size d = 2. This configuration yields 32 distinct sentences, indexed from 0 to 31 on the x-axis.
The vocabulary contains two words and 6 and the 32 sentences on the x-axis are ordered first by ascending count a tokens, with sentences containing identical numbers of tokens further sorted lexicographically. The target distribution in the left panel thus explicitly designates only the extreme cases (aaaaa and bbbbb) as legal sentences, while the right panel's target distribution considers sentences with 2 or 3 tokens valid. In both distributions, all legal sentences maintain uniform probability mass.
Our sampling experiments employ slep 3 without any training phase, as the target distributions be analytically computed. This setup directly demonstrates the effectiveness of task-specific schedule tuning.
36

---

Table 9: Comparison of beta reparameterized and manually designed schedules on code-generating task HumanEval [9]. Sampling steps set as 64,
Prompt def bi(planctl, planet2):
There are eight planets in our solar system: the closerst to the Sun Mer- cury, the next one is Venus, then Earth, Mars, Jupiter, Saturn, Uranus, Nep- tune. Write function that takes two planet names as strings planetl and planet2. The function should return tuple containing all planets whose bits are located between the orbit of planet and the orbit of planet2, sorted by the proximity the sun. The function should return empty tuple if planetl or planet2 are not correct planet names. Examples [[[""upiter", "Neptune")
("Saturn", "Uranus") bf("Earth", "Mercury")=sz: ==> ("Venus") ((""Mrrrry", "Uranus")=== ==> ("Venus", "Earth", "Mars". "Jupiter", "Saturn")
Beta Parameterized Schedule = CDF
Manually Designed Schedule ag = t:
planets = ("Mercury", "Venus", "Earth", "Mars", "Jupiter", "Saturn", "Uranus". "Neptune") if planet1 and planet2 and planets weteessindex(pplarncee (planet1) cts.index(planet2 return
planets = = ("Mercury", "Venus".
"Earth", "Mars", "Jupiter", "Saturn", "Uranus", "Neptune")
planets "Earth". "Mars". "Jupiter", "Saturn" "Uranus", "Neptune")
("Mercury", "Venus"
plan-
planetsplan-
planets = = ("Mercury", "Venus". "
ets. wws.indxx(laaee.indee
(planetplan
"Earth". "Mars". "Jupiter", "Saturn"
planets. blanets.index(( (planet1)
+ 1: plan-
"Uranus", "Neptune")
end.i (anexx(lanet2)
planets "Earth", "Mars", "Jupiter", "Saturn".
("Mercury", "Venus",
return O
"Uranus". "Neptune")
planets = = ("Mercury", "Venus".
"Earth", "Mars", "Jupiter", "Saturn".
"Uranus". "Neptune")
planets = = ("Mercury", "Venus".
"Earth". "Mars". "Jupiter". "Saturn".
"Uranus". "Neptune")
planets ("Mereury", "Venus",
"Earth". "Mars", "Jupiter", "Saturn",
"Uranus", "Neptune")
planets==
37


Reviews:
[{"review_id": "oYuBHOnjTD", "reviewer": "Reviewer_dGWj", "summary": "The submission proposes a theoretical framework for Masked Diffusion Models (MDM) based on energy minimization in optimal transport and from that derives an optimal sampling schedule. Kinetic energy, conditional kinetic energy, and geodesic energy are shown to be equivalent under the MDM formulation, which connects Kinect energy as a sampling-efficiency view with a geometric interpretation (geodesic paths) in discrete diffusion. The authors derive a closed-form condition for the optimality of the mask schedule $alpha_t$, namely $alpha_t^* = \\sin^2(\\pi/2\\cdot \\gamma_t)$, where $\\gamma_t$ is a continuous interpolation schedule. For this condition, the MDM minimizes all three energy functionals simultaneously, i.e. MDMs achieve optimal transport trajectories when the condition holds.\nFrom that theoretical result, the authors derive a technique to tune mask schedules after training, by parameterizing $\\gamma_t$ as the CDF of a Beta distribution with two parameters, which allow for search of a near-optimal schedule in 2D without retraining of the model. This is possible, because MDM training objectives are shown to be invariant to the chosen schedule. Experiments on synthetic data and real benchmarks confirm that the optimized schedules yield faster sampling without a reduction in output quality. For example, on a code generation task, the optimal schedule matches the quality of a linear schedule using only half the number of denoising steps. Likewise, on a math reasoning task, it reached qualitative parity with a quarter of the steps, showing improved few-step generation performance.", "questions": "- The theorems assume a continuous-time interpolation from 0 to 1. In practice, a discrete sequence of masks must be chosen. Is the sine-squared schedule still exactly optimal in a discretized sense, or only in the continuous limit? Currently this is only done implicitly by tuning Beta parameters for a fixed step count. A brief clarification of the theoretical connection would be helpful.\n\n- How are the parameters of the Beta-CDF selected for a given task? Did you perform a grid search on (a, b) and pick the best performing schedule? If yes, what are the parameters of this search? It would be useful to know the number of evaluations needed and whether it is feasible to apply to very large models. Discussing the tuning strategy would also help readers to adapt your method more easily and extend this part of it.\n\n- Do you have hypotheses why GSM8K and BBH did not benefit from optimized Beta schedules? Is it because those tasks require very specific refinement at the end, so a linear schedule is already as good as it gets? Providing some discussion on these cases would be valuable.", "limitations": "yes", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Quality:\nThe work bases its applications on a thoroughly derived mathematical foundation. The proofs appear sound and the central result, Theorem 3.6, in particular is non-trivial (minimizing geodesic energy does not guarantee minimizing kinetic energy) but the authors address this carefully in their argumentation.\nEmpirically, diverse benchmarks were chosen (MBPP, HumanEval for code, GSM8K and Math for reasoning) and the authors show clear evidence that the schedules found via Beta tuning significantly reduce sampling steps without reducing quality. The methodology is solid and supported by informative plots (e.g. Fig 5). In addition, showing tasks where their method does not outperform the baseline (GSM8K, BBH) goes a long way in presenting a complete picture, including the limitations. Enough details are given to ensure reproducibility.\nWhile mathematically sound, the authors don’t explicitly discuss the connection between the assumed continuous-time diffusion setting in the theoretical part and the finite-step discrete sampling in real settings. For example, Theorem 3.6 proves the optimality in the continuum and some discussion of how closely the discretized schedule approximates the continuous optimum would strengthen the rigor of the derivation.\n\nClarity:\nThe paper is written clearly, despite its theoretical depth. The introduction motivates the problem and lays out the key points of the paper logically. Fig. 1 is helpful in both summarizing and guiding the reader through Section 3. Notation is consistent and deferring the heavy mathematical proofs to the appendix helps the main text stay focused.\nWhile the notation is mostly standard, expressions like in the conditional kinetic energy definition (Eq. 6) might require readers that are not familiar with OT to study the appendix thoroughly before being able to follow the results and keeping track of the notation. In addition, some key ideas, like the mask schedule invariance of the ELBO objective are mentioned somewhat briefly and more emphasis on such points in the main text could improve the reader’s understanding of why post-training schedule tuning is possible.\n\nSignificance:\nThe paper addresses a significant gap in the understanding of MDMs, which prior to this work showed promising results, but used an ad hoc sampling process based on manually chosen mask schedules without theoretical justification. The contributions of this work significantly advancing the understanding by grounding the sampling of MDM in OT theory. This can guide future design of diffusion-based generative modes as it is possible to reason about MDM behavior using energy minimization and geodesics. And practically, the faster sampling is highly important for most downstream tasks. In addition, not requiring retraining is a big advantage for deployment, because pretrained models can be tuned cheaply for new tasks.\nNevertheless, the empirical performance gains are not universal, as some benchmarks showed no improvements over linear schedules. This shows that further research is needed to gain a better understanding, which the authors acknowledge as a future direction. Nonetheless, the work provides a good foundation for future research to build upon.\n\nOriginality:\nThe contributions of this work are highly novel. While prior work was drawing connections between DMs and OT, to my knowledge this is the first work to interpret MDMs as energy-minimizing transport processes. The closed-form optimal schedule condition was hinted at in earlier analysis (Jo & Hwang, 2025: Continuous Diffusion Model for Language Modeling), but it did not connect it to sampling energy minimization. The proposed Beta-CDF parameterization is a clever choice to simplify the schedule search and the authors show how known schedules are specific points in that space. Related work is cited clearly and the differences are made clear. While this area of research is very active, the paper presents a fresh combination of concepts, extending prior research in non-trivial ways.", "was_revised": true, "final_justification": "I appreciate the detailed responses from the authors, which answered all my questions and concerns. I am happy to raise my rating to “strong accept.”", "rebuttal": "Thank you for your extremely thorough review and acknowledgment of our contributions. Below is our point-by-point response.\n\n### Weakness (Quality) & Question 1:\n\n> While mathematically sound, the authors don’t explicitly discuss the connection between the assumed continuous-time diffusion setting in the theoretical part and the finite-step discrete sampling in real settings.\n> \n> The theorems assume a continuous-time interpolation from 0 to 1. In practice, a discrete sequence of masks must be chosen. Is the sine-squared schedule still exactly optimal in a discretized sense, or only in the continuous limit? Currently this is only done implicitly by tuning Beta parameters for a fixed step count. A brief clarification of the theoretical connection would be helpful.\n\nWe appreciate this insightful question. Indeed, **under Riemann discretization, MDM trajectories still minimizes a corresponding discrete energy functional that converges to the continuous formulation as the number of steps $N \\to \\infty$.**\n\nDue to the equivalence of the three energy functionals, we focus our analysis on the geodesic energy. Recall that $\\mathcal{E}_ g$ can be equivalently expressed as:\n\n$\\mathcal{E}_ g = \\int_{0}^1 \\frac{||\\dot{y}_t||^2}{\\dot{\\gamma_t}} dt.$\n\nTherefore, we can discretize $\\mathcal{E}_g$ as:\n\n$\\mathcal{E}_ g^N = \\sum_{n=0}^{N-1} \\frac{||y_{n+1}-y_{n}||^2 /(\\Delta t)^2}{|\\gamma_{n+1}-\\gamma_{n}| /\\Delta t} \\Delta t = \\sum_{n=0}^{N-1} \\frac{||y_{n+1}-y_{n}||^2 }{|\\gamma_{n+1}-\\gamma_{n}|},$\n\nwhere $\\Delta t=1/N.$ For simplicity, consider the special case where $\\gamma_t = t$. Then we have $|\\gamma_{n+1}-\\gamma_{n}| = \\frac{1}{N}$. We now demonstrate that the $N$ equidistant points along the geodesic still minimize this functional. First, we note that the squared chord length on the unit sphere satisfies $||y_{n+1}-y_{n}||^2 = 2(1-\\cos(\\theta_n))$, where $\\theta_n$ denotes the angle between $y_n$ and $y_{n+1}$. This transforms the minimization problem into finding $N-1$ intermediate points along the great circle arc between $y_0$ and $y_1$ that minimize $\\sum_n 2(1-\\cos(\\theta_n))$.\n\nFrom the properties of spherical geodesics, we know that placing all intermediate points along the geodesic path enforces the constraint $\\sum\\theta_n = \\frac{\\pi}{2}$. Given that $1-\\cos(x)$ is convex over $(0, \\pi/2)$, Jensen's inequality establishes:\n\n$\\sum_n 2(1-\\cos(\\theta_n)) \\geq N\\cdot 2(1-\\cos(\\frac{\\pi}{2N})),$\n\nwith equality achieved when points are uniformly distributed along the geodesic. Any deviation from the geodesic path would further increase the cumulative angular separation $\\sum\\theta_n$ beyond $\\frac{\\pi}{2}$, consequently raising the total energy. This proof extends to arbitrary $\\gamma_t$ schedules through proportional point distribution based on $|\\gamma_{n+1}-\\gamma_{n}|$.\n\nWe will add the above discussion in Section 3.2 of the paper in the final version. We are grateful for this constructive question that allows us to strengthen the theoretical grounding of our approach.\n\n### Weakness (Clarity):\n\n> While the notation is mostly standard, expressions like in the conditional kinetic energy definition (Eq. 6) might require readers that are not familiar with OT to study the appendix thoroughly before being able to follow the results and keeping track of the notation.\n>\n> In addition, some key ideas, like the mask schedule invariance of the ELBO objective are mentioned somewhat briefly and more emphasis on such points in the main text could improve the reader’s understanding of why post-training schedule tuning is possible.\n\nThank you for the constructive suggestions. We will improve the presentation by providing clearer explanations of the definitions in the main text to enhance readability. Also, the ELBO's invariance to mask scheduling indeed plays a crucial role in our framework, therefore we will emphasize this fundamental property more prominently in the main text to improve conceptual clarity.\n\n### Question 2:\n\n> How are the parameters of the Beta-CDF selected for a given task? Did you perform a grid search on (a, b) and pick the best performing schedule? If yes, what are the parameters of this search? It would be useful to know the number of evaluations needed and whether it is feasible to apply to very large models. Discussing the tuning strategy would also help readers to adapt your method more easily and extend this part of it.\n\nWe appreciate this insightful question. Our theoretical analysis suggests that different downstream tasks may require generated text to possess specific intrinsic structures, which in turn necessitates the generation process to emphasize particular temporal phases. These temporal preferences are captured through different $\\gamma_t$ schedules, which induce corresponding optimal $\\alpha_t$ through Condition 3.4. \n\nTherefore, we hypothesize that the schedule preferences are mostly inherent to task nature rather than data specifics. To validate this hypothesis, we conducted the following experiments demonstrating that **randomly chosen small subsets (50-150 instances) of test data suffice for reliable schedule selection.** Specifically, we compare schedule performance between small test subsets and full evaluations (generation length=256, steps=64):\n| Beta-CDF Parameters | (0.5,0.5) | (1,1) | (0.9,0.3) | (0.3,0.9) |\n| :--- | :---: | :---: |:---: | :---: |\n| **task=GSM8K($\\uparrow$)** |||||\n| Random subset 1 (n=132)| **56.06**  | 53.79 | 52.27 | 0.00 |\n| Random subset 2 (n=132)| **55.30**  | 52.27 | 50.76 | 0.76 |\n| Full test set (n=1319) | **52.69** | 50.72 | 51.86 | 0.45 |\n| **task=humaneval($\\uparrow$)** |||||\n| Random subset 1 (n=82) | 17.07 | 35.37 | **43.90** | 3.66 |\n| Random subset 2 (n=82)| 6.10 | 15.85 | **19.51** | 2.44 |\n| Full test set (n=164) | 11.59 | 22.56 | **24.39** | 1.83 |\n\nWhile HumanEval evaluations exhibit greater variance due to smaller test populations (n=164 total), the relative performance rankings remain consistent across subsets - a critical indicator of our method's robustness. This empirical validation confirms that schedule preferences are very probably induced by intrinsic task attributes rather than specific data instances. \n\nTherefore in practice, we recommend conducting initial grid searches using small random test data subsets (about 50~150 instances) across parameters $a,b \\in \\\\{ 0.1, 0.2, ..., 1.0 \\\\} $ to find a set of well-performed schedules. After the coarse search, we can perform finer-grained selection on shortlisted candidates using larger data subsets. This approach achieves comprehensive parameter space exploration while maintaining computational feasibility.\n\nWe will expand this discussion in Section 4 to better guide practitioners in adapting our methodology. We appreciate this valuable question that helps clarify implementation details.\n\n### Weakness (Significance) & Question 3:\n\n> Nevertheless, the empirical performance gains are not universal, as some benchmarks showed no improvements over linear schedules. This shows that further research is needed to gain a better understanding, which the authors acknowledge as a future direction.\n>\n> Do you have hypotheses why GSM8K and BBH did not benefit from optimized Beta schedules? Is it because those tasks require very specific refinement at the end, so a linear schedule is already as good as it gets? Providing some discussion on these cases would be valuable.\n\nOur theoretical analysis suggests that different tasks may prefer different schedules, therefore the observed parity with linear baselines in certain tasks aligns with our theoretical framework.\n\nSpecifically, the linear schedule $\\alpha_t = t$ corresponds to $\\gamma_t=CDF(\\mathcal{B}(0.5,0.5))$, which has higher derivatives on both initial and final generation phases while being relatively static in the middle. Since $\\dot{\\gamma}_t$ appears in the denominator of our energy functional expression, this implies that tasks requiring sustained refinement throughout generation (particularly middle phases) rather than the starting or ending phases might inherently favor the linear schedule.\n\nAmong the six tasks we experimented on, BBH is only one that focuses on general reasoning problems. As for GSM8K, although it is mathematics-focused, its prompt requires the generated text to use detailed natural language explanations compared to other math benchmarks. This might be the reason why sustained refinement in the middle phases are preferred. Considering this similarity, it is reasonable that GSM8K and BBH have similar preference on schedules although it is difficult to mathematically deduce the exact expression of the energy functionals. \n\nConsidering the complexity of real-world tasks, it is indeed a challenging direction to systematically study the relationship between tasks and their prefered schedules. A promising research direction involves designing intermediate benchmarks that balance practical relevance with analytical tractability - more sophisticated than our current toy examples yet simpler than real-world tasks.\n\nWe will incorporate this analysis in Section 4 of our paper in the final version. We appreciate this insightful line of questioning and we hope our response addresses your concerns. Further discussion is welcomed if needed."}, {"review_id": "S1hqf8BCg0", "reviewer": "Reviewer_nSvp", "summary": "This paper presents a theoretical framework interpreting Masked Diffusion Models (MDMs) as minimizing energy functionals rooted in discrete optimal transport. It establishes that three distinct formulations of transport cost—kinetic energy, conditional kinetic energy, and geodesic energy—are mathematically equivalent within the MDM structure. A key result is the derivation of a closed-form optimal mask schedule, showing that MDMs minimize these energies under specific conditions.", "questions": "I only have one questions: the proposed method demonstrates strong improvements in few-step sampling on HumanEval and Hendrycks Math. However, the gains are less significant or inconsistent on other benchmarks like BBH and GSM8K. Can you elaborate on why the proposed energy-based schedule performs better on some tasks than others?", "limitations": "Yes.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Strengths: The paper offers a highly novel theoretical interpretation of MDMs, unifying different energy functionals and drawing on principles from information geometry and optimal transport. A major contribution is the identification of a closed-form condition that governs optimal masking schedules, bridging discrete probability flows and geodesic trajectories. The method is tested on strong baselines, including LLaDA 8B, showing significant improvements in low-step sampling efficiency on tasks like HumanEval and Hendrycks Math.\n\nweaknesses: While experiments show improved sample efficiency, the performance gains are limited to a subset of benchmarks (notably, not significant on BBH or GSM8K).", "was_revised": true, "final_justification": "All of my concerns have been addressed and so I raised my score.", "rebuttal": "Thank you for recognizing our contributions and providing thoughtful feedback. Please find our response below.\n\n### Weakness & Question:\n\n> While experiments show improved sample efficiency, the performance gains are limited to a subset of benchmarks (notably, not significant on BBH or GSM8K)\n>\n> The proposed method demonstrates strong improvements in few-step sampling on HumanEval and Hendrycks Math. However, the gains are less significant or inconsistent on other benchmarks like BBH and GSM8K. Can you elaborate on why the proposed energy-based schedule performs better on some tasks than others?\n\nWe appreciate this insightful question. We clarify that our methodology provides both a theoretical foundation for understanding schedule preferences and a practical mechanism for task-specific optimization, rather than proposing a universally superior schedule.\n\n**In fact, linear schedule $\\alpha_t = t$ is a special case in our framework since it corresponds to $\\gamma_t=CDF(\\mathcal{B}(0.5,0.5))$, a point in our parameter space.** Therefore, the empirical parity on certain benchmarks confirms that linear schedules already serve as near-optimal candidates for specific task categories. However, our framework's value lies in enabling systematic task-specific schedule optimization, overcoming the limitations of previous one-size-fits-all approaches.\n\nSpecifically, the linear schedule $\\alpha_t = t$ corresponds to $\\gamma_t=CDF(\\mathcal{B}(0.5,0.5))$, which has higher derivatives on both initial and final generation phases while being relatively static in the middle. Since $\\dot{\\gamma}_t$ appears in the denominator of our energy functional expression, this implies that tasks requiring sustained refinement throughout generation (particularly middle phases) rather than the starting or ending phases might inherently favor the linear schedule.\n\nAmong the six tasks we experimented on, BBH is only one that focuses on general reasoning problems. As for GSM8K, although it is mathematics-focused, its prompt requires the generated text to use detailed natural language explanations compared to other math benchmarks. This might be the reason why sustained refinement in the middle phases are preferred. Considering this similarity, it is reasonable that GSM8K and BBH have similar preference on schedules although it is difficult to mathematically deduce the exact expression of the energy functionals. \n\nConsidering the complexity of real-world tasks, it is indeed a challenging direction to systematically study the relationship between tasks and their prefered schedules. A promising research direction involves designing intermediate benchmarks that balance practical relevance with analytical tractability - more sophisticated than our current toy examples yet simpler than real-world tasks.\n\nWe will incorporate this analysis in Section 4 of our paper in the final version. We appreciate this insightful question and we hope our response addresses your concerns. Further discussion is welcomed if needed."}, {"review_id": "gPXwSkprlZ", "reviewer": "Reviewer_5LPg", "summary": "The paper proposes a unified theoretical framework that interprets masked diffusion models (MDMs) as solutions to discrete optimal transport problems. Within this framework, three energy formulations—kinetic, conditional kinetic, and geodesic—are shown to be mathematically equivalent in the sense that their corresponding losses are proportional under the MDM structure. Leveraging this insight, particularly the geodesic perspective, the paper introduces an optimal mask schedule with analytically derived coefficients along the diffusion trajectory. This leads to practical sampling improvements while deepening the theoretical understanding of MDMs. Experiments on both synthetic and real-world benchmarks demonstrate that the proposed energy-inspired schedules consistently outperform hand-crafted baselines, particularly in low-step sampling regimes.", "questions": "Please refer to the Weaknesses section above for Questions.", "limitations": "Please refer to the Weaknesses section above for Limitations.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Overall, the main strength of this paper lies in its unified theoretical framework for understanding MDMs through kinetic, conditional kinetic, and geodesic formulations. Additionally, the proposed masking schedule is well-motivated, as it aligns with the geodesic properties established by the theoretical analysis. Finally, the schedule also demonstrates practical improvements in MDM sampling performance. However, I list some potential weaknesses of the paper below:\n- Although the study introduces a kinetic viewpoint, the proposed masking schedule relies heavily on geodesic insights already explored in [1]. This makes the kinetic formulation feel more like a minor addition, with the main contribution closely following prior work—thereby limiting the novelty of this work.\n- The use of the Beta CDF as the masking schedule $\\gamma$ is reasonable, as it generalises existing schedules such as linear and squared sine. It also enables inference-time optimisation in MDMs. However, the choice of the Beta distribution parameters $a, b$ appears to be heuristic. Could the authors elaborate on how to systematically search for or optimise $a, b$ rather than relying on trial and error? Additionally, how can varying $a, b$ values lead to improvements over the linear and squared sine schedules in certain scenarios (especially in low-step sampling regimes)? Any insights or empirical observations would be appreciated.\n\nIn my view, the unified theoretical framework appears somewhat disconnected from the proposed schedule. The most practical contribution lies in the new masking schedule, which is primarily inspired by prior work—limiting the overall novelty. Furthermore, the choice of parameters $a, b$ in the Beta distribution across different settings lacks sufficient justification. I currently lean towards a borderline reject, though I may revise my evaluation depending on the authors’ responses.\n\n[1] Jaehyeong Jo and Sung Ju Hwang. Continuous Diffusion Model for Language Modeling. 2025.\n\n**[After Rebuttal] I thank the authors for their responses. My concerns have been addressed, and I am happy to revise my evaluation from 3 to 4.**", "was_revised": true, "final_justification": "I thank the authors for their responses. My concerns have been addressed, and I am happy to revise my evaluation from 3 to 4.", "rebuttal": "Thank you for your careful review and thoughtful comments. Below is our point-by-point response.\n\n### Weakness 1:\n\n> Although the study introduces a kinetic viewpoint, the proposed masking schedule relies heavily on geodesic insights already explored in [1]. This makes the kinetic formulation feel more like a minor addition, with the main contribution closely following prior work—thereby limiting the novelty of this work.\n>\n> The unified theoretical framework appears somewhat disconnected from the proposed schedule. The most practical contribution lies in the new masking schedule, which is primarily inspired by prior work—limiting the overall novelty.\n\nWe appreciate this critical perspective. We would like to clarify that **Our framework fundamentally operates within the optimal transport paradigm for efficient sampling beyond the geodesic perspective.** As optimal transport connects directly with efficient sampling [1*,2*], we establish provably optimal schedules for sampling efficiency in masked diffusion models. More importantly, unlike DFM's decoupled optimization of $p_t$ and $Q_t$, MDM's unique architecture – where both $p_t$ and $Q_t$ are jointly determined by $\\alpha_t$ – natually raises fundamental questions about  whether such a constrained setting still can reach optimality. Our analysis demonstrates that MDM intrinsically encodes optimal rate matrices within its constrained structure. This theoretical guarantee strengthens confidence in MDM's elegant framework, particularly given its empirical success. **Such a perspective is absent in [1].**\n\n**Technically, our core objective of our analysis is to minimize the kinetic and conditional kinetic energy, which cannot be interpreted easily within the geodesic perspective.** Therefore, the geodesic method only gives us a mathemetical tool rather than a core goal. By utilizing the geodesic perspective, our main theorems are successfuly proven and understanding of MDM's capabilities are deepened. As noted by Reviewer dGWj, such synthesis \"presents a fresh combination of concepts, extending prior research in non-trivial ways.\" \n\n**Even within the geodesic perspective, it is unclear how the pure geometric method in [1] can be used to prove the optimality of Condition 3.4.** By introducing the geodesic energy, our proof of Lemma 3.5 is proved using functional minimization methods, not only extending the original proof, but also successfully connected to our main goal in optimal transport.\n\nTherefore, although the closed-form condition was hinted in [1], our theoretical contributions beyond prior work is still non-trivial. We will clarify these points in the revised version and we hope that this response has addressed your question.\n\n[1*] Cédric Villani. Topics in optimal transportation. Vol. 58. American Mathematical Soc., 2021.\n\n[2*] Neta Shaul et al. “On Kinetic Optimal Probability Paths for Generative Models”. In: International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA. Ed. by Andreas Krause et al. Vol. 202. Proceedings of Machine Learning Research. PMLR, 2023, pp. 30883–30907.\n\n### Weakness 2:\n\n> The use of the Beta CDF as the masking schedule $\\gamma$ is reasonable, as it generalises existing schedules such as linear and squared sine. It also enables inference-time optimisation in MDMs. However, the choice of the Beta distribution parameters $a, b$ appears to be heuristic. Could the authors elaborate on how to systematically search for or optimise $a, b$ rather than relying on trial and error? \n>\n> The choice of parameters $a, b$ in the Beta distribution across different settings lacks sufficient justification.\n\nWe appreciate this insightful question. Our theoretical analysis suggests that different downstream tasks may require generated text to possess specific intrinsic structures, which in turn necessitates the generation process to emphasize particular temporal phases. These temporal preferences are captured through different $\\gamma_t$ schedules, which induce corresponding optimal $\\alpha_t$ through Condition 3.4. \n\nTherefore, we hypothesize that the schedule preferences are mostly inherent to task nature rather than data specifics. To validate this hypothesis, we conducted the following experiments demonstrating that **randomly chosen small subsets (50-150 instances) of test data suffice for reliable schedule selection.** Specifically, we compare schedule performance between small test subsets and full evaluations (generation length=256, steps=64):\n| Beta-CDF Parameters | (0.5,0.5) | (1,1) | (0.9,0.3) | (0.3,0.9) |\n| :--- | :---: | :---: |:---: | :---: |\n| **task=GSM8K($\\uparrow$)** |||||\n| Random subset 1 (n=132)| **56.06**  | 53.79 | 52.27 | 0.00 |\n| Random subset 2 (n=132)| **55.30**  | 52.27 | 50.76 | 0.76 |\n| Full test set (n=1319) | **52.69** | 50.72 | 51.86 | 0.45 |\n| **task=humaneval($\\uparrow$)** |||||\n| Random subset 1 (n=82) | 17.07 | 35.37 | **43.90** | 3.66 |\n| Random subset 2 (n=82)| 6.10 | 15.85 | **19.51** | 2.44 |\n| Full test set (n=164) | 11.59 | 22.56 | **24.39** | 1.83 |\n\nWhile HumanEval evaluations exhibit greater variance due to smaller test populations (n=164 total), the relative performance rankings remain consistent across subsets - a critical indicator of our method's robustness. This empirical validation confirms that schedule preferences are very probably induced by intrinsic task attributes rather than specific data instances. \n\nTherefore in practice, we recommend conducting initial grid searches using small random test data subsets (about 50~150 instances) across parameters $a,b \\in \\\\{0.1, 0.2, ..., 1.0\\\\}$ to find a set of well-performed schedules. After the coarse search, we can perform finer-grained selection on shortlisted candidates using larger data subsets. This approach achieves comprehensive parameter space exploration while maintaining computational feasibility.\n\nWe will expand this discussion in Section 4 to better guide practitioners in adapting our methodology. We appreciate this valuable question that helps clarify implementation details.\n\n\n### Weakness 3:\n\n> Additionally, how can varying $a, b$ values lead to improvements over the linear and squared sine schedules in certain scenarios (especially in low-step sampling regimes)? Any insights or empirical observations would be appreciated.\n\nThank you for this insightful question. As mentioned in the response to Weakness 2, different intrinsic structures of tasks may induce different weight function $\\gamma_t$, therefore inducing different $\\alpha_t$ preferences. Since $\\dot{\\gamma}_t$ appears in the denominator of our energy functional expression, temporal phases where $\\gamma_t$ has larger derivatives get less weight and phases where $\\gamma_t$ is relatively flat get more focused. By varying the parameter $(a,b)$, we fundamentally alter the schedule's concavity and inflection points, therefore changing the temporal focus. When the temporal focus is aligned with intrinsic demand of the specific task, sampling performance is thereby promoted.\n\nTake the first toy experiment in Figure 4 as an illustrative case. Since the target distribution only have positive probability on extreme points (all zeros or all ones), precise trajectory control becomes paramount in initial generation phases to prevent simultaneous 0-1 generation. As the generation proceed to the ending period, however (for example 4 zeros have already been unmasked), the distribution itself (which is parameterized by a neural network in real world experiments) will prohibit the generation of ones since it will lead to invalid combinations. This eliminates the need for intensive late-stage focus, enabling efficient resource allocation to the decisive initial phases. Such temporal preference make the model prefer the parameter (0.6,0.3), where the starting period of generation is more weighted.\n\nIn low-step regimes, strategic allocation of limited computational resources becomes more critical. This explains why our tuning method lead to improvements especially under low step sampling scenarios.\n\nWe hope that our response above has addressed your concerns, and we welcome further discussion if needed."}, {"review_id": "Kl9hlQQy2o", "reviewer": "Reviewer_2RNf", "summary": "In this paper, the authors worked on Masked diffusion models(MDM) for discrete data. They proposed  to define three different energy formulations—kinetic, conditional kinetic, and geodesic energy, and then prove they are mathematically equivalent, thus minimize all three with the same  optimal sampling schedules that based on a Beta-CDF parameterization which they believe enables efficiency.", "questions": "The paper is difficult to read. The following are some of my questions:\n1. It seems there is some mistake when the authors introduced the MDM. Page 2: line 62: The authors mentioned: \"a masking process (operating backward in time, t = 1 → 0) and its reverse, the unmasking process (t = 0 → 1) ...\", this seems opposite to what was in the previous literature where forward diffusion is masking and backward is unmasking. As a result,  in  Line 57: The authors assume p0(z) follows a uniform distribution, but they did not specify what is the target distribution p_1(z). Again, this contradicts with existing literatures. For example, ref[8] of the paper discussed the forward diffusion to a uniform distribution; and ref [43] works on the case that x_1 follows a Dirac distribution.\n\n2. Page2: Eqns. (1) (2) did not make sense with (z not equal to x) in between the two lines, seems like a mistake.\n\n3. The paper did not explicitly mentioned what the target distribution is; Figure 4 showed some, but did not explicitly explained it.", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "It is interesting to see they proposed energy functions for characterizing MDM; however, the optimal sampling schedule defined to minimize the  the three formulations that the authors proposed sounds not to improve anything because the authors also showed the \"Invariance of Training Loss to the Mask Schedule\" in Appendix C. So, I wonder if the three are equivalent formation of the same concept.\n\nAlso, in the three definitions, what if the denominator, say in the eqn. (5), equals 0? I am further concerned when I read Page6, Figure 2: according to the pictures, the color indicates energy values, it seems like discontinuous for Beta(a,b). This suggested the energy function defined did not consider possible 0 in the denominator.\n\nWhile authors focused on Q_t(z, x), I suggest they should provide p_t(z, x) when defining MDM. Also, they use the model in ref[43] for experiments; Further, it seems there is some issue with their definition of MDM sees not the same, see the section of Questions.", "was_revised": true, "final_justification": "Thanks to the authors answering my questions. I would like to raise my rating to 5.", "rebuttal": "Thank you for the recognition of our contributions and the thoughtful comments. Below is our point-by-point response.\n\n### Weakness 1:\n\n> However, the optimal sampling schedule defined to minimize the the three formulations that the authors proposed sounds not to improve anything because the authors also showed the \"Invariance of Training Loss to the Mask Schedule\" in Appendix C. So, I wonder if the three are equivalent formation of the same concept.\n\nWe appreciate this important question. The three energy formulations proposed in our work quantify the **sampling** process's \"energy cost\" from distinct perspectives, which fundamentally differ from the **training** loss function. To be more specifically, The training phase aims to learn a neural network approximation of the target distribution whereas the sampling phase addresses the operational trajectory for sample generation from the learned distribution. \n\nTherefore, while Appendix C demonstrates training loss invariance to mask schedules – enabling post-training schedule tuning – the sampling process remains highly schedule-dependent, as evidenced by our empirical results.\n\nWe will adapt the above clarification in the main text of the final version to make our theoretical framework clearer.\n\n### Weakness 2:\n\n> Also, in the three definitions, what if the denominator, say in the eqn. (5), equals 0?\n\nWe thank you for this meticulous observation. **We would like to clarify that such cases are natually excluded from the summation, therefore does not affect our framework.** Consider the kinetic energy definition:\n\n$$\\mathcal{E}_ k(p_t,Q_t;\\gamma_t)=\\mathbb{E}_ {t,p_t(z)}\\sum_ {x:x\\neq z}\\frac{1}{\\dot{\\gamma}_tp_t(x)}Q_t(z,x)^2.$$\n\nTo calculate this formulation, we first sample $z$ that have positive probability, i.e. $p_t(z)>0.$ Due to the detailed balance condition, as long as $Q_t(z,x)>0$, we must have $p_t(x)>0$ due to the continuity of the probability distribution. Therefore, the case $p_t(x)=0$ can only occur when $Q_t(z,x)=0$.\n\nSince the summation is of all probability flows between states, such $(z,x)$ pairs carrying zero transition probability ($Q_t(z,x) = 0$) are inherently and natually excluded from the summation, as it suggests that the flow between $z$ and $x$ does not exist at all. Our theoretical framework and proofs implicitly incorporate this exclusion principle throughout the analysis. We acknowledge the need to explicitly state this detail and we will revise accordingly in our final version in caption to enhance clarity. However, such concern does not affect our theoretical framework. Thank you for highlighting this important presentation consideration.\n\n### Weakness 3:\n\n> I am further concerned when I read Page6, Figure 2: according to the pictures, the color indicates energy values, it seems like discontinuous for Beta(a,b). This suggested the energy function defined did not consider possible 0 in the denominator.\n\nWe appreciate your careful examination. The apparent discontinuity in Figure 2 stems from our discrete parameter sampling strategy for visualization purposes rather than inherent energy function discontinuity. Energy values actually vary continuously with respect to $(a,b)$ parameters. Figure 2 is to emphasize the relationship between different weight functions and corresponding energy minima positions, therefore we only selected sparsely distributed $(a,b)$ pairs to show the coarse relationship. As addressed in our previous response regarding denominator zeros, this potential concern does not affect the theoretical validity or experimental results presented.\n\n### Question 1:\n\n> It seems there is some mistake when the authors introduced the MDM. Page 2: line 62: The authors mentioned: \"a masking process (operating backward in time, t = 1 → 0) and its reverse, the unmasking process (t = 0 → 1) ...\", this seems opposite to what was in the previous literature where forward diffusion is masking and backward is unmasking. As a result, in Line 57: The authors assume p0(z) follows a uniform distribution, but they did not specify what is the target distribution p_1(z). Again, this contradicts with existing literatures. For example, ref[8] of the paper discussed the forward diffusion to a uniform distribution; and ref [43] works on the case that x_1 follows a Dirac distribution.\n\nThank you for this insightful observation. Our temporal parameterization follows conventions from energy-based sampling literature [1*,2*] rather than standard diffusion frameworks. Specifically: $t=0$ denotes the initial sampling state (fully masked Dirac distribution in MDM) and $t=1$ marks the final generated state.\n\nThis orientation aligns the forward temporal direction with the unmasking process – the core focus of our energy analysis. The reversed timeline offers two key advantages: 1) It simplifies energy functional expressions by synchronizing time progression with generation advancement, and 2) Ensures mask/interpolation schedules naturally increase from 0 to 1. While differing from some diffusion literature, this convention maintains internal consistency and better serves our analytical framework's objectives.\n\n[1*] Neta Shaul et al. “On Kinetic Optimal Probability Paths for Generative Models”. In: International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA. Ed. by Andreas Krause et al. Vol. 202. Proceedings of Machine Learning Research. PMLR, 2023, pp. 30883–30907.\n\n[2*] Neta Shaul et al. “FlowMatching with General Discrete Paths: A Kinetic-Optimal Perspective”. In: CoRR abs/2412.03487 (2024). DOI: 10.48550/ARXIV.2412.03487. arXiv: 2412.03487.\n\n\n### Weakness 4 & Question 3:\n\n> While authors focused on Q_t(z, x), I suggest they should provide p_t(z, x) when defining MDM.\n>\n> The paper did not explicitly mentioned what the target distribution is; Figure 4 showed some, but did not explicitly explained it.\n\nWe appreciate this technical clarification. In continuous-time Markov chains, the marginal distribution p_t(z) evolves via the Kolmogorov forward equation $\\frac{dp}{dt} = p_tQ_t$. For MDM's specific $Q_t(z,x)$ structure, $p_t(z)$ generally lacks closed-form solutions but remains fully determined through $Q_t$ specifications. Therefore following MDM traditions [1*,2*], we only specify $Q_t(z,x)$ and $p_t(z)$ is then fully determined.\n\nAs for the target distribution, our main experiments focus on practical scenarios where target distributions correspond to task-specific data distributions without explicit analytical forms. As for the synthetic experiment in Figure 4, the target distributions indeed have analytical forms and details are presented in Appendix E.4.\n\n[1*] Jacob Austin et al. “Structured denoising diffusion models in discrete state-spaces”. In: Advances in Neural Information Processing Systems 34 (2021), pp. 17981–17993.\n\n[2*] Jingyang Ou et al. “Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data”. In: ICLR. 2025.\n\n### Question 2:\n\n> Page2: Eqns. (1) (2) did not make sense with (z not equal to x) in between the two lines, seems like a mistake.\n\nThank you for this meticulous observation. The $z \\neq x$ condition rigorously reflects continuous-time Markov chain conventions [1*]. Off-diagonal elements $Q_t(z,x) (z \\neq x)$ & represent transition rates between distinct states, while diagonal elements satisfy $Q_t(z,z) = -\\sum_{x\\neq z}Q_t(z,x)$ to conserve probability mass. Explicitly specifying $z \\neq x$ ensures mathematical precision when defining transition dynamics, as diagonal elements derive from off-diagonal terms rather than being independently specified.\n\n[1*] Jingyang Ou et al. “Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data”. In: ICLR. 2025.\n\nWe hope that our response above has addressed your concerns, and we welcome further discussion if needed."}, {"review_id": "h88rhpEBSV", "reviewer": "Reviewer_gvvr", "summary": "The authors propose to examine the kinetic energy of Masked Diffusion Models and to explore the influence of the masking process schedule on it. Three formulations of kinetic energy were proposed and it was proven that they are equivalent up to a multiplicative constant and can be simultaneously minimized w.r.t. masking schedule. Given this insight, an explicit masking schedule solution for the mask kinetic energy minimization was derived. However, this explicit masking schedule solution depends on the time-dependent energy weights $\\gamma_t$ that define the kinetic energy functional.\n\nOn the practical side of the paper, the authors suggest to parameterize the time-dependent energy weights  $\\gamma_t$ as a CDF of the Beta distribution then optimize over the Beta distribution parameters and receive the masking schedule as an explicit kinetic energy minimization solution given the optimized time-dependent energy weights  $\\gamma_t$. The practical algorithm is tested on LLaDA 8B [1] masked diffusion model in a post-training manner on a set of general reasoning, mathematical problem solving, and code generation benchmarks.", "questions": "- Although I think the presented benchmarks are relatively solid it would be nice to the see evaluation of the proposed method for the optimization of masking schedules of basic MDLM [2] on unconditional generation problems and measure test ppl (e.g. Figure 1 [2]).\n- It would be nice to see the full end-to-end algorithm and corresponding evaluations for tuning the masking schedule using the proposed parameterization.\n- In both subfigures at Figure 2 there are quite sharp bound after which energy start to drastically grow (green to purple edge). Why is that?\n- In Figure 2 caption the following said “Axes represent the beta-parameterization of $\\alpha_t$ (see Sec. 3.3)”. What exactly does it mean, what is beta-parameterization of $\\alpha_t$ ? I could not find such information in Sec. 3.3\n- In Def 2.1 the *unmasking* process rate matrix is used, while as I understand it there also could be *masking* process rate matrix, am I right? If that is the case and in general these energy formulations (Def 2.1, Def 2.2, Def 2.3) can be defined in both ways (forward and backward) I think it should be mentioned somewhere in the main part of the paper.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n\n- Novel theoretical insights. Interesting and novel theory investigations of the kinetic energy in the light of masking schedule. Closed-form masking schedule minimizer for the kinetic energy gives valuable insight into the design of masked diffusion.\n- Simple yet effective method. The suggested method for the post-training schedule change is rather simple and, according to experimental results, effective.\n- Relatively robust evaluation. The experimental evaluation is done on the state-of-the-art model LLaDA 8B [1] and six diverse benchmarks, which seem to be a relatively robust evaluation.\n\nWeaknesses:\n\n- Unnecessarily detailed theory. A lot of space in the paper is dedicated to kinetic energy and related theory but in the end, the conclusion is rather simple. Not sure that all that theory is particularly needed to be presented in the main part of the paper given that it doesn’t provide any other insight besides optimal masking schedule form.\n- Shifted complexity. The optimal schedule still hinges on time-dependent energy weights $\\gamma_t$, effectively transferring the optimization challenge from  $\\alpha_t$ to $\\gamma_t$ without simplifying the overall problem.\n- Missing comparison of parameterizations. In the light of previous weakness, it would be nice to see the comparison of the $\\alpha_t$ parametrization as an explicit minimizer of the kinetic energy given energy weights $\\gamma_t$, which in turn are parametrized as CDF of Beta distribution and parametrization of $\\alpha_t$ directly as CDF of Beta distribution.\n- Limited impact scope. With contributions split between a bit of theoretical material and a rather simple post-training optimization of masking schedule, I think the overall paper impact may be limited.\n- Lack of a practical algorithm. Seems like the paper explores how various schedules influence MDM inference but stops short of delivering a concrete end-to-end procedure for optimizing the proposed masking schedule $\\alpha_t$ parameterization. Would be nice to see one.\n\n[1] Nie, S., Zhu, F., You, Z., Zhang, X., Ou, J., Hu, J., ... & Li, C. (2025). Large language diffusion models. *arXiv preprint arXiv:2502.09992*.\n\n[2] Sahoo, S., Arriola, M., Schiff, Y., Gokaslan, A., Marroquin, E., Chiu, J., ... & Kuleshov, V. (2024). Simple and effective masked diffusion language models. *Advances in Neural Information Processing Systems*, *37*, 130136-130184.", "was_revised": true, "final_justification": "The authors have addressed most of my concerns.  However, there is still a gap between theory and practice, since authors have not explored counterparametrizations, not involving $\\gamma_t$, and hence it is not particularly clear wherever the energy parameterization is necessary for good practical performance. Overall, I find the theory a bit unconvincing, but the practical benefits are valuable regardless.", "rebuttal": "Thank you for your extremely thorough review and constructive suggestions. Below is our point-by-point response.\n\n### Weakness 1: Unnecessarily detailed theory \n\nWe appreciate this critical perspective. We would like to clarify that although our optimal schedule induced by Condition 3.4 is rather simple, our theoretical framework serves crucial purposes beyond deriving the optimal schedule.\n\nSpecifically, MDM's constrained architecture - where both $p_t$ and $Q_t$ are simultaneously determined by $\\alpha_t$ - naturally raises questions about its capacity to achieve optimal transport compared to DFM's decoupled optimization. Through Theorem 3.6, we demonstrate that MDM intrinsically encodes optimal rate matrices within its structural limitations. This theoretical guarantee resolves fundamental doubts about MDM's capabilities while explaining its empirical success, thereby strengthening confidence in the framework's mathematical soundness.\n\nTherefore, we believe that our theoretical analysis establishes foundational insights into MDM's operational mechanisms, being helpful in guiding future works.\n\n### Weakness 2 & 3: Shifted complexity & Missing comparison of parameterizations\n\nThank you again for the critical perspective. We would like to clarify the importance of introducing $\\gamma_t$ empirically. \n\nFirst, The $\\alpha_t$-$\\gamma_t$ duality provides **complementary perspectives** for schedule design. Common $\\alpha_t$s correspond to complex $\\gamma_t$s without closed-form expressions (e.g. $\\alpha_t=\\sin(\\frac{\\pi}{2}t)$ corresponds to $\\gamma_t$ that do not have a closed form). Conversely, Beta-CDF parameterized $\\gamma_t$s generally yield non-trivial $\\alpha_t$s (e.g. $\\gamma_t=CDF(\\mathcal{B}(0.9,0.3))$ correspond to nontrivial $\\alpha_t$).\n\nMore importantly, **our empirically effective Beta-CDF tuning method stems directly from the energy perspective**. The discovery that standard schedules (linear/sine-squared) correspond to CDF of Beta distribution (Proposition 3.7) motivated systematic Beta-CDF exploration. This connection remains obscure when operating purely in $\\alpha_t$ space.\n\nWe appreciate your idea in directly parametrizing $\\alpha_t$, and although such parameterization might be able to yield empirical improvements as well, the introduction of $\\gamma_t$ is still indispensible due to the reasons above.\n\nWe hope that our clarification above has addressed your problem.\n\n### Weakness 4: Limited impact scope\n\nAs clarified in our response to weakness 1-3, despite our rather simple closed-form expression, our theoretical analysis induces unique and nontrivial insights. Also, our effective schedule tuning method, especially the Beta-CDF parameterizing method cannot be done without our theoretical framework. Moreover, as detailed in the response to your next question, we have proposed a more thorough explanation on how to effectively tune the parameters in real-world applications. Such method are also highly motivated by our energy perspective.\n\nConsidering these clarifications, we kindly ask you to reassess the significance of our work.\n\n### Weakness 5 & Question 2: Lack of a practical algorithm\n\nWe appreciate this insightful question. Our theoretical analysis suggests that different downstream tasks may require generated text to possess specific intrinsic structures, which in turn necessitates the generation process to emphasize particular temporal phases. These temporal preferences are captured through different $\\gamma_t$ schedules, which induce corresponding optimal $\\alpha_t$ through Condition 3.4. \n\nTherefore, we hypothesize that the schedule preferences are mostly inherent to task nature rather than data specifics. To validate this hypothesis, we conducted the following experiments demonstrating that **randomly chosen small subsets (50-150 instances) of test data suffice for reliable schedule selection.** Specifically, we compare schedule performance between small test subsets and full evaluations (generation length=256, steps=64):\n| Beta-CDF Parameters | (0.5,0.5) | (1,1) | (0.9,0.3) | (0.3,0.9) |\n| :--- | :---: | :---: |:---: | :---: |\n| **task=GSM8K($\\uparrow$)** |||||\n| Random subset 1 (n=132)| **56.06**  | 53.79 | 52.27 | 0.00 |\n| Random subset 2 (n=132)| **55.30**  | 52.27 | 50.76 | 0.76 |\n| Full test set (n=1319) | **52.69** | 50.72 | 51.86 | 0.45 |\n| **task=humaneval($\\uparrow$)** |||||\n| Random subset 1 (n=82) | 17.07 | 35.37 | **43.90** | 3.66 |\n| Random subset 2 (n=82)| 6.10 | 15.85 | **19.51** | 2.44 |\n| Full test set (n=164) | 11.59 | 22.56 | **24.39** | 1.83 |\n\nWhile HumanEval evaluations exhibit greater variance due to smaller test populations (n=164 total), the relative performance rankings remain consistent across subsets - a critical indicator of our method's robustness. This empirical validation confirms that schedule preferences are very probably induced by intrinsic task attributes rather than specific data instances. \n\nTherefore in practice, we recommend conducting initial grid searches using small random test data subsets (about 50~150 instances) across parameters $a,b \\in \\\\{0.1, 0.2, ..., 1.0\\\\}$ to find a set of well-performed schedules. After the coarse search, we can perform finer-grained selection on shortlisted candidates using larger data subsets. This approach achieves comprehensive parameter space exploration while maintaining computational feasibility.\n\nWe will expand this discussion in Section 4 to better guide practitioners in adapting our methodology. We appreciate this valuable question that helps clarify implementation details.\n\n### Question 1:\n\n> Although I think the presented benchmarks are relatively solid it would be nice to the see evaluation of the proposed method for the optimization of masking schedules of basic MDLM [2] on unconditional generation problems and measure test ppl (e.g. Figure 1 [2]).\n\nThank you for this question. Our extended experiments demonstrate the broader applicability of our schedule optimization framework to unconditional generation tasks. We still conducted evaluations using Llada due to its unprecedented 8B scale with generation length fixed at 256 tokens and sampling steps at 64. Perplexity measurements were computed via Quen2-7B ([1*]) on 100 generated samples per schedule configuration:\n\n| Beta-CDF Parameters | (0.5,0.5) | (1,1) | (0.9,0.3) | (0.3,0.9) |\n| :--- | :---: | :---: |:---: | :---: |\n| test ppl ($\\downarrow$) | 408.62 | 187.25 | **169.98** | 4874.87 \n\nThe (0.3,0.9) parameter configuration exhibits catastrophic performance degradation (PPL >4800), producing semantically incoherent outputs that align with its poor performance across multiple tasks in our main experiments. This consistency underscores our method's robustness in identifying non-viable schedules across different generation paradigms.\n\nWe hope that our experiment has addressed your problem and further discussion is welcomed.\n\n[1*] Qwen Team. \"Qwen2.5: A Party of Foundation Models\". 2024.\n\n### Question 3:\n\n> In both subfigures at Figure 2 there are quite sharp bound after which energy start to drastically grow (green to purple edge). Why is that?\n\nOur computational analysis reveals that schedules deviating significantly from theoretical optima may induce the energy functionals to dirverge. The Beta distribution parameters systematically govern critical schedule characteristics including concavity-convexity properties and inflection point locations. When these parameters stray beyond critical thresholds from their optimal configurations, the energy functionals will may diverge to infinity.\n\nOur large-scale experiments corroborate this phenomenon: certain schedule may exhibits catastrophic performance degradation on specific tasks while maintaining normal functionality on others. This bifurcation originates from the energy landscape's acute sensitivity to $\\gamma_t$ variations. Different $\\gamma_t$ schedules create distinct energy topographies where non-optimal configurations reside in high-energy areas, leading to suboptimal generation performance.\n\n### Question 4:\n\n> In Figure 2 caption the following said “Axes represent the beta-parameterization of $\\alpha_t$ (see Sec. 3.3)”. What exactly does it mean, what is beta-parameterization of $\\alpha_t$ ? I could not find such information in Sec. 3.3\n\nThank you for highlighting this presentation gap. The beta-parameterization refers to constructing $α_t$ schedules through the Beta CDF: $α_t = CDF(\\mathcal{B}(a,b))(t)$, where $(a,b)$ parameterize the Beta distribution. Figure 2 illustrates how these parameters influence energy landscapes under fixed $\\gamma_t$ schedules through Condition 3.4. Each axis coordinate corresponds to specific $(a,b)$ pairs that generate distinct $α_t$ via this CDF mapping. The visualization demonstrates how Condition 3.4-guided parameter selection identifies energy minima. We will clarify this methodology explicitly in revisions to prevent misinterpretation.\n\n### Question 5:\n\n> In Def 2.1 the unmasking process rate matrix is used, while as I understand it there also could be masking process rate matrix, am I right? If that is the case and in general these energy formulations (Def 2.1, Def 2.2, Def 2.3) can be defined in both ways (forward and backward) I think it should be mentioned somewhere in the main part of the paper.\n\nThank you for your suggestion. While energy formulations could indeed be constructed for both masking and unmasking processes theoretically, our framework deliberately focuses on the unmasking rate matrix since the unmasking process directly governs the practical sampling trajectory – the primary optimization target for real-world applications. \n\nWe hope that our response above has addressed your concerns, and we welcome further discussion if needed."}]

Author rebuttal:

Meta-review: This paper presents a theoretical framework that interprets Masked Diffusion Models (MDMs) through the lens of energy optimization in optimal transport. The authors demonstrate the mathematical equivalence of three distinct energy formulations - kinetic, conditional kinetic, and geodesic energy - within the MDM structure. A key theoretical contribution is the derivation of a closed-form optimality condition for the mask schedule, showing that MDMs minimize all three energies when this condition is satisfied. Building on this theory, the authors propose a practical method for schedule optimization by parameterizing the interpolation schedule via the CDF of a Beta distribution. Empirical results on synthetic and real-world benchmarks (e.g., code generation, math reasoning) demonstrate that their energy-inspired schedules can outperform standard hand-crafted baselines.

The topic of MDM models is currently of very high interest to the community. The paper received five reviews which, after discussion, all resulted in positive scores highlighting many aspects of the paper (novelty of the theory and interpretation, practical insights about schedule optimization, etc.). At the same time, one of the reviewers (gvvr) raised a valid and crucial point, questioning the profound theoretical significance of the work and suggesting that the theory might serve a primarily "decorative" role for what is essentially a practical contribution of a better schedule search method. After reading the paper, the meta-reviewer also agrees with this concern. Nevertheless, the meta-reviewer believes that the paper provides a valuable contribution to the community, with the merits (primarily the practical part) outweighing the flaws. The paper will be of interest to the community, and its acceptance is recommended.
Decision: Accept (poster)