Any-Order Flexible Length Masked Diffusion
Authors: Jaeyeon Kim
Lee Cheuk Kit
Carles Domingo-Enrich
Yilun Du
Sham M. Kakade
Timothy Ngotiaoco
Sitan Chen
Michael Samuel Albergo
Venue: ICLR 2026 Poster (2026)
Keywords: Diffusion Model
Generative Model
Discrete Diffusion
Stochastic Interpolant
Primary area: generative models

Abstract: Masked diffusion models (MDMs) have recently emerged as a promising alternative to autoregressive models over discrete domains. MDMs generate sequences in an any-order, parallel fashion, enabling fast inference and strong performance on non-causal tasks. However, a crucial limitation is that they do not support token insertions and are thus limited to *fixed-length* generations. To this end, we introduce **Flex**ible **M**asked **D**iffusion **M**odels (FlexMDMs), a discrete diffusion paradigm that simultaneously can model sequences of flexible length while provably retaining MDMs' flexibility of any-order inference. Grounded in an extension of the stochastic interpolant framework, FlexMDMs generate sequences by inserting mask tokens and unmasking them. Empirically, we show that FlexMDMs match MDMs in perplexity while modeling length statistics with much higher fidelity. On a synthetic maze planning task, they achieve $\approx$ 60\% higher success rate than MDM baselines. Finally, we show pretrained MDMs can easily be *retrofitted* into FlexMDMs: on 16 H100s, it takes only three days to fine-tune LLaDA-8B into a FlexMDM, achieving superior performance on math (GSM8K, 58\%$\to$67\%) and code infilling performance (52\%$\to$65\%).

Published as a conference paper at ICLR 2026
ANY-ORDER FLEXIBLE LENGTH MASKED DIFFUSION
Jaeyeon Kim
Lee Cheuk-Kitl.
Carles Domingo-Enrich
Yilun Du.12
Sham Kakade
Timothy Ngotiaoco 1.2
Sitan Chen 1.V
Michael S. Albergo' 2,
'Harvard University
*Kempner Institute Miicosoott Research New England
"Institute for Artificial Intelligence and Fundamental Interactions, MIT
ABSTRACT
Masked diffusion models (MDMs) have recently emerged as a promising alterna- tive to autoregressive models over discrete domains. MDMs generate sequences in an any-order, parallel fashion, enabling fast inference and strong performance on non-causal tasks. However, crucial limitation is that they do not support token in- sertions and are thus limited to fixed-length generations. To this end, we introduce Flexible Masked Diffusion Models (FlexMDMs), discrete diffusion paradigm that simultaneously can model sequences of flexible length while provably re- taining MDMs' flexibility of any-order inference. Grounded in an extension of the stochastic interpolant framework, FlexMDMs generate sequences by insert- ing mask tokens and unmasking them. Empirically, we show that FlexMDMs match MDMs in perplexity while modeling length statisties with much higher fidelity. On a synthetic maze planning task, they achieve 60% higher suc- cess rate than MDM baselines. Finally, we show pretrained MDMs can easily be retrofitted into FlexMDMs: on 16 H100s, it takes only three days to fine-tune LLaDA-8B into a FlexMDM, achieving superior performance on math (GSM8K, 58% 58%-67%) and code infilling performance (52%~>55%) We open-source our code at httpss//gttabbcomm/man//////////
INTRODUCTION
While diffusion models (Ho et al. 2020 Song et al. 2020 Sohl-Dickstein et al. 2015) are now the
leading paradigm for generative modeling continuous domains, recent work has begun to expand their scope to discrete spaces, The prevailing approach, Masked Diffusion Models (MDMs) (Shi ctt al. 2024 Sahoo et al. 2024 Gat et al., 2024), generates sentences a non-left-to-right, any-order Tashion. Compared to autoregressive models, this any-order generation ability yields substantially faster inference and strong downstream performance on non-casual tasks such as planning Ye et al.
2024 code Nie et al. 2025 Ye et al. (2025). and reasoning (Nie et al.) 2024).
Despite these successes, current MDMs cannot (I) model distributions supported on sequences of variable length and (2) insert tokens during generation (Figure I left). Both capabilities natural desiderata for generative models over discrete domains. We therefore ask: Can we model variable-length data while preserving MDMs' any-order generation power?
We answer in the affirmative by proposing the Flexible Masked Diffusion Model (FlexMDM). FlexMDMs start from empty string and gradually insert mask tokens and then unmask them (Figure [ right). Beyond learning the usual unmasking posterior-the the distribution of a clean to- ken masked positions-we introduce an insertion expectation: the expected number of tokens to insert conditioned on the current sequence. Crucially, we show that FlexMDM is theoretically grounded (i.e., under perfect training, samples from the true data distribution) and retains the any- order sampling property of MDMs, thereby directly addressing the question above. Empirically, we demonstrate that MM offers significant new upgrades to the MDM paradigm,
"equal contribution, randomized ordering; Lee Cheuk-Kit led the theoretical development of the method and engincering effort; Jacyeon Kim led experiment development and paper presentation. lead senior authors.

---

Published as a conference paper at ICLR 2026
MDM
FlexMDM
Inserted
Unmasked
mask
mask
mask
mask
mask
mask
Xo
Empty
mask
blue
mask
mask
mask
X
mask
love
skies
fx
[)
fo [
9 8)
go
I
mask
bluc
skies
green
mask
X
love
mask
mask
skies
love
blue
skies
green
cows
X1
love
Rio's
blue
skies
Variable Length
X
Token Insertion
Variable Length
Token Insertion
Figure 1: Flexible Masked Diffusion Model (FlexMDM) addresses MDMs' inability to handle variable-length sequences and token insertion while preserving any-order generation power. At each step, FlexMDM performs insertion and unmasking by predicting the expected number of mask tokens to insert (ge) and the posterior over clean tokens (fe), respectively.
A FlexMDM pretrained on OpenWebText is able to model the length distribution with substan- tially higher fidelity while matching the perplexity of an MDM counterpart.
On planning tasks, FlexMDM achieves markedly better results, beating the success rate of MDMs by nearly 60% natural synthetic baseline.
MDMs can be retrofitted into FlexMDMs at 8B+ scale: We fine-tune LLaDA-8B (Nie et al. 2025), open-source MDM, into FlexMDM using only 16 H100s for three days. The model transfers cleanly from its MDM initialization and, with its newly acquired variable- length capability, attains notably better performance GSM8K (58%~67%% and Code infilling (52%~>65-%)
Theoretically, our construction relies on the machinery of continuous-time Markov chains (CTMCs) and in particular introduces the new notion of a joint interpolant, novel extension of stochastic interpolants (Albergo & Vanden-Eijndee
2022 Albergo et al.) 2023b Lipman et al. 2022). Recent work (Zheng et al. |2224 Ou al., (024) established an equivalence between MDMs and any-order the need for CTMCs. In contrast, we prove that FlexMDMs also possess the flexibility of any any-order generation, yet the continuous time perspective is absolutely essential for them to accurately model the length distribution. Accordingly, we re-derive the connections between MDMs and stochastic interpolants and use them to ground the design of FlexMDMs.
language prodels-obviatte
Roadmap. We begin in Section 2 with a broadly accessible review of CTMCs and the connection between MDMs and discrete flow matching. Building on this, Section 3 derives the FlcxMDM training objective, and Section introduces our inference procedures. Section presents exper- imental results.
Concurrent work. Concurrent works (Wu et al. 2025b; Havasi al. 2025 attempt to tackle the same problem. Ww et al.] (2025b) introduces an auxiliary expand token in traiing and heuristically replaces each expand token with two mask tokens inference. Havasi al. (2025). also based the discrete flow matching framework, shares a similar theoretical grounding as our result. The main differences lie in our particular choice of interpolant that leads to the development of an any-order sampling algorithm. For clarity, we provide a detailed comparison in Appendix A
2 PRELIMINARIES: CONTINUOUS- -IME MARKOV CHAINS AND MASKED DIFFUSIONS
In what follows, we provide an overview of continuous-tin Markov chains (CTMCs), their role in defining discrete diffusion models, and link them to the MDM framework. As mentioned in the introduction, this theme is essential to defining FlexMDMs Section
Transport with continuous-time Markov Chains. Given a target distribution P1 sequences with a finite vocabulary set (c.g., text), our aim to learn to transport samples from a reference dis- tribution Po through continuum of distributions (pt)ue(011 such that pt= = Pr. This type of trans-
2

---

Published as a conference paper at ICLR 2026
port can be realized by continuous-time Markov chain, which is stochastic process (Xi) with Xg Po governed by time-dependent transition rate matrix (Re( bejo.11 satisfying
(x(
Σ Re(x,y), Re( 2 0, + yfz
Intuitively, the rate matrix determines the infinitesimal likelihood that X, transitions any other state y via
P P - =) [m=w) .R (2,, v) + o(h),
where we denote the conditional probability measure P(-|X, = x) of a new state given the present one. Here o(h) remainder that vanishes faster than h h 0. generative modeling for these discrete distributions, our aims are to (a) specify path of marginal distributions {pe (eeouu connecting Po P1 and (b) learn the associated Rg such that these marginals collectively satisfy the Kolmogorov forward equation:
O.p,(x) = Σ (v, Pr= Po-
(3)
This ensures that at time 1, the evolution specified by 2 results in a sample from the target distribution P1. The rate matrices defined in this paper sparse: therefore, we assume that the unspecified entries are 0 and the diagonal entries are defined through Equation (T).
2.1 MASKED DIFFUSION MODELS
We briefly review MDMs (Sahoo et al. 2024 Shi et al. 2024 and discrete flow matching with the masked construction (Gat et al.. 2024), through the lens of stochastic interpolants. The target distribution P1 assigns probability to length L sequences. The base distribution pu employed by these models is the point mass distribution at the fully masked length- sequence (m..... m), where m an auxiliary mask token.
To define the intermediate (Pe )e(11) that bridges the base and the target, make of stochas- tic interpolant (ar))=(o, collection of random variables whose marginal distribution defines the continuum { )ee(,, 1]. i.e., 2t, Pt. Although the previous notion of stochastic interpolant (Albergo (ct al. 2023b) is defined in a continuous space, it naturally extends to a discrete space, and we deier formal exposition Appendix C
MDM Interpolant
Design of distribution path. The stochastic inter- polant relies on a smooth and monotone unmasking schedule 0₁ [0, 1] + [0, 1] with boundary condi- tion (ag, a1) = (0, 1) and time derivative denoted by dr. To draw Me, first sample clean sequence Ii ~ pr: then, independently for every coordinate we draw an unmasking time 7" from density årdt
Frr
mask
skies
love
skies
Unmasked
Mask
mask
skies
T
umasking time
and set
m t T
mask
mask
mask
T
Figure 2: To draw sample 3t, can equivalently sample the clean sequence Iı ~ Pt. draw unmasking times, and then accord- ingly unmask or mask each coordinate's 's to- ken.
This process is illustrated in Figure clean token stays masked with probability 1 1-a, in coordinate- coordinate-indeeendenn fashion, defining Pr(- 21).
2 Hence, each
We then write by marginalizing over #1 P1.
MDM training. We now derive the MDM rate matrix that induces a CTMC whose marginals coincide with {Pr (peoe) and how it learned in practice. The central object is the unmasking pos- terior: the posterior on the clean token z for masked index given 24 = x and time step t, i..., P( = v|æ, = IT). We model this posterior with a neural network fo t) E (A(E))" where A Σ) denotes simplex of probability distributions over the vocabulary E.

---

Published as a conference paper at ICLR 2026
FlexMDM Interpolant
FlexMDM Inference
90 (Try T)
Tr
mask
skies
S,
love
Rio's 's
blue
skies
No. insert at Posesonn 1. (o.-00(--7)))00
Unmasked
Mask
mask
skies
mask
skies
(empty)
T umasking time
Unmask at [od
love
was
71
inscrtion time
6-0
fo(Er, [o]
Figure 3: Left (FlexMDM interpolant). To draw sample Zt, one can equivalently draw sample 3i ~ Pr, and for each token unmask, mask, or remove it according to the unmasking and insertion times T2). An auxiliary interpolant 8, gives closed-form expressions for the FlexMDM rate matrices. Right (FlexMDM Inference). Learned unmasking posterior and insertion expectation later used inference.
For every position where x' = m. the network aims to predict fo(æ, t) 2) se P(af = u/0, = x), and is trained by minimizing the following variational loss:
Le
/ 0
Σ
log folor 11, dt.
1
Olt
Here, E denotes the expectation 21 P1 and xa pr(-121). The minimizer of this loss is the ground-truth unmasking posterior, which fully determines the MDM's rate matrix below. Precisely, fort E [0, 1], the rate matrix time is given by: for partially masked sequence (EU (EU(m))
de
R1(2, xlx < v|)
P( v/u, = x),
V VE VE mm,
1 Oie
nomasking posterior
where xra=² +- u] denotes the sequence obtained from z by replacing its i-th token with Therefore, once So has learned the unmasking posterior, one can simulate the CTMC using the rate matrix in (5). The variational loss (4) quantifies the sampling guarantee of this estimated CTMC. Let p be the terminal distribution of the estimated CTMC. CTMC.Then, the loss function bounds the KL-divergence:
DKL (pillp Les
where L. is the global minimum of C. When the loss is in its infimum, the KL divergence vanishes, resulting in the ground truth distribution.
Connections to other MDM frameworks. Connecting to prior works MDM (Si et al. 2024 Sahoo et al. 2024 Campbell et al., (022), defining an interpolant is similar to defining forward process for the case of diffusion models or probability path in the case of flow matching. The modeled quantity is identical to the unmasking posterior across all frameworks. For inference, a common scheme is proceed by: each intermediate time step, (a) selecting a subset of indices to unmask and (b) sampling clean tokens from the learned posterior. In the infinitesimal limit, this procedure equivalent to simulating the CTMC of (5). Mcanwhile, subsequent work Kim et al. (202)); Nie et al. (2025) shows that MDMs also allow theoretically grounded any-order inference: tokens can be unmasked in an arbitrary order without necessarily following the CTMC at (5). We will revisit this aspect in Section and show that our FlexMDM preserves this advantage.
3 VARIABLE LENGTH MASKED DIFFUSIONS: TRAINING
In this section, we introduce Flexible Length Masked Diffusion Model (FlexMDM): a discrete diffusion that models distribution P1 assigning probabilities to sequences of different lengths. Fol- lowing the MDM's recipe, we aim to introduce a stochastic interpolant z, whose marginal distribu- tion defines the path (P:)ee(o.1) and learn the corresponding CTMC. Everything hinges selecting

---

Published as a conference paper at ICLR 2026
an interpolant that is (a) easy to sample t and (b) equipped with a closed-form rate matrix amenable to neural network training.
Challenge. Reusing the MDM interpolant is inadequate: t O, the base distribution Po would consist of fully-maskes sentences of variable lengths, which is impossible to sample since we do not know the length statistics of P1 in advance. On the other hand, one consider an interpolant constructed by masking and removing tokens from a clean sequence. However, this complicates the rate matrix characterization-ookee indices shift as insertions occur. To bridge this gap, we introduce the joint interpolant, an extension of the stochastic interpolant that augments the process with an auxiliary variable explicitly tracking token positions. This enlarged state space allows us construct a broader class of rate matrices while preserving an easy-to- to--spple base distribution.
Design of distribution path. We now introduce our FlexMDM's joint interpolant that allows us to model the variable length p1. This construction relies on two smooth, monotone schedules an insertion schedule a: [0. 1] + [o, 1] and an unmasking schedule 3:: [0. 1] - [0, 1], with the boundary conditions (ag. 01) = (So. 31) = (0,1) 1) and time derivatives denoted by á..
To draw 24, we first sample clean sentence æ1 P1- Independently for each coordinate i, draw an insertion time Tí and an unmasking time T3 with 77 < Tī according to the density below. Accordingly, we either unmask, mask, or remove zi to obtain z/: (empty).
0<b<T st Tj TiS111
3, Tj (121)))) 1- Br,
Ti di dt,
dt.
(6)
Here, 1 denotes the indicator function. We obtain xa by concatenating the symbols zi. and dropping at (empty). Consequently, the length of x4 equal to or less than that of x (see Figure left). As we mentioned above, augment 2, with an index-tracking variable st, forming the joint interpolant 8t). Let len(x;) denote the length I;; then
8 = {i E {0, len(x)) - 1) Ti <..
i.e. the set of indices whose clean tokens have not been deleted. Equivalently, the positions in ø referenced by 2, 's each index. By regarding 8, as a list and ordering its elements in ascending order, we also have I2 = (z:"((0) ,...,.i (1(1en(s)--1),
We revisit 8; shortly show how it enables an explicit rate matrix. Since (zz,, 82) is governed by the sampled unmasking and insertion times, we write (Iz, $,) ~ P- x Marginalizing (- x1) over # ~ Pi yields pt- Since the boundary condition sets 0o 8o 0, all tokens are deleted at i O; po is the point mass the empty string.
FlexMDM training. We now explain how we train our FlexMDM to learn the desired rate matrix. We first discuss what the CTMC looks like at high level: recall from (6 that when l. increases, to- kens are progressively inserted and unmasked, Indeed, one can show that a CTMC that generates the interpolant can be characterized by two quantities that govern the rate of insertion and unmasking:
Unmasking posterior (modeled by fo(a, t)]i]] E (E)):: for each index i that x = m, the posterior distribution over the underlying clean token.
Insertion expectation (modeled by go (x, [)[ E Rgo): for all indices i in x; the expected number of tokens that remain be inserted in between ---1 and af.
fe resembles the familiar unmasking posterior from MDMs, whereas go is new: it predicts how many tokens need to be inserted. Intuitively, modeling a variable-lenge p1 is harder than the fixed-length insertion expectation allows us to parameterize more complicated
setup of CTMC for FlexMDM: its rate matrix will appear soon in Proposition
2, To define the training loss, we set the boundary values of 8, as se[ se[-1] := -1 and 8, [len(s:)) len(x1), and let o(x; v) y x x log & denote a scalar Bregman divergence.
y
de len(xe)
Le
E jo
log fo( af
(ss - s - 1] - 1, go (w,, ()()
di
3
=
somunsking dove
insertion lous
(7)
'Writing xς mean the symbol derived from source position i is this a mild abuse of notation since the superscript refers to position in 21 rather than valid index of the (shorter) sequence Ft-
5

---

Published as a conference paper at ICLR 2026
Here, the expectation is taken over 21 P1, (zt, pe(((x)). Proposition Texactly the unmasking posterior and insertion expectation and shows they uniquely minimize 7
exxcxst characterizes
Proposition 1 (FlexMDM training loss). The loss La in is uniquely minimized
fa () u] = Pa = u|zz = x).
90 (z, [
[[s - se - 1] 2
wimazking posterior
insertion expectation
These quantities match the explanation above: the posterior over the clean token together with the expected number of insertions. They precisely determine the FlexMDM rate matrix stated next. Proposition 2 (FlexMDM Rate Matrix), Let the rate matrix R, be defined as:
Unmask: Rg (I,, x//² < u))
1 1-8. (xx w/ar x),
9,
2 VE VE
(8)
Insert: Rg (z, 4 m)
4, I-a, E [s s, - 11 =
where x 4 m the sequence obtained from x by inserting mask token between (x'-1_ x²). Then R₁ solves the KFE (equation (3) with Pt as the probability mass function of the FlexMDM interpolant Tt.
Proposition I thus implies that minimizing the loss yields exact recovery of the rate matrix. In practice, we could simulate the CTMC using the learned networks (fe. ge) in place of the ground- truth quantities in (8) By denoting the resulting terminal distribution as pf. the variational loss quantifies the terminal-time KI. divergence:
(PI 11p)) S Ce
We defer formal demonstration of propositions and the KL divergence guarantee to Appendix D Definition of the joint interpolant reinstated in definition D.2 the rate matrix in proposition D.3
the loss and variational bound in proposition D.4
Remark. Our FlexMDM interpolant introduces only one extra quantity beyond MDM's unmasking posterior: the insertion expectation, simple scalar per position. This stems from our design choice to gradually insert and then unmask a token. As shown in Section 5.2 this enables efficient task transfer of pretrained MDM weights. In contrast, alternative interpolants would require modeling more complex objects, such as a full token distribution, adding unnecessary training burden.
4 VARIABLE LENGTH MASKED DIFFUSIONS: INFERENCE
In this section, we outline inference algorithms for FlexMDM, focusing on two variants: vanilla inference and adaptive inference. We begin with brief overview of inference in MDMs.
Adaptive inference in MDM. For the case of MDM, MDM inference proceeds by simulating the rate matrix entries in From high-level one way this can be done is by (a) independently sampling a subset of masked tokens to unmask and (b) sampling clean tokens from the unmasking posterior. Crucially for what follows, the same guarantee holds for -on-independent adaptive choices of un- masking indices, e.g., confidence-based posterior, not following the matrix's unmasking entries. This adaptive inference strategy is widely used due to its empirical performance. We adopt this template and show that FlexMDM inherits the same any-order property.
correctness hinges on using the ground-truth unmasking
Vanilla inference. We begin with the vanilla inference of FlexMDM, which is obtained by discretiz- ing the CTMC in (8) using trained neural networks (fo. Choosing an appropriate discretization scheme is crucial, as different schemes can lead to markedly different empirical behavior. We adopt r-leaping- r-leaping-oriiinaeee originating in chemical physics and shown outperform naive Euler discretization for MDMs CCmpbell et al. [2022) whhch batches all events occurring within a fixed interval [1. ( r). At a high level, for each discretized step, we simultaneously (Figure 3 right):
Unmasking: For each mask token, sample for every unmasking a number according to the un- masking intensities in the rate matrix. Unmask only if a non-zero entry returned.
Insertion: Sample the number of mask-token insertions from a Poisson distribution parameterized by the insertion rate, then apply those insertions.

---

Published as a conference paper at ICLR 2026
Subroutine 1: VLMDM inference
Subroutine 2: Unmasking Step
Require: Learned functions (fo.go)
1: if vanilla inference:
Require: Discretization 0 = < iN
2: for E {ilX7, m) and E E
Require:
Insertion, Unmasking schedule ot, 3r
3: Set rate t Aie T
1: Initialize X 8-
1-80, k, ~ Poi(r fo(X tw ul) il such that ki = Set X K-
2: for k = to -1
4:
3: T tx++1 t
5:
4: Invoke Subroutine 2 for unmasking
6:
5: for i in [l.)(
7: if adaptive inference:
6: Set rate I-act
deg
8: Select K (the size of |(S)
9: for i E {il m)
7:
Sample é ~ Poi (r- go eo(Xx, tx [x)
10: Compute confidence Ci
8: Insert masks between X and Xi
11: for i in argmaxk
9: return XIN
X Cat(fo(XL, tt)
Algorithm 1: VLMDM inference. At each step we perform unmasking and insertion. For un- masking, unmask by r-leaping (vanilla) by confidence-basee selection (adaptive). The number of mask tokens to insert is drawn from a Poisson distribution. Notation: Cat, Poi imply the cat- egorical and Poisson distribution, respectively. argmaxk(C) is the indices set of the K largest components of C. We provide more details in Appendix E
As the number of steps - 0C, this inference algorithm recovers the CTMC and the discretization error vanishes. Algorithm [ details the full sampler.
Adaptive inference. Notably, one can choose the positions unmask adaptively. Precisely, each inference step select the unmasking positions according heuristic rule that prioritizes the most confident indices, where confidence is computed either from the model's unmasking posterior via a semi-autoregressive rule (prioritizing leftmost masks). We find such adaptive choice substantially boosts performance; Section Since unmasking indices in adaptive no longer trace the
transitions described by the rate matrix entries defined in (8), one might ask whether sampling still guarantees sample from the target distribution P1 in the infinitesimal limit. The following proposition answers in the affirmative.
Proposition 3 (Any-order inference, informal). Consider any sampling scheme that, at each step: (i) unmasks arbitrary subset of masked positions but draws revealed tokens from the ground-truth unmasking posterior; and (ii) applies insertion CTMC governed by the ground-trute -ruth rate matrix. Then the resulting process samples from the target distribution p1.
The formal statement and the proof of Proposition 3 are given in Appendix F In words, following the unmasking entries of the rate matrix corresponding the schedule used in training is not necessary to preserve the sampling guarantee. Moreover, the samplers as + oo in Algorithm i subsumed by the class in Proposition 3 therefore, assuming access to the ground-truth unmasking posterior and insertion expectation, the corresponding class of algorithms Algorithm samples from p1 up discretization error.
Remark. A key technical ingredient underlying the rigor of our adaptive inference is that the re- spective entries of the unmasking posterior of the ground truth matrix in Proposition 3 do not depend on the choice of unmasking schedule 3, (the proof is given in Appendix [..2.T]. This inde- pendence allows a single model fe learn all possible unmasking transitions arising along different paths that ultimately connect Po to thereby enabling adaptive unmasking at inference time. This feature is the same mechanism enabling adaptive inference for MDMs, but for FlexMDMs, proving that it interfaces correctly with insertions is quite subtle. Note that a similar notion independence of the choice of path-hac been introduced in continuous spaces Albergo al. (202a)) Negrel et al. 2025). We defer further discussion to Appendix B
5 EXPERIMENT
In this section, we present experimental results for FlexMDM, demonstrating the following:
MDx is effective veriahle-length learner: length modeling, planning, local edits.

---

Published as a conference paper at ICLR 2026
100
70
OpanWel/Toct
FlexMDM
80
MDM
MDM (266)
Generative Perplesity
60
MDM
FiecM(MM (254)
Quantile
FlodEDM (1024)
50-
40
40
20
30-
0
256
512
708
102
266
512
1024
2048
4096
Path
Subgozis
Invalid cells
Leugth
Sampling steps
(a) Maze task illustration. The model is given subgoals and is re- quired to connect them.
(b) Length modeling. FlexMDM recovers the true length distribution Open WebText training data.
(c) Perplexity. FlexMDM achieves generative perplexity par with MDM.
FlexMDM is scalable: 8B FlexMDM is obtainable by initializing from a pretrained MDM.
Section 5.1 presents from-scratch results for FlexMDM on text and planning tasks distributions, confirming its practical efficiency. Next, Section 5.2 provides an 8B 8B-scale FlexMDM's training recipe, initialized from LLaDA-8B [Ne al. (2025), and evaluates in math and code infilling tasks. We begin with the architectural and scheduling choices used throughout.
Training design. Recall from Section 3 that FlexMDM models the unmasking posterior fo and insertion expectation go given state and time step f. We adopt DiT (Peebles & Xiel 2023). bidirectional transformer that enables additional embedding, as a backbone. To learn both quantities jointly, we attach two output heads: standard posterior head for fo and scalar softplus head for go. Moreover, we choose our unmasking and insertion schedule to be both linear, ar = Be =
5.1 PRETRAINING
In this section, we evaluate FlexMDM's ability to learn variable-length data from scratch. Our base- line is MDM, which is fixed-length but can handle variable-length sequences by padding fixed maximum length with an auxiliary pad token. This padding setup is widely used in instruction fine- tuning when variable-length answers are desired (Nie et al.) 20225] (2024] Ye et al., [2225] Gong al., 2024). For fair comparison, we use vanilla inference for both MDM and FlexMDM throughour. Further experimental details appear in Appendix
5.1.1 PRETRAINING ON TEXT DATA
We first construct a training dataset from the raw OpenWebTex, WebTeet corpus (Gokaslan et al. 2019), split- ting each article into paragraphs to preserve semantic coherence and yield variable-length sequences. Models pretrained on this data, therefore, generate variable-length text.
Results. We train 175M FlexMDM and MDM with a maximum sequence length 1024 for 500K iterations and batch size 1024. Using the pretrained models, we vary the number of sampling steps and measure (a) generative perplexity as a proxy for text fluency, and (b) the induced length distribu- tion. Figure 4c shows comparative generative perplexity for the two models, improving sampling steps increase, indicating no fluency degradation for FlexMDM despite its more involved loss ob- jective. Crucially, we observe that FlexMDM matches the true length distribution far more closely (Figure 4b): with only 256 steps it tracks the ground truth distribution (red line, whereas MDM remains miscalibrated even at 1024 steps (blue line). In Appendix F.I.II we provide additional ex- perimental results: the entropy of given sequences to ensure the generative perplexity is measured under similar conditions and the perplexities for larger sampling steps (8192, 16384, 32768).
Remark. We remark that our pretraining pipeline differs from prior MDM setups that truncate the corpus fixed maximum length. Also, one might ask why do provide additional metrics on text benchmarks, such as validation perplexity. This is because MDM and FlexMDM use differ- ent objectives (see equation 4 and equation
making likelihood comparisons hard interpret.
TThe ground-truth unmasking posterior is independent of 3r, we condition the network o only; under the linear choice o, t, this coincides with the usual time emhedding.
8

---

Published as a conference paper at ICLR 2026
Albeit, we experimenaally confirm that both loss functions lihood, although with different scales, by evaluating both pretrained models on a downstream text benchmark (Appendix F.I.I). We address the concern about the absence of the metric by evaluating scaled models on downstream benchmarks Section 5.2
reliable proxy of the like-
5.1.2 PLANNING TASK
We further evaluate FlexMDM's ability planning task in a discrete space. Motivated by an ear- lier study Janner et al. (2022) that investigated the ability of continuous diffusion in maze tasks, we design grid-maze benchmark: the maze is fixed but unknown to the model, with a subset of cells invalid. Given a sequence of subgoal grids (g)...., .K)). the model must connect this sequence with- out entering invalid cells (see Figure 4a This subgoal structure aligns naturally with FlexMDM: starting from (01 (01..... 9K ). inference inserts mask tokens between subgoals and then unmasks to generate a feasible path. Theoretically, this can be seen as augmenting the base distribution to con- tain data-dependent distribution Albergo et al. (2024). This is in stark contrast to MDM, where it must preassign each subgoal to a specific position, which is difficult to know a priori. We provide additional details Appendix F2
Results. We use 41 X maze and control the task's difficulty via varying the number of subgoals K e 12]. As increases, MDM performance degrades markedly, while FlexMDM maintains robust success rates, reaching a gap of up to 60% at K = 12. These results firmly support FlexMDM planning, where preallocating token positions is inherently challenging for fixed-length models.
Difficulty
MDM
FlexMDM
Easy
68.4%
92.3%
Medium
29.3%
90.4%
Hard
24.2%
90.0%
principled approach for subgoal-based based
Table 1: FlexMDM outperforms MDM the subgoal-style maze-
planning task.
5.2 SCALING FLEXMDM
In this section, we address FlexMDM's scalability by scaling 8B parameters and observing no- table improvements over an MDM baseline. We start from the observation that MDM and FlexMDM both share the unmasking posterior as core component, suggesting effective task transfer from a pretrained MDM might be possible. To demonstrate this, concretely, we initialize from LLaDA- Base (Nie et al. 2025> and make the following modifications: (a) add time-embedding layers and scalar head to model the insertion expectation; (b) attach LoRA adapters. Altogether, the resulting number of trainable parameters is 4000M. To cover both natural and mathematical language, we train on the 50:50 mixture of OpenWebText (Gokaslan et al. 2019) and Proof-Pill Pilef- (Azerbayev) et al. 2023). Surprisingly, we observe rapid transfer: within three days on 16 H100 GPUs, the model generates variable-length sentences. Notably, the number of tokens used is 1=3.1 B, and this is in contrast much smaller than the number of LLaDA-Bas pretraining tokens (s11.T) We then instruction-finettue (IFT) this base FlexMDM to evaluate on downstream tasks. See Appendix for more details.
Results. For comparison, we train FlexMDM and LLaDA-Baas from the same number of IFT pairs. For math and code, respectively, we IFT on the GSM8K train Split(Cobbe et al. (2021); 88000 pairs) and the educational split of ope-sft-stage-2 (Huang al., 2024) ( pairs), for which IFT-ed models are evaluated in the GSM8K test split and Eval-nnEvaliine (single line) (Bavarian et al.) 2022] in zero-shot. Sampling done by confidence-baat basedd sampling with a sliding window. Notably, as the number of sampling steps increases, FlexMDM continues to improve, highlighting its strength in reasoning tasks given sufficient compute-whereet the IFT-ed LLaDA's performance remains flat. Although in this experiment more diverse instruction intruction-answer answer pairs with sufficient compute will yield a more generalized model.
IFT task-specific pairs, expect that training much
For fair comparison, since FlexMDM is not IFT-ed, we IFT LLaDA-BBsse rather JLaDA-instruc this differs from Zhao al. (025). We employ zero-shot evaluation, which also differs from Nie al. (2024

---

Published as a conference paper at ICLR 2026
GSM8k Performance
HumanFval In-filling Performance
65
FlexMDM
FlexMDM
66-
LLaDA
LLaDA
Pass01
60
62-
55
60
58
128
256
512
1024
4096
128
256
512
1024
4096
Sampling step
Sampling step
Figure FlexMDM performance exhibits superior scaling when more sampling steps are allocated.
CONCLUSION
In this work we proposed Flexible Masked Diffusion Models (FlexMDM), discrete diffusion framework over variable-length sequences. Theoretically, via a joint interpolant viewpoint, we pro- vide rigorous guarantees for both training and inference of FleMMDM. Empirically, FlexMDM learns variable-length structure across diverse scenarios, scales to 8B parameters, trains in only a few GPU- hours, and yields substantial improvements math and coding infilling tasks. Further exploration of FlexMDM's 's capabilities is a promising direction for future work.
Beyond these results, goal is to align generative modeling with how humans and nature compose discrete sequences. Instead filling fixed positions; they insert, revise, and reorder tokens. We hope that work takes step in this direction.
REFERENCES
Michael S. Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic inter- polants, 2022.
Michael S. Albergo, Nicholas M. Boffi, Michael Lindsey, and Eric Vanden-Eijndenn Multimarginal generative modeling with stochastic interpolants, 2023a. URL http:::://exiv..rr//aaa.. 2.310.06693
Michael Albergo, Nicholas Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: unifying framework for flows and diffusions. arXiv preprint urXiv:2303.08797, 2023b.
Michael Samuel Albergo, Mark Goldstein, Nicholas Matthew Boffi, Rajesh Ranganath, and Eric Vanden-Fijnnen. Stochastic interpolants with data-dependent couplings. In Forty-firs/ Interna- tional Conference on Machine Learning. 2024. URLiittps://peeneviiee net /formm?
id FFILRGDO
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in neural information processing systems, 14:179811-799999999 11993,, 2021.
Zhangir Azerbayev, Hailey Schoelkopf, Keiran Paster, Marco Dos Santos, Stephen McAleer, Al- bert Q Jiang, Jia Deng, Stella Biderman, and Sean Welleck. Llemma: An open language model for mathematies, 2023.
Mohammad Bavarian, Heewoo Jun, Nikolas Tezak. John Schulman, Christine McLeavey, Jerry Tworek, and Mark Chen. Efficient training of language models to fill in the middle, arXiv preprint 22X.:22207..4255 2022.
Heli Ben- Ben-Hamu, Itai Gat, Daniel Severo, Niklas Nolte, and Brian Karrer. Accelerated sampling from masked diffusion models via entropy bounded unmasking. arXiv preprint urXiv:2505.24857 2025.
10

---

Published as a conference paper at ICLR 2026
Joe Benton, Yuyang Shi, Valentin De Bortoli, George Deligiannidis, and Arnaud Doucet. From denoising diffusions to denoising markov models, 2024. URL http://aaxivv.oooo abs/
22111 03595
Julius Bemnr, Lorenz Richter, and Karen Ullrich. An optimal control perspective on diffusion-based generative modeling. Transactions Machine Learning Research, 2024. ISSN 2835-8856. URL
net www.wforrrroorrooemmo fo
Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du, Zhe Fu, et al. Deepseek Ilm: Scaling open-source language models with longtermism. arXiv preprint arXiv:2401.02954 2024.
Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. A continuous time framework for discrete denoising models. Advances in Neural Information Processing Systems, 0112222662227777 2022.
Andrew Campbell, William Harvey, Christian Weilbach, Valentin De Bortoli, Tom Rainforth, and Arnau Doucet. Trans-dimenn Trans-dimensional generative modeling via jump diffusion models, 2023. URL
hittps://axxx http://awwwi.oooo.//////////////
www./oo.////.oooooo..........o.r 2305. 16261
Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design. arXiv preprini 240::::0000099977,, 2024.
Huiwen Chang, Han Zhang. Lu Jiang, Ce Liu, and William T Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 1 11315-1122,, 2022.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichir Nakano, et al. Training verifiers to solve math word problems. arXiv preprint 1Xiv:2110.14168. 2021.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances neural information processing systems, 35:16344-16359 2022.
Google DeepMind. Gemini diffusion,, 2025. URL httpss//bbbeg..ooole//eeooooooo
google deepmind//gemini www-ini--iffuu htw.://d//eemiii--.ffiiioioomm diffusion/
Juechu Dong, Boyuan Feng, Driss Guessous, Yanbo Liang, and Horace He. Flex attention: A programming model for generating optimized attention kernels, arXiv preprint arx0v:2412.054966 2024.
Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky TQ Chen, Gabriel Synnaeve, Yossi Adi, and Yaron Lipman. Discrete flow matching. Advances in Neural Information Processing Systems, 37: 2024.
Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex. Openwebtext corpus.
http:
//Sky]1on007.ghtt git.ub.ttuuuu.... wttp:hhtt......o.eeeeeeeeeeeeeee httppppeeebTextcorrreeoooooo 2019.
Shansan Gong, Shivam Agarwal, Yizhe Zhang, Jiacheng Ye, Lin Zheng, Mukai Li, Chenxin An, Peilin Zhao, Wei Bi, Jiawei Han, et al. Scaling diffusion language models via adaptation from autoregressive models. arXiv preprint arXiv:2410.1 17891, 2024.
Shansan Gong, Ruixiang Zhang, Huangjie Zheng, Jiatao Gu. Navdeep Jaitly, Lingpeng Kong, and Yizhe Zhang. Diffucoder: Understanding and improving masked diffusion models for code gen- eration. arXiv preprint arXiv:2506.2066 20639, 2025.
Marton Havasi, Brian Karrer, Itai Gat, and Ricky TQ Chen. Edit flows: Flow matching with edit operations. arXiv preprint (tt.:2506000//8.. 2025.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.

---

Published as a conference paper at ICLR 2026
Peter Holderrieth, Michael Samuel Albergo, and Tommi Jaakkola. LEAPS: A discrete neural sam- pler via locally equivariant networks. In forty-second International Conference on Machine Learning. 2025a. URL http:///eeaaaaeeeee.ooooo/oooood
net wetp//oormmidddlolaatt.....
Peter Holderrieth, Marton Havasi, Jason Yim, Neta Shaul, Itai Gat, Tommi Jaakkola, Brian Karrer, Ricky T. Q Chen, and Yaron Lipman. Generator matching: Generative modeling with arbitrary markov processes. In The Thirteenth International Conference on Learning Representations, 2025b. URL httpps://ppeneeealt/ net Lttee/oommmmiiidlPlaaalaa
Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling. Argmax flows and multinomial diffusion: Learning categorical distributions. Advances in neural information processing systems, 34: 12414-12466 12465, 2021.
Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Con- ference on Learning Representations, 2022. URL https: /openrevieww...eeooo net /
forum?
id=nzevkeeFYYE
Siming Huang, Tianhao Cheng, Jason Klein Liu, Jiaran Hao, Liuyihan Song, Yang Xu, J. Yang, J. H. Liu, Chenchen Zhang, Linzheng Chai, Ruifeng Yuan, Zhaoxiang Zhang, Jie Fu, Qian Liu, Ge Zhang. Zili Wang, Yuan Qi. Yinghui Xu, and Wei Chu. Opencoder: The open cookbook for top-tier code large language models. 2024. URL httpttttts//axi..... http://aww...o/pp////1.1..00. 04900
Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, 2022.
Jaeyeon Kim, Kulin Shah, Vasilis Kontonis, Sham Kakade, and Sitan Chen. Train for the worst, plan for the best: Understanding token ordering in masked diffusions. arXiv preprint arXiv:2502.06668, 2025.
Inception Labs Samar Khanna, Siddhant Kharbanda, Shufan Li, Harshit Varma, Eric Wang, Sawyer Birnbaum, Ziyang Luo, Yanis Miraoui, Akash Palrecha, et al. Mercury: Ultra-fast language models based on diffusion. arXiv preprint arXiv:2506. 17225006.72998 /7298, 2025.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben- Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. 2022.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning generate and transfer data with rectified flow, 2022.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 1rXiv:1711.05 2017.
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. arXiv preprint 1iXx:23300.18344 2023.
Long Ma, Fangwei Zhong, and Yizhou Wang. Reinforced context order recovery for adaptive rea- soning and planning, arXiv preprint arXiv:2508. 13070, 2025a
Xinyin Ma, Rumpeng Yu, Gongfan Fang, and Xinchao Wang. dkv-cache: The cache for diffusion language models. arXiv preprint arXiv:2505. /578/, 2025b.
Hugo Negrel, Florentin Coeurdoux. Michael Albergo, and Eric Vanden- Multitask learn- abs / 508 04605
ing with stochastic interpolants, 2025. URL
Shen Nie, Fengqi Zhu, Chao Du, Tianyu Pang, Qian Liu, Guangtao Zeng Min Lin, and Chongxuan Li. Scaling up masked diffusion models on text. arXiv preprint vrXiv:2410.185/1 2024.
Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, Jun Zhou, Yankai Lin, Ji-Rong Wen, and Chongxuan Li. Large language diffusion models. arXiv preprint arXiv:2502.0999 099922 2025.
Jingyang Ou Shen Nie. Kaiwen Xue, Fenggi Zhu, Jiacheng Sun, Zhenguo Li. and Chongxuan Li. Your absorbing discrete diffusion secretly models the conditional distributions of clean data. arXiv preprint arXiv:2406.037388 03733 2024.
12

---

Published as a conference paper at ICLR 2026
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the JEEE/CVF international conference on computer vision, pp. 2023.
Stefano Peluchetti. Non-denoising forward-time diffusions, 2022.
URL https://
openreview wwt.:/eeeeooommeaaaeeeeee.e
Fred Zhangzhi Peng. Zachary Bezemek, Sawan Patel, Jarrid Rector-Brooks, Sherwood Yao, Avishek Joey Bose, Alexander Tong, and Pranam Chatterjee. Path planning for masked diffu- sion model sampling. arXiv preprint arXivv2502.03540, 2025.
Qwen. Qwen2 technical report. arXiv preprint arXiv:2407.10677 2024.
Alec Radford and Jeffrey Wu. Rewon child, david luan, dario amodei, and ilya sutskever. 2019. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Litu Rout, Constantine Caramanis, and Sanjay Shakkottai. Anchored diffusion language model. arXiv preprini arXiv:2505./84 2025.
Subham Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. Advances in Neural Information Processing Systems, 37:1 1001661101844 2024.
Subham Sekhar Sahoo, Zhihan Yang, Yash Akhauri, Johnna Liu, Deepansha Singh, Zhoujun Cheng, Zhengzhong Liu, Bric Xing, John Thickstun, and Arash Vahdat. Esoteric language models. arXiv preprint arXiv:2506.0/99 2025.
Neta Shaul, Itai Gat, Marton Havasi, Daniel Severo, Anuroop Sriram, Peter Holderrieth, Brian Kar- rer, Yaron Lipman, and Ricky TQ Chen. Flow matching with general discrete paths: kinetic- optimal perspective. arXiv preprint arXiv:2412.034877 2024.
Neta Shaul, Itai Gat, Marton Havasi, Daniel Severo, Anuroop Sriram, Peter Holderrieth, Brian Kar- rer, Yaron Lipman, and Ricky T. Q. Chen. Flow matching with general discrete paths: A kinetic- optimal perspective. In The Thirteenth International Conference on Learning Representations, 2025. utLpttts:
Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, and Michalis Titsias. Simplified and general- ized masked diffusion for discrete data. Advances in neural information processing systems, 37: 103131-103167, 2024.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learn- ing, pp. 2256-2266 pmlr, 2015.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint 1rXiv:2011.13456 2020.
Yuxuan Song, Zheng Zhang, Cheng Luo, Pengyang Gao, Fan Xia, Hao Luo, Zheng Li, Yuehang Yang, Hongli Yu, Xingwei Qu, et Seed diffusion: large-scale diffusion language model with high-speed inference. arXiv preprint 20Xi::2508.02193. 2025.
Alexander Swerdlow, Mihir Prabhudesai, Siddharth Gandhi, Deepak Pathak, and Katerina Fragki- adaki. Unified multimodal discrete diffusion. arXiv preprint arXiv:2503.20885 2025.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao. Vedanuj Goswami. Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet. Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nic, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh
13

---

Published as a conference paper at ICLR 2026
Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, lliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv: 2307.09288 2023.
Francisco Vargas, Shreyas Padhy, Denis Blessing, and Nikolas Nüsken. Transport meets varia- tional inference: Controlled monte carlo diffusions, 2025. URL /tL_ittts///rrxv.rr////a.
2307.010501
Dimitri von Rütte, Janis Fluri, Yuhui Ding, Antonio Orvieto, Bernhard Schölkopf, and Thomas Hofmann. Generalized interpolating discrete diffusion. arXiv preprint arXiv:2503. S50.::5533044422, 2025.
Zhe Wang, Jiaxin Shi, Nicolas Heess, Arthur Gretton, and Michalis Titsias. Learning- order autoregressive models with application molecular graph generation. arXiv preprint arXiv:2503.05979, 2025.
Chengyue Wu, Hao Zhang, Shuchen Xue, Zhijian Liu, Shizhe Diao. Ligeng Zhu, Ping Luo, Song Han, and Enze Xie. Fast-dllm: Training- free acceleration of diffusion Ilm by enabling cache and parallel decoding. arXiv preprint arXiv:2505. 226188 2025a.
Zirui Wu, Lin Zheng, Zhihui Xie, Jiacheng Ye, Jiahui Gao, Yansong Feng, Zhenguo Li, Victoria W., Guorui Zhou, and Lingpeng Kong. Dreamon: Diffuusion language models for code infill- ing beyond fixed-size canvas, 2025b. Uttpss:::///kku hwtppp.itttbuu....... 1otpbo///..ooll//25/////// dreamon
Jiacheng Ye, Jiahui Gao, Shansan Gong, Lin Zheng, Xin Jiang, Zhenguo Li, and Lingpeng Kong. Beyond autoregression: Discrete diffusion for complex reasoning and planning. arXiv preprint arXiv:2410.14157, 2024.
Jiacheng Ye, Zhihui Xie, Lin Zheng, Jiahui Gao, Zirui Wu, Xin Jiang, Zhenguo Li, and Lingpeng Kong. Dream 7b, 2025. URL htthtttpp////iiiit
github http:/uuw....oolloo////////////e 20ee/
Siyan Zhao, Devaansh Gupta, Qinqing Zheng, and Aditya Grover. d1: Scaling reasoning in diffusion large language models via reinforcement learning. arXiv preprint arXiv:2504. /22 2025.
Kaiwen Zheng, Yongxin Chen, Hanzi Mao, Ming-Yu Liu, Jun Zhu, and Qinsheng Zhang. Masked diffusion models are secretly aine-sgn masked models and exploit inaccurate categorical sampling. arXiv preprint Att.:2409.02008, 2024.
14

---

Published as a conference paper at ICLR 2026
CONTENTS
Introduction
2
Preliminaries: Continuous- Time Time Markov Chains and Masked Diffusions
Masked Diffusion Models
3
Variable Length Masked Diffusions: Training
Variable Length Masked Diffusions: Inference
5
Experiment
5.1
Pretramiig
5.1.1
Pretraining text data
8
5.1.2
Planning task
5.2
Scaling up FlexMDMM
6
Conclusion
10
Related Works
17
Notation
17
Discrete Stochastic Interpolants: Definitions and Propositions for Section 2
18
Discrete Stochastic Interpolant
18
The Masked Diffusion Interpolant
19
Joint Discrete Stochastic Interpolants: Definitions and Propositions for Section
D.I Joint Interpolant
21
D.2 Flexible-Length Masked Diffusion
22
Details for Section
25
E.
Precise detail the inference algorithms
25
Proof of FlexMDM's any-order inference capability
26
E.2.1
Proof preliminaries)
26
Formal guarantee for adaptive inference
26
Proof of Theorem
27
Proof of Lemmal
28
E,6 Proof of Lemma
28
F
Experimental details
29
Pretraining on WebTee
29
E1.1 Additional experiments
30
15

---

Published as a conference paper at ICLR 2026
F2
Pretraining the Maze planning task
30
Weight Initialization training from LLaDA
32
16

---

Published as a conference paper at ICLR 2026
A RELATED WORKS
Discrete diffuuson and flows. Early diffusion models were formulated as continuous-time Markov chains over continuous spaces with Gaussian transition kernels (Sohl-Dickstein et al., 2015: Ho
et al. 2020), and were later connected continuous- time formulations
stochastic differential equations, offering a unifying perspective on score-based generative modeling (Song et al. 2020). In parallel, discrete diffusion has been developed from the viewpoint of Markov chains over discrete space (Hoogeboom et al.) 202I). Notably, Austin et al. (202I) introduced D3PM with several fami- lies of discrete transition kernels, and Lou et al. (2023) proposed SEDD, which adopts score-based training objectives. A complementary line of work studies discrete flows (Campbell et al. [2024] Gat al. 2024p, aiming to understand continuous time Markov chains (CTMCs) that interpolate be- tween data and base distributions; this perspective aligns with ours. Subsequent extensions consider token-wise paths and path-wise structure within such flows (Shaul al. 2024).
Masked Diffusion Models. Among discrete-transition designs, absorbing-state (a.k.a. masking) kernels have become popular and strong-performing choice. Recent work shows that this yields simple and principled training recipe, referred to as Masked Diffusion Models (MDMs) (Sahoo 2024). A growing body of results demonstrates the scalability of this approach across problem settings and modalities, including large-scale natural language modeling (Nie et al.
ct al. 2024 Shi et al.
2024 [025 Ye et al. 2025 Song et al., 2025| DeepMind] (205), (025)) code generation (Labs al. 2025
Gong al. 2025). and multimodal learning (Swerdlow et al. 2025).
Any-order inference in MDMs. With the advent of MDMs, subsequent work has established that they admit theoretically grounded any-order inference, wherein tokens can be unmasked in arbitrary orders rather than following a fixed CTMC schedule (Kim al. 2025 Peng et al. 2025). Practical loken-ordering rules span spectrum of heuristics based on model confidenee and uncertainty e.g., maximum- maximum-probability logits (Chang et al.) [022 Zheng al., 2024), probability margin (Kim al.)
2025),
semi-autoregressive schedules (Nie et al. 2024), and entropy- -barop criteria (Ben-Hami et al., as well as strategies that leverage reference models to guide the unmasking trajectory (Peng
2025
al.
2025). Beyond heuristics, another thread trains auxiliary modules to anchor or adapt the
generation order (Rout et al. 2025), while recent work directly learns token orders end-to-end (Mal
al.
2025a Wang al. 2025
Stochastic interpolant. Stochastic interpolant (Albergo & Vanden-Hijnden 2022 Albergo al, 2023b) is general framework for building measure- transsortt based generative models on contin- uous state space. While building off different philosophical grounds, it can be seen as equivalent to flow matching (Lipman et al. 2022: Liu et al.) 2022). Extensions of the interpolant have been proposed for conditional generation through data-dependent coupling (Albergo et al. 2024), which adopt for infilling task design in Section 5
Descriptive overview on concurrent work. The most notable concurrent work is Edit- Flow (Havasi et al. 2025), where the primary mathematical machinery that enabled their construc- tion referred to as "Flow Matching with Auxiliary Process". We note that this can be interpreted as mathematically equivalent to the notion of joint interpolant in this work, e.g., our Proposition D.1 is equivalent to Theorem 3.1 in Havasi et al. (2025).
The main differences are (1) the choice of interpolant and (2) the guarantee of any-order inference. Whereas EditFlow built around an explicit probability path, we instead define a pair of coupled random variables that implicitly induce this path, leading to a different choice of intermediate. As discussed in Section 3 our choice of interpolant yields a distinct training objective for an unmask- ing posterior and insertion-expeccatio term. Consequently, it enables the any-order inference guarantee established in Section!
B NOTATION
In this section, we reiterate the notations used in the main body and introduce auxiliary notations that are used in the proofs of the appendix.
17

---

Published as a conference paper at ICLR 2026
Strings. Let denote the empty string, E vocabulary of words, special mask token. We write the i-th element z with O-baesd indexing, ø|s the string indexed by an index set, e.g. abc[[0,2]] 2|| ac To insert a token before position i in string æ, write < v, e.g., to prepend a token abc 4 = dabe and to append a token abc 4g d = abed. To replace the i-th token in x with v, write x: xxx² v/. As much of the work involves masking, write C y if x can be constructed by partially masking We write mask(æ), unmask(r) S [Ien(x)] for the set of indices corresponding to mask and clean tokens.
C DISCRETE STOCHASTIC INTERPOLANTS: DEFINITIONS AND PROPOSITIONS FOR SECTION
In continuous spaces, common approach to define generative transport is the stochastic interpolant framework, which implicitly defines the interpolation distribution pt by specifying interpolant (=(.))))1) and regressing the required quantities to realize the transport.
In the section C.I we introduce a discrete analogue of the stochastic interpolant. To illustrate this framework, we reformulate the widely used masked diffusion model for sequences of length within our setup. Briefly, masked diffusion defines the interpolation pt by progressively unmasking tokens in sentences drawn from the data distribution. At time t 0, all tokens are masked, so Po a point mass the fully masked sequence, the token repeated n times. The transition rates driving the generative transport can be characterized as functions of per-token posterior probabilities conditioned on time I. These are typically learned by minimizing a variational objective in the form of weighted cross-entropy loss.
C.1 DISCRETE STOCHASTIC INTERPOLANT
To obtain the target rate matrix, the discrete stochastic interpolant relies on an interpolating rate matrix that drives a sample from sample drawn from a sample from Po a sample from pl. defined as follows:
Definition C.1 (Discrete Stochastic Interpolant and Interpolating Rate). Let Fo ~ Po and Tı ~ Pr. A discrete stochastic interpolant is family of random variables (1=1) defined
common probability space and satisfying the boundary conditions It=0 = Io and It=1 = I1, for which there exists continuous-tin
Markov chain with bounded, time-dependent transition rate matrix K2 such that, for each t E [0, 1), Law( Vs To, æ1) coincides with the marginal distribution at time of that Markov chain started at Xu- We refer K2o. interpolating rate matrix.
With an interpolating rate of an interpolant, the target rate matrix can then be obtained through Proposition C.I
Proposition C.1 (Target Rate), Given a discrete stochastic interpolant xs and an interpolating rate matrix K2o the continuous-timg Markov chain with initial distribution po and target transition rate matrix R₁ defined as,
R2(x, y) = Exoo, [K (o,, (r.y))a;
.r]
has marginals equal to Law(x:).
Proof. Writing Pe probability mass function (pmf) of 2, and pe(:/20, 21) pmf of æ; conditioned on Zo. 21. We further write g( 21) the joint pmf of in and xj and qi (Zo, 21 [x;) to be the joint pmf conditioned on It-
18

---

Published as a conference paper at ICLR 2026
It suffices to show R2 satisfies the Kolmogorov Forward Equation with pmf pt follows:
Σ
Emo.a mi [K (2,2))2 ylpe(y)
K7o (r0, 21 (y)p((y)
20,21
Kt
Ktomn (y, x1 21)9 21)
Z0
Σ
Kr
K (v. 21)
y (0)r
Osp((x)
This concludes the proof.
Remarks. While written considerably differently, the framework is mathematically equivalent to discrete flow matching Gat et al. (2024). The difference is only philosophical: discrete flow match- ing relies on the notion of conditional probability path that the interpolating rate should induce, whereas define such probability path only implicitly through the definition of the interpolant.
C.2 THE MASKED DIFFUSION INTERPOLANT
As concrete example of the discrete stochastic interpolant, we reformulate the masked diffusion model and its learning in the framework. As masked diffusion starts from point mass, we drop the dependence of in writing.
Definition C.2 (The Masked Diffusion Interpolant). Let Fn ~ p1 be a sentence of length n drawn from the data and ag smooth unmasking schedule that interpolates from az 0 to 0z=1 1.
Define the unmasking times {T JiEjo...... as:
Vi iO...., 1): T ae dt,
Then, the masked diffusion interpolant is defined as:
m ift < Ti,
o ift 2 T.
In other words, at each time f, x, reveals a subset of the tokens of 21, with each token of indepen- dently unmasked at its associated time T
Proposition C.2 (The Masked Diffusion Interpolating Rate). One interpolating rate K of the masked diffuusio interpolant Fe is given by:
ái
Ya C 1I,U E. m K2(2, x|o" + +v) 1 - -a
1 =1)-
Proof. Let pe(-(21) as the pmf of x, conditioned on æ1. From the definition of the interpolant, notice that:
len(21)-1
(x(x(1))=
II
(1
1z) = m) + tallfo =
1=0
19

---

Published as a conference paper at ICLR 2026
We verify that K, satisfies the Kolmogorov Forward Equation 33 under the conditioned pmf as follows,
L.H.S
Oepe(al (1:(21)
len(21)-1
= 0
II
(1 (1-a;)1(s's = m) fa,lfx'- =
1=0
len(z))-1
Σ
+ = II o.)1 m) + = =x/=
1=0
17i
R.H.S
Σ 1fx² = 2]}
pe(a|a' + m| m||x1] 1f2² = m) at
pr (1(2)11)
1 aq
1 1
lon(z))-1
Σ 1f 21 aa
(1 at)
II (1--a))12a m) +a =z}
jfi
len(m.)-1
1xx = m)
at
(1 at
I a. 162 = m) + =
Ot
L.H.S.
This concludes the proof.
Note that since each T" is sampled independently from a continuous distribution, the probability that two unmasking times coincide is zero. Thus, only single token is unmasked in any infinitesimal transition almost surely.
Proposition (The Masked Diffusion Target Rate), By proposition target rate Rg that induces Law(x;) is:
de az
Vz C E m
R<(2,0/2' t u)) P( w/2re- = x.).
Proof. Following proposition C.2 the result follows from invoking proposition
To learn an approximation RG, we now parameterize an approximate target rate of the form <x <- v)) 1-a, fo t)(i, v) where fo(a, t), is learned approximation to the poste- rior P(z u xe ===x)) xx)
The target rate can then be characterized by a variational objective that measures the discrepancy between the true and approximate path measures.
Proposition 5 (Variational Loss for Masked Diffusion). The loss function is defined as:
L[R,] = [|/ Ezizz
de
Σ t=0
1
m) log fe(z,, t) xi] dt,
at
uniquely minimized when Rg = R4, and is connected to the terminal KL-divergence by:
(pi(Pilée)) < L[R,].
where Pi is the approximate data distribution generated by R.
Proof. Let P and P be the path measures associated with the continuous-tima Markov chain of the target rate matrix Rd in proposition C.1 and approximation through neural network ft. The
20

---

Published as a conference paper at ICLR 2026
variational loss follows by expanding the KL-divergence between the two path measures.
24) (oe, dt + log (
(IP P) Ee
Re Re(
t: pa
Re( v) log Rd Re( w) Re(xr. v) dl
Eire Σ
Ra(etty)
E
Re(
ái
1 Σ
1 m) log fo (i, de
where the first line takes the expectation the Radon-Nikodym derivative between the two path measures. The statement of Radon-Nikodyne derivative between two CTMCs can be found in the Appendix Campbell et al. (2024). A discrete-time equivalent derivation can also found in Shaul et al. (025)..
The terminal KL bound then follows directly from the data processing inequality, that is:
(PL Φ1) PKL (PI[P)
This technique is standard, as shown in Vargas et al. (2025) for the case of path reversal- construction of diffusion generative models, and in Holderritth et al. (2025a) for discrete diffusion.
D JOINT DISCRETE STOCHASTIC INTERPOLANTS: DEFIITTIONS AND PROPOSITIONS FOR SECTION 3
Building on the discrete stochastic interpolant, we proceed to construct a discrete diffusion that models probability distribution whose supports span variable-length sequences.
On a high level, we would like to define an interpolant constructed by deleting and masking sen- tences from the data distribution. However, the corresponding interpolating rate becomes cumber- some to characterize, as it is no longer clear what each mask token should unmask to.
To this end, we introduce the joint interpolant that allows us to construct broader class of inter- polants and interpolating rate matrices by augmenting the interpolant with auxiliary information that allows us to specify more flexible interpolation path. We then leverage this newfound freedom to construct the flexible-length masked diffusion model.
D.1 JOINT INTERPOLANT
By introducing an auxiliary variable coupled with the interpolant, the joint interpolant expands the class of interpolating rates that can be defined.
Definition D.1 (Joint Interpolant and Joint Interpolating Rate). Let To ~ Po and D1 P1. A joint interpolant is family of coupled random variables {(x,, 8t)),, defined on common probability space and satisfying the boundary conditions It=o = In and It=1 = F1, for which there exists a continuous-timg Markov chain with bounded, time-dependent transition on the joint state space such that, for each t E [0, 1]. the conditional law Law(It, St. To. 21) coincides with the marginal distribution time of this Markov chain started (Iu, 80). We call K jojoitt interpolating rate matrix.
matrix K K2
Proposition D.1 (Joint Interpolant Target Rate), Let {(r,, St) seeoo be a joint interpolant with joint interpolating rate matrix K Consider the continuous time Markov chain with initial distribution po and target transition rate matrix RG defined by
R₂ v) Esgaro
Σk. Ke ((x, (1,
L'ES
where S denotes the discrete state space of the auxiliary variable st. The marginal of the chain at time is then equal to Law(x.).
21

---

Published as a conference paper at ICLR 2026
We note that this result is equivalent to "Flow Matching with Auxiliary Process" in concurrent work Havasi et al. (2025).
Proof. Let pr(s rro, x1) the joint pmf of zt, 8t conditioned on zo, 41, let (o 2 sx/24) be the joint of cp. 11, 9t conditioned on Ig, and let @ (No, x1) to be the joint of Tu, T1. We proceed to verify R₂ satifies the KFE in Equation 3
Σ Re(y,x)p((v)
Σ
Σ
( (z. at = pi(y)
KRowi
y
Σ
Σ
Σ
K ((V, (1r, 8 ( selv)p((u)
y
$1.20:11 s'ES
Σ
Σ
K
((), st), (z, s x1, 2z
y StiZa S'ES
Σ
Σ
K (1), (2,, 8 20 (2
$1,70.11 s'ES
Exo.a1
Edipi(z, |0. 120.21)
This concludes the proof.
D.2 PLEXIBLE-LENGE MASKED DIFFUSION
We then instantiate the joint interpolant obtain the length-aware masked diffusion model, using a sorted list of indices that has been inserted. Again, we drop 2o in the writing as the model in- terpolates between a point mass at an empty sentence the full data distribution. We redefine the interpolant in equation for clarity.
Definition D.2 (Flexible-Length Masked Diffusion Joint Interpolant). Let x = (af, ...,. - ~ P1 be sequence length Let a4 and 3t be monotone and differentiable schedules on [0. 1] such that ogg Ao 0 and a₁ 31 Define insertion and unmasking times [Ti (T) )-
follows:
Ti de dt,
3.
1 1 -rg
dl.
At each time t E [0, 11, define the sorted index set 8t as:
9; {i {0, ... n 1} t Ti}, 11. and boundary values:
with ascending order sc[ < <...
(le(()(s)) st) <...<
s.|-1] 11, se (e(ee((s))) n,
and define the interpolant state 24 per-coordinate as:
it < T 1
ift Z 1,
The process (st,2))ee(o0) the flexible length masked diffusion joint interpolant.
Here, 8: tracks the ordered set of indices whose tokens have been inserted. The interpolant reveals the true token of only after both insertion and unmasking. Given access to this ordered set, one interpolating rate is:
Proposition D.2 (FlexMDM Interpolating Rate), A joint interpolating rate matrix Q7 for the joint interpolant above is given by:
22

---

Published as a conference paper at ICLR 2026
1. Unmask: For index set &ils x² = m. and E E:
3.
Q ((x,, 8), (x(a + u]. = 1 B
1 1(a1 4/1
= v)
Insert:
For index set 8, xi xila $ 8, and position such that s|i-1] < < a[1]:
Q (z) (a & m, (1) 1- at
Proof. We first write down the p=(:;-1a1) the conditioned pmf of (s.x) given 21. Let n
len(2,) (zt), then
2 x1) A(t)) II 4(t),
SE
A(I) = (1 a4) --11
4(f) = /'' du
1 31
Be Bu
1 1xx² m3 +
11 zil du.
Bu
Bu
Differentiate using the product rule:
0,pr(s, æ1) A(t) II ((()) + A(t) Σ
(i,(t) II L(t))
iE*
jEs
For A(t) = (1- ax) -aaaa)nle n-1a we have
(n (n-1k) 18(]dan
A(t) (n - (s)ddel az) 11a) =A(t(
1-
For each i E apply the Leibniz rule to [t)):
i((t) de 1f2 m) f du
- Br 11x² m) +
3e
1(x zi} du
1 Bu
1 Ba
à 1(2" m] ( (111(xx m] +1(2' x13) 1/' du du.
1 Br
Substituting A(t) and it) into the product-rule expansion yields
(spr(5, x2)
=-(n -Is))á a4) a- 1- (tt)
(d11fr' m] + Bo 1fz m) + 1f22 = ri)) /-' àn
+(1 - a1) n-|=1 n-1ol Σ
du
1 - 3a
jEs
II I(t)
vEADD)
(n
|s])
d+ A(t) (,(t)
Σ
3.
A(t) II T,(t)
1
1- B
jE
IE9
Σ
3 A(L)(
1 3r
du
T₂(1) + Σ
d+
(1 (x) A(1)
II 7,(1)
1 B
1
Bu
04
J6 fm
S (3)
iEs
tE
23

---

Published as a conference paper at ICLR 2026
This can then be rewritten term by term as,
0 [11) =
Σ
dt
Pr
x1) Σ
A.
pr(s, x x1)
04
1 - Br
x m
Σ
Pt
(8,
ml
[x1) + Σ z
di
Pc(s
[0[1]] remove(s, i) (1)
Be
az
xt fm
where remove(x; i) refers to the string constructed by removing the i-th element of
Notice that this is equivalent to the R.H.S of the KFE (Eq. B) if one uses the rate matrix Q The four terms correspond 1) Outgoing mass from insertion; Outgoing mass from unmasking; 3) Incoming mass from unmasking: 4) Incoming mass from insertion. This concludes the proof.
[
Proposition D.3 (FlexMDM Rate Matrix (Restated from Proposition 2 By Proposition induced marginal target rate Rg is:
D. the
E: 1. Unmask: For x²
B.
Re(xxx00 < ul) 1 3.
P( =v xe = x))
Insert: For position i (0,..../<<11:
Rd (2, 4; m) = Ot
Exg [se se(i - 11
Cz
Proof. The proof follows by noting proposition D.2 and invoking proposition D.I
Performing an approximate target rate matrix Rg in terms of an approximate posterior by token x æ) and an approximate number of insertions go (I, [)] RJ = = x.. The target rate matrix can be learned by minimizing the follow- ing variational objective. Note that the variational loss objective below is the same as one we defined
E., [1] s s ( (+
fo(x,t) ( u] 2
[
equation] 7
Proposition D.4 (FlexMDM Loss (Restated from Proposition Ty. The loss function is defined as:
L[Re] = =1, Exiise
1s m) log fo(z,, t) 2
dt,
B
k=0
+ $f( Ex
(s - seli - -1] -11 g0 t)(t))
dt,
Se,
Ol
i=0
where (1x, y) log 2 g' uniquely minimized when R. R₁ and is connected by terminal KL-divergence by:
DRa(pi (1pi) S I(fee,
where Pi is the approximate data distribution induced by Rt.
Proof. Let P and P be the path measure of continuous time Markov chain starting with the empty string with rate matrix RG and &. respectively.
24

---

Published as a conference paper at ICLR 2026
Consider the KL-divergence between path measures P and PP,
1=1
Rs 5t) log ( IT)
D
PP Ex
Re( 22) (24, dt +
1:
Σ Re v) - Re(x,, v) log Re v)
Exg
Σ Re(
dt
Re Re( v)
len(
-1
Be
1(x = m] log Sa(1,, (i, sf(l0]] 100
dl.
1
31
+
Exi,
szli - 1] 1. go (2 t)( ((s
The terminal KL-bound then follows from the data processing inequality, that is:
DEL (p1) DEL (PIIP)
DETAILS FOR SECTION
PRECISE DETAIL ON THE INFERENCE ALGORITHMS
In this section, provide details the unmasking steps of vanilla and adaptive inference for FlexMDM, summarized in Algorithm [ Subroutine 2. Suppose at inference time we are given the discretization step size т, partially observed sequence X1, and the current time step tg.
Vanilla inference. For each masked position (i.e., X m) and each clean token E, sample unmasking events from Poisson distribution Poisson( (R,T), where Rs is the unmasking Deg rate toward token Concretely, R. I-Btz fo 4x) [i, v). so that the event count is distributed
as k, Poi I
A. I-
(i, ul) masked position is unmasked only exactly one [o -
token produces count k. while all others produce zero. This tau-leaping scheme batches all events that occur within the interval [4x, L + 7].
Adaptive inference. We first draw the number of tokens to unmask, denoted by an integer K. shows that the choice of K does not affect the theoretical guaran- tees, in practice, we set K to match the expected number unmasked tokens under vanilla inference
While (
6rg # masked tokens in Xx,
yields stable behavior. Accordingly, we sample ~ Poi I-P.,
Next, we compute confidence score for each masked position, based heuristics such as:
Top-K probability (Chang et al. 2022 Zheng et al. 2024): For state x at time I, the confidence at
position i is given by max.es Je(z, t
2 oj.
Top-K probability with sliding window: We further restrict sampling to the leftmost L tokens, where
= min( L2,
with 71 and ?22 hyperparameters. This approach is related to sani-autoregressive strategies used in Nie et al. (2025).
Finally, we select the subset positions to unmask as the Top-K masked indices with the highest confidence scores.
25

---

Published as a conference paper at ICLR 2026
E.2 PROOF OF FLEXMDM's ANY-ORDER INFERENCE CAPABILITY
E.2. PROOF PRELIMINARIES
Form of posterior. We first compute, for each x/. the probabilities of being masked or deleted in
equation 6 This follows from a straightforward calculation using the joint distribution of (T). T:):
p(a) (empty)) = p(T > f) Ckty
Be 1- ( xá, dsdu =: 1-
p(a = m) = p(T S tt T3 >
Here we define % as 11-
m) Therefore, the process in equation (6) is equivalent to
observing partially masked subsequence 24 obtained by sampling x1 and, for each position of 31, independently deleting it with probability 1 - a,, masking it with probability -t or leaving unchanged with probability a, + 7 1. Note that at and 7t both increase from 0 to increases from 0 to 1.
The posterior is given by
p(11= = x= 22 = z)
x p(x") ) p(ax = x x1 = x²)
= -(x*) (1 - -a.) (enl(e((a)) )ol)(()) (1 #mask(m) (are + - #ummask(a)) #{s: S Cx"1+
x p(xx)) - (1-a,) (en)))en""" # ss: : C x1).).
Importantly, as in the vanilla MDM setting, the posterior does not depend on the unmasking schedule (1t) (thus 3,). which will enable us perform unmasking in adaptively chosen positions. Note also that if all sequences in the support of were of the same length, this posterior would also be independent of (a:); while we do not prove it, in this case this would allow us to choose an arbitrary order of unmaskings and insertions.
Extension of posterior to t 1. Motivated by the form of the posterior above, we define the following:
(x xx) [p(x*) Laga. (p(x*) (a,-,,)ee(et ) -leu((x) #1s : x C cx otherwise
if =
this the same as p(11 x 24 x)) We will denote the marginals of (- .) by (- | x) for E E. The reason for extending the definition of the posterior to t = that an adaptive FlexMDM sampler (see Definition I). because we are entirely decoupling unmasking from the schedule of insertions, after the final insertion step the time parameter / may be even though there are still tokens left to unmask. We will assume oracle access to q1 | x) as in practice these simply the anyerrder marginals for and furthermore in practice these are already well-approximated by the learned posterior marginals p(z) 21 x:) for arbitrarily small 8 > 0.
Note that for
Index-tracking variable. Recall that in the definition of the joint interpolant we defined an index- tracking variable 8 which essentially tracked which indices of ø correspond to the tokens in ø;- While our analysis below will not use the language of stochastie interpolants, idea tracking St. with slightly modified notation. Specifically, for any 0 S f. < 1, we will the notation s)) and (e) to denote probability and expectation with respect to the distribution given by sampling .1 from qe(- I), and then sampling uniformly random from subsets for which x C (z1)ls- When we only care about the marginal distribution over s, write Praja,=a and Egj,, Lala,mar ma- Given such subset and i (0....,Is| - 1}, use s; to denote its i-th element in sorted order; as before, we also define the boundary values 8-1 = - and $len(s) len (a))) ). The insertion expectations which we had denoted by [[+e[i] s, - 1] 1] in the main body thus given by E,/2,= 8;-1 1] in the notation of this section.
will still use the
E.3 FORMAL GUARANTEE FOR ADAPTIVE INFERENCE
Definition E.1. Given query access to the posterior marginals g- x, = x) and to the insertion expectations Eajz,zz 8i-1 1), an adaptive FlexMDM sampler is any algorithm which produces
26

---

Published as a conference paper at ICLR 2026
sequence of iterates It15...,.24 Bt. where 0 t <... in 1, by starting at Dt, and
arbitrarily alternating between steps of the following form:
Any-order unmasking step: Starting from Lt,, if ((aask(tt)) is nonempty, pick arbitrary index i
therein (possibly probabilistically) sample o from qi, ( &t,)) and Tt, St, 18,, +- v).
Insertion step: Starting from Ft,- run the CTMC with rate matrix
E Es (8t 8:-1 1] T-as
if = <; m
R (I,
len(z) 2im Ri x 4 m)
y = x otherwise
fort, Ststttt obtain Degts Iftit1 1, then apply any-order unmasking until mask is empty, and terminate.
Note that the rate matrix in the second bullet point above is identical to the one in the main body except that we only consider transitions given by insertions.
Formally, will show the following:
Theorem E.1. Any adaptive FlexMDM sampler for will generate a sequence of iterates Cey..... St, such that λt, is exactly sample from p.
E.4 PROOF OF THEOREM
To show that adaptive sampling works, we inductively prove an even stronger statement: at any intermediate step of the sampler after it has produced fty. the final output é.. is sample from 4ty (: Bt2).
The following two lemmas provide the inductive steps for unmasking and insertion respectively: Lemma E.1 (Inductive step for unmasking). Let O S S 1 and let æ be a partially unmasked Let = (#z))emask(z) denote any distribution over masked indices of x. 1. Sample index from
subsequence of length
Suppose that one runs the following:
Sample u from the posterior marginal gi(: (: xz = x)
3. Sample from qe(- x/æ² +- vl).
The output of this procedure is a sample x² from qi x))
Lemma F.2 (Inductive step for insertion). Let O < t < 1 and let x be a partially unmasked subsequence of length m. Let 0 S S be any duration of time. Suppose that one runs the following:
Starting from x, run the CTMC with rate matrix given by Eq. for time h to oblain x'
2. Sample from ( x²).
The output of this procedure sample from the posterior q |x)).
We defer the proofs of these to Sections E.5 and below. The idea for the former is identical to the proof of the folklore fact that vanilla MDMs can sample in any order (Kim et al. 2025). The proof for the latter is more involved and involves explicitly verifying that the Kolmogorov backward equation is satisfied by the rate matrix we have constructed.
Here we verify that these Lemmas are enough to establish Theorem
Proof of Theorem E.1 We show more generally that starting from any intermediate time step Lt, (not just j 1), any adaptive FlexMDM sampler outputs sample from qt (2,)) We do this by inducting on the total number of insertion steps that remain.
As the base In this case, we can further induct on the number unmasking steps and apply Lemma E. with therein set to conclude that the final output is sample from qı ).
for the induction, if no more insertion steps remain, then must have 1.
27

---

Published as a conference paper at ICLR 2026
For the inductive step, we have < 1 and suppose have shown that for any FlexMDD sampler that makes at most m insertion steps, starting from any žt, at intermediate time t,. it samples from Jty (- B2y) ). Now consider FlexMDM sampler that makes most m insertion steps starting from ft, at intermediate time 1-- If in the next step it performs an insertion step, i.e. it runs the CTMC with rate matrix defined above for total time h
1,+1 l,. then by Lemma F.2 and the inductive hypothesis, it samples from 4t, (- Zt,). ). Alternatively, suppose the sampler performs some sequence of & unmasking steps before performing an insertion step. Then by further inducting on L, we conclude by Lemma E.1 that the sampler eventually outputs sample from 4t, (- Zz2).
Finally, the theorem follows from the special case where t, and Bt, E.
E.5 PROOF OF LEMMA
Proof. Fix any index i E mask(r). The marginal g x) is given by 3:(11) E.,,s((n))) p(x) ((a1)),,, 24 (em(ee)) #(( #(S (3))ls)
Ex,
Pr (2r1)ss; =
En ((x))(1 ar) [en(m)))oon()) #(S (x1)
(10)
The posterior qe(- xfxx* t- u|) is given by
p(x*) (1 0x) (e,))le(("))eenee (S xxxx o| C x ss1
qe(x" a|lx + o]) Ex, p(z1) ) (1 - ax) lom) ton(z)
(11)
#(S x ex t w| S (xi)ls}
Note that the numerator of Eq. (10) and the denominator of Eq. (TI are the same. So conditioned on unmasking index i, the above procedure outputs x" with probability
Σ
Pr S))z, (x11)l: u) 90(x" |x|x² v))
p(x*) (1 lon() Eg #1S x/00 < u| Ca'ls) Σn p(x1)(1 )en((( len((()) #IS C (xı)ls}
9e(2 x))
This holds conditioned on unmasking any E mask(z), so regardless of the choice of a over such positions, we will sample from the correct distribution q-(- ar).
E.6 PROOF OF LEMMA
Proof. It suffices to show that the rate matrix satisfies the Kolmogorov backward equation
len(xr)
0gge(a (x)
R (a, xr < (x) x 4 m) Ri (ar, x) (x" [x).
==
First note that the rate Rinx 4 m) can be expressed
En, p(x1): (1 ag )1o((1)) Es - - ā
En, p(x1) )- a, (en))e((1)-een(t len(a)) #IS (æ1)ls) 1 a
Furthermore, E, - 3,-1 - 1) = len(11) - len(x), so
Σ R (I, <; m)
Σn, p(x1): (1 a, (en))t #1S S (æı)ls) (len(z1) - len(x))
(12)
p(1)) (1 (en0s))))) (SS S (a1)ls) 1 1- az
28

---

Published as a conference paper at ICLR 2026
Let us compute deqe xx)
p(x*) (1 ax) (lol))" #(S x [s} (len(x*) -(e(())))ee(r))
dee
En p(x1) -1 ) # z _ (x1)
1 00
Ex, p((1)) (
(2) |ss} (len(x)) len(x))
En, p(x1) aa) (,))) #
(2i)ls}
1 0x
X p(x*) (1 - (,,))enne ( #1S x C x [s}
(13)
my p(r1) on)) # (S J C (#1) ls}
Note that by Eq. 12). the term in the parentheses in Eq. 13 is exactly
len(z) R R x 4; m)q((=" (x) = - -12 12in (x,, x) 1x), (z,
It remains to verify that the first term in Eq. 13 is equal to E, (x, m) | m).
To that end, we must show that
len(xr) En, p(x) (1 or (e.))eee))))nnii Jen(()) Σ Sta (wn) (8: 8:-1
#1S i m C
Ex, p(x1)(((()))))) )) an) [lo((z)) lon(z) #(S 4 S (21)ls}
= (S : x S (x*)]s} (len(x*) - len(x)) (14)
The key combinatorial step is as follows. For any x, in the support of p. consider subset for which x C. (xi)ls- Note that for every such S, we can uniquely associate exactly 8; 8i-1
different subsets S' of size |S|+1 1 for which x 4, m S (#1)ls- Therefore, Es.5( (s, S,-11
1) = #1S : 4 m C (s1)ls}, and the left-hand side of Eq. 14 thus becomes
le)(
len(x))
#1S: 4 m [z
Σ
(st - 3 - 1) =
Σ Sxx((s)
(len(x") - (x+))
1=0
S
which completes the proof of Eq.
F EXPERIMENTAL DETAILS
F.1 PRETRAINING ON OPENWERTEXT
Dataset preparation. As mentioned in Section 5.1 to obtain a variable-length dataset, we split OpenWe articles paragraph-we wiss using the GPT-2 Radeforr & Wu (019). This can be implemented by locating the token index corresponding to the delimiter \n An. Sequences longer than 1024 tokens are then chunked recursively by splitting by the delimiter closest to the middle sequence, yielding a variable-length dataset with maximum sequence length 1024.
FlexAttention. To handle variable-length sequences during training, we pad each batch to the maximum sequence length. In MDM, by design, pad tokens also enter the model input. In contrast, FlexMDM is designed to receive only clean or mask tokens as inputs. Ideally, QKV attention should attend to pad tokens; however, current FlashAttention (Dao al. (2022) does not support this for non-causal attention (our setup of interest). We therefore adapt Flex (Dong et al.] 2024). A side benefit is improved training speed, since FlexMDM performs attention over fewer tokens than MDM's full-sequence attention. Note that in the LLaDA experiment, we did not implement this optimization; pad tokens can therefore attend other tokens, though we expect the impact to be negligible.
Training configuration. As in Section 5.1 we model FlexMDM with DiT Peebles Xie (2023) backbone and embed the insertion schedule ay. For MDM, we use the same DiT configuration with time step embedding but without the softplus scalar head. Transformer configuration is: hidden ze: 768, heads: conditions dimens on: 1 128, dropout: 0. 05, layers: 12.
29

---

Published as a conference paper at ICLR 2026
We train both models 16 H100 GPUs with global batch size of 1024 and max training iter- ation 1M. We use the AdamW (Loshchilov & Hutter 2017) optimizer with weight decay 0.03, learning rate 3e-4, 2000 warmup steps, and an EMA factor of 0.9999. Additionally, discrepancy time-step sampling to reduce variance: one f. is drawn uniformly from each interval [i/T, (i + 1)/TT as in prior MDM training (Shi et al. 2024
low-
Metric. For evaluation, we take sequences generated by both models and retain the clean tokens by removing padding (e.g., the leading pad token). We adopt LLaMA-2-7B (Touvron et al. 2023) as the reference model to compute likelihoods. We notice that the pretrained MDM generates short sentences with unreasonably large (worse) perplexities. Therefore, filter overly short sequences with < 10 tokens when calculating average perplexity.
F.1.1 ADDITIONAL EXPERIMENTS
Sequence entropy. We follow the setup from (Zheng et al. 2024) to measure the entropy of given sequence. As we can observe in the table below, both models exhibit comparable sequence entropy across different sampling steps.
Sampling steps
256
512
1024
2048
4096
MDM
4.4732
4.5691
4.4417
4.2259
4.1876
FlexMDM
4.4586
4.5173
4.5660
4.4712
4.4304
Generative perplexity for extended sampling steps. Below, extend the generative perplexity measurement to more sampling steps. We observe plateaus around 16,384 steps. This is likely be- cause discretization errors for the CTMC and the estimation errors for each transition entry become negligible, leading to diminishing generative perplexity.
NDM
255
0.2
10:4
20.6
4236
1000
1034
3574
SempRacuige
Downstream evaluation on the text understanding. We follow the evaluation protocol of von
Rütte al. (2025).
Model
HellaSwag
PIQA
ARC-E
ARC-C
BoolQ
OBQA
WinoGrade
FlexMDM
28.75 29.70
51.69 52.55
25.26
25.75
50.50
24.60
51.78
MDM
25.26
26.09
50.05
23.80
49.01
F.2
PRETRAINING ON THE MAZE PLANNING TASK
Task construction. We generate mazes with fully random, recursive division procedure (the code is provided Listing [). on a 41 x 41 grid, with some invalid cells. As described in Sec- tion 5. we consider subgoal-conditioned planning task: the model is given sequence of sub- goals (g1...., gK) and must produce a valid path that connects them in order. To construct training examples for given K, sample 15000 start- goal pairs (g1, 9k), compute the shortest path for each pair via breadth-fi first search, and then add controlled detours to obtain up to 10 distinct valid paths per pair. Subgoals are formed by selecting K intermediate cells uniformly at random along
30

---

Published as a conference paper at ICLR 2026
chosen path (start and goal already fixed). We use 95% of the pairs for training and hold out 5% for validation to evaluate generalization to unseen pairs and subgoal sets.
Training data construction. For MDM, the training sequence ((g1..... ........9.. 9K) [SEP] Path), where [SEP] is special separator and Path denotes the tokenized path. During training, the prompt (g),..... 99K) bypasses the interpolant so that, at inference time, the model can condition on (g),.... . 9 [SEP] and generate the path. For FlexMDM, we use an interpolant in which the subgoal tokens are exempt from the process in (6): that is, tokens corresponding to each 9i are kept clean at all times. This also enables generation to start from (g1. (91..... 9K Although this conditional generation template changes the base distributions Po for both MDM and FlexMDM, we note that the theoretical guarantee from Section 3 and Section 4 still holds once the training is perfect (under the access to the ground truth unmasking posterior and insertion expectation), the inference algorithms recover the ground truth distribution Pt.
Training configuration. We use the same architectural design in the OpenWebText pretraining, but with a smaller model: hidden size 256, heads : 8, conditional dimensionn ion: 128, dropout: 0 1, Layers: 8. Both models have approximately 8.5M parameters. We train them on A100 GPUs with a global batch size 256 for up to 100 epochs. We use A (Loshchiloy & Hutter 2017) with weight decay 0.01, learning rate x 10 20 warmup epochs, and an EMA factor of 0.9999.
Metric. Given the final conditionally generated sequence, we report the success rate under two criteria: (1) all visited cells are valid (no collisions with invalid cells), and (2) the path connects the subgoals consecutively order. We perform inference both models with the number of sampling
steps 256.
Listing Code for the maze Construction
RECURSIVE DIVISION (perfect maze)
def _divide (g, top, left, h, w):
if h <= 2 w <= 2:
return
horizental if horizental:
split the longer dimension
random. randrange (top 1, top + h 1, 2) gap random. randrange (left, left 2)
gly, left:left
gly, gapl =
divide top, left, - top, w) divide + 1, left, top + h - 1, w)
clsc:
x random, randrange (left 1, left w 1. 21 gap random, randrange (tep, top + h, 2) g[top:top h, x] =1
g(gap, x]
divide (g, top, left, h, left) divide top, x 1, h, left + 1)
WRAFFER with extra doorways
def divisicn_mant (m, seed=2025, extra_door frac=0.
m, seed
size in scells*
(not bitmap coords)
int or None
dsor__door_frae
# 0, perfect maze; s0 flicks more doors (loops)
need (seed) H, 1,
31

---

Published as a conference paper at ICLR 2026
g ne.zeros( (H, W,, atype=np.uint
g(o, :1, g(8 1, :1, gli, 0l, gl:, w 1, 1,
divide (g, 1, 1, H 2, W 2)
optional imperfection: punch extra doorways if extra extrr_ddor__ree ddor_fraaa > 0:
candidates = 1]
for r in range (1, H - 1) for in range (1, if glr, c] # Check if wall separates two passages orthogonally
1):
if g[r - 1, cl -- 0 and gir + 1, cl == 0: aendiddaess.ppendde
c))
elif giz,
1] 0 and giz,
0:
candidates. append c))
k = int (len (candidates) * extra_door_frae for (r, C) in random. sample (candidates, k-k)
return glr,
F.3 WEIGHT INTTIALIZATION TRAINING FROM LLADA
In this section, we describe the procedure for adapting the pretrained LLaDA-8B base model into the FlexMDM.
Training configuration. LLaDA uses bidirectional Transformer without an additional time em- bedding, leveraging the fact that MDM does not require an explicit time signal (Zheng et al. 2024). For FlexMDM, to model the insertion expectation, we inject a time-embedding pathway via AdaLN (Peebles & Xie 2023). For parameter efficiency, we tie (share) the four AdaLN parameter sets across intermediate Transformer layers. We also attach softplus scalar head parameterize the insertion expectation.
Next, we attach LoRA adapters (Hu et al. 2022) to every attention projection (q-proj, k.proj. v-proj) and to the unmasking-posiin the unmasking posteriors differ between MDM and FlexMDM: FlexMDM, unmasking must ac- count for token shifts induced by insertions. This fine-tuning of the last head is crucial for enabling variable-length generation.
head. We include LoRA the unmasking posterior head because
The LoRA configuration that use is 128, 128, and dropout 0.1. Training runs for 200, 000 gradient steps with a batch size of 64 on 16 H100s, which took approximately 3 days. We optimize with AdamW (Loshchilov & Hutter 2017) learning rate x 10-4 and weight decay 0.1, using a cosine warm-restarts scheduler.
Evaluation GSM8K. We instruction-fine-tune (IFT) the FlexMDM base model the GSM8K training split. To preserve the intrruction-answee format, we modify the interpolant in Eq. 6 that tokens corresponding to the instruction are excluded from the interpolant- these tokens remain fixed for all time steps. We apply the same strategy to obtain the baseline (that is. IFT from modifying the MDM interpolant so that instruction tokens remain fixed. This choice is equivalent to the IFT recipe used in Nie et al. (2025); Ye et al.] (2025). Both models are IFT-ed for 1000 epochs. (Other IFT hyperparameters match those used in our base setup unless otherwise noted.)
For FlexMDM inference, we start from the instruction prompt / = 0 and run adaptive inference to Concretely, we use Top-K probability with sliding window (Appendix F with 71 = 5.0 and ?/2 = 64. For LLaDA, report the best result under the sutm--uttorerrssivve confidence-basee sampling of (Nie et al. (025).. For both models, we set the token sampling temperature to 0.0, which confirm be important for strong Pass @ 1. Overall, adaptive inference substantially improves performance over vanilla inference.
For fair comparison, since FlexMDM is not IFT-ed, IFT LLaDA-Base. rather LLaDA-instruct, this differs from Zhao et al. (2025). We employ zero-shot evaluation, which also differs from Nic al. 2024).
32

---

Published as a conference paper at ICLR 2026
Evaluation HumanEVal-infill. Code infilling conditions prefix and suffix, and asks the model to complete the middle part of the code. For FlexMDM, we format training exam- ples as (Prefix; [SEP]; [Middle]; [SEP ;; Suffix), where [SEP] is separator token and Instruction ion describes the infilling task for a model. We modify Eq. (6) so that tokens for Prefix, Suffix x, and [SEP] are fixed in the interpolant (i.e., excluded from the interpolant). Thus, at f = 0 the state is (Prefix; [SEP]: ]; SEP]; Suffix).
For MDM, we the format (Instruct io;; Prefix; [SUF Sufufix [SEP]; ; with Instruction puuciion prompting infill after prefix and suftix, along with [PRE] and [SUF] sepa- rate the prefix and suffix, respectively. Here too, the tokens without Middle are held fixed by the modified interpolant. This difference in formatting reflects the fixed-length nature of MDMF MDM (no token insertion). This formatting has been used the code infilling tasks for autoregressive models in Bavarian et al. (2022). Naively masking the Mi ddle e span yields base state at / = 0 of (Pref x; Masked; Suf x), where Masked This leaks length information -atemattoonmmttriilllly materially simplifying the task- and renders comparisons to FlexMDM unfair, since FlexMDM does not observe the target span length. For fair evaluation, therefore avoid length-revealing masks and require methods to infer the span length during inference.
fully masked sequence of length [Middee] el
We IFT both models the educational-instie split of opc-sft-stages the architecture and optimization configurations match those used for GSM8K IFT. At evaluation, we initialize from the base distributions: for FlexMDM, (Instruction; Prefix; [SEP]; [uf for MDM, rInstruction rectructiin ion; Prefix; Suffix; [SEP]). )). We use the same Top-K adaptive inference for both and temperature 0.0. Final outputs scored with the HumannVvllinfflll verifier toolkit compute success rates.
33


Reviews:
[{"review_id": "3nccVHl1Uw", "reviewer": "Reviewer_J8NK", "summary": "This paper proposes Flexible Masked Diffusion Models (FlexMDM), a discrete diffusion framework that adds token insertion to standard masked diffusion models while provably preserving any-order decoding. It introduces a joint interpolant over token values and indices. The proposed model learns both the usual unmasking posterior and a new insertion expectation, and derives a CTMC-based training objective with guarantees. Empirically, FlexMDM matches MDM perplexity while modeling length distributions far more faithfully. Initializing from a pretrained MDM model (LLaDA-8B), fine-tuned FlexMDM exhibits significant performance improvements on GSM8K and HumanEval-infill as the number of sampling steps increases.", "questions": "Please see Weaknesses", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper proposes a novel and principled approach to addressing variable-length generation, which is a known limitation of existing masked diffusion models.\n- The proposed approach is theoretically grounded: a joint interpolant is defined over both token values and the number of insertions, then a neat training objective is derived based on CTMC, as a natural extension to the standard MDM training loss. \n- This paper also proves the compatibility of insertion prediction with adaptive inference in MDM\n- The empirical evidence is compelling. FlexMDM models the length distribution much more faithfully than MDM without sacrificing perplexity, and demonstrates better scaling performance on GSM8K and HumanEval-infill", "weaknesses": "- It would help practitioners more if the paper could provide more descriptions of the implementation details. E.g., after inserting/deleting tokens, will the padding be dynamically adjusted?\n- Some additional questions \n    - Figure 3 left: I guess $S_{\\tau}$ shoud be $[0, 3, 4]$ right?\n    - In algorithm 1, positions are unmasked first before insertion prediction, but I suppose the newly unmasked tokens should not be visible to the insertion module $g_{\\theta}$ right?\n    - Given the flexibility of combining any-order unmasking with insertion prediction, I wonder whether it's possible to simply reuse the unmasking posterior from a pretrained MDM and only separately learn the insertion expectation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JeHd28tyfz", "reviewer": "Reviewer_VJ5y", "summary": "This paper proposes **Any-Order Flexible-Length Masked Diffusion (FlexMDM)**, a generalization of masked diffusion models (MDMs) that supports variable-length generation through a continuous-time Markov chain formulation. The authors introduce an *insertion process* parameterized by a learned insertion expectation (g_\\theta) in addition to the usual denoising posterior (f_\\theta). This enables length-adaptive, any-order generation while preserving theoretical guarantees and compatibility with pretrained MDM weights. Experiments demonstrate strong results on text, maze planning, and code/math reasoning tasks, including scaling to an 8B-parameter model.", "questions": "See weakenesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. **Clear theoretical formulation.** The continuous-time treatment and derivation of the rate matrix and training objective are elegant and rigorous. The paper provides the right amount of mathematical depth while maintaining readability.\n2. **Minimal yet effective extension.** Introducing a scalar insertion expectation per gap is an elegant way to achieve variable-length modeling without compromising the core MDM formulation.\n3. **Scalability and empirical validation.** The ability to retrofit large pretrained MDMs and obtain consistent gains across domains (text, code, math) is a strong practical signal.\n4. **Well-designed sampler.** The adaptive unmasking and τ-leaping integration are both theoretically justified and empirically efficient.", "weaknesses": "0. **Missing support for addition operation**.\nCompared to editflow, it supports addition while this work doesn't. \n\n1. **Missing discussion of Seed Diffusion.**\n   The paper should reference and compare to *Seed Diffusion: A Large-Scale Diffusion Language Model with High-Speed Inference* (arXiv:2508.02193). Both aim to enable flexible-length or insertion-based generation, but take different approaches—Seed Diffusion uses a single-model canvas growth mechanism, whereas FlexMDM introduces a learned insertion intensity. A brief conceptual comparison would clarify relative advantages (FlexMDM’s theoretical rigor vs. Seed Diffusion’s simplicity).\n\n2. **Apparent need for two models.**\n   The description suggests two learned functions—one for denoising and one for insertion—which might appear “not elegant.” It would help to emphasize that these are *two heads of the same model* rather than two independent models. The authors could also discuss whether a unified parameterization (e.g., sharing a latent event-rate representation) is feasible, as done in Seed Diffusion’s single-network formulation.\n\n3. **Limited evaluation in real-world settings.**\n   While the presented results are promising, additional experiments on **broader math and coding benchmarks** would strengthen the claims.\n\n   * Include *MATH-500/5000* or *AIME24/25* for compositional reasoning.\n   * Extend code evaluation to *HumanEval*, *HumanEval+*, or *MBPP*.\n   * Report **pass@10** and **pass@k vs. compute-time** to better capture stochastic sampling performance.\n\n4. **Minor clarity issues.**\n\n   * An architecture figure showing the shared backbone and dual heads would avoid the “two-model” confusion.\n   * A brief qualitative example illustrating insertion dynamics would improve intuition.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "l3H0RY39fI", "reviewer": "Reviewer_77ke", "summary": "Current masked diffusion models are constrained to fixed-length sample generation, often resorting to tricks (e.g. padding) to deal with variable-length inputs. This paper addresses this issue by introducing a flexible-length masked diffusion model based on token deletion/insertion in addition to masking/unmasking operations. This behavior is achieved by learning to predict an “insertion expectation” (expected number of tokens yet to be inserted after a given position) in addition to the canonical unmasking posterior. The insertion expectation is parameterized directly as an additional scalar prediction head.\n\nThe resulting method is able to accurately match the ground-truth sequence length distribution, more so than the fixed-length baseline utilizing padding tokens. It also outperforms the fixed-length baseline on a subgoal-conditioned maze solving task, and can be cheaply retrofitted onto a 8B masked diffusion model (MDM), improving math and coding performance for longer sampling horizons.", "questions": "- Q1. What are the generative PPLs of MDM and FlexMDM (175M) without filtering by sequence length? (also see W1.1)\n- Q2. What is the sequence entropy (as in Zheng et al., 2024) of MDM and FlexMDM (175M) for different numbers of denoising steps? (also see W1.1)\n- Q3. What is the training/validation loss (throughout training/final) of MDM and FlexMDM (175M)? (also see W1.2)\n- Q4. What is the performance of MDM and FlexMDM (175M) on some relevant text benchmarks (e.g. HellaSwag, ARC-E, ARC-C, WinoGrande, PIQA, etc.)? (also see W1.3)\n- Q5. What is the training and inference speed of MDM vs. FlexMDM (both 175M and 8B)? (also see W1.4)\n- Q6. How many tokens were used for training LLaDA on FlexMDM? (also see W1.5)\n- Q7. Why are the reported numbers on GSM8k (Fig. 5) much lower compared to LLaDA-base (Nie et al., 2025) despite explicitly training on the task? (also see W1.6)\n- Q8. Will the model weights and/or training code be released? (also see W2)\n\nAdditional questions:\n- Q9. Does the FlexMDM training objective constitute a likelihood bound (i.e. ELBO)? If so, this does not appear to be immediately obvious based on L278.\n- Q10. Were any alternative ways to parameterize the insertion expectation considered? For example, one may consider modeling the number of insertions as a categorical distribution over the number of to-be-inserted tokens and calculating the expectation thereof.\n- Q11. For large-scale pretraining, it is a requirement to have statically-shaped batches for the sake of model compilation and good FLOP utilization. Therefore, in order to apply FlexMDM at scale it will be crucial to have a recipe to do fixed-length training, e.g. through sequence packing. Can FlexMDM support such a scenario?\n- Q12. Is it possible to skip the distinct masked case and directly go from empty to unmasked? This could avoid the apparent overhead of first inserting masked tokens and only unmasking them later.\n- Q13. Why is a poisson distribution the correct choice for inserting mask tokens? As far as I can tell, this is not explained in the paper.\n\nNits (not considered for final score):\n- Some equations are unnumbered (e.g. L184, L230, L259, L278). For the sake of referenceability, it is good practice to number all equations (and not just the ones that are referenced in the text itself).\n- Time notation breaks the convention of t=0 being noise-free and t=1 being complete noise. This is confusing for people familiar with the diffusion literature, especially for those unfamiliar with flow-matching. For the sake of consistency with the literature, I recommend keeping the diffusion convention. Alternatively, a note on this break of convention may help prevent unnecessary confusion.\n- L294: Missing parentheses around equation number. Alternatively, it is good practice to reference equations as such, i.e. “Eq. 5” or “Equation 5” or “Eq. (5)”.\n- L317: Seemingly missing some words in “Since unmasking indices in an adaptive [inference setting] no longer trace [...]”\n- L375: The footnote mark makes the noise schedule appear to be quadratic. Swapping the position of the period and footnote may help, i.e. `$\\alpha_t = \\beta_t = t$.\\footnote{...}`\n- L452: Missing period at the end of the line.\n\n---\n\n### References\n- Nie et al. (2025): https://arxiv.org/abs/2502.09992\n- Zheng et al. (2024): https://arxiv.org/abs/2409.02908", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "### S1. Solves an important problem of MDMs\nVariable-length generation has been a major limitation of MDMs, and the proposed solution is a valuable addition to both theory and practice of MDMs.\n\n### S2. Theoretically grounded approach\nThe proposed approach is grounded in theory and provides some guarantees on the likelihood of the trained model.\n\n### S3. Applicability to existing MDMs\nThe proposed approach can be retrofitted to existing pre-trained MDMs with minimal training (as little as 13B tokens), enabling flexible-length generation for both improved efficiency and accuracy.", "weaknesses": "(ordered by descending severity)\n\n### W1. Meager experimental results\nWhile the set of considered tasks is well-rounded, the selective reporting and omission of results as well as non-standard choices that skew the picture in favor of the proposed method need to be rectified:\n1. Generative PPL should always be reported together with sequence entropy to control for diversity collapse. Additionally, filtering out short sequences (L1575) is non-standard and skews the picture in favor of FlexMDM. This is an inherent limitation of using the gen. PPL metric and should be controlled for via sequence entropy.\n2. Validation loss, and perhaps also training curves, should be reported even if the losses aren’t directly comparable (L416). If the FlexMDM objective indeed is a likelihood bound, then both losses are meaningful and, to some extent, comparable.\n3. Accuracy on text benchmarks should also be reported for the sake of consistency and comparison with the literature. These numbers are comparable even if the losses aren’t.\n4. Numbers on training/inference speed of FlexMDM vs. MDM are missing. In principle, this should be favorable for FlexMDM due to its smaller average sequence length.\n5. The required training time to convert an existing MDM to FlexMDM (L450) should (also) be reported in terms of training tokens, as the wallclock time highly depends on the level of optimization of the given codebase.\n6. The reported performance of LLaDA on GSM8k (Fig. 5) is considerably lower than what’s reported in the original paper. The original paper (Nie et al., 2025) reports 70.3% for the base model and 69.4% for the IFT model. It is unclear how this number could get worse by directly training on GSM8k.\n\n### W2. Reproducibility\nFor the sake of reproducibility, I strongly urge the authors to release all artifacts required to reproduce the results presented in this paper, including trained model checkpoints as well as training and inference code.\n\n### W3. Some fundamental limitations of MDMs remain\nWhile the proposed method enables flexible insertion of tokens during the generation process, some other fundamental limitations remain: Namely, the resulting model is unable to delete or revise existing tokens once filled in. Therefore, the model still risks accumulating errors throughout sampling, just like traditional MDMs.\n\n### Conclusion\nDespite the limited experimental results (W1) and concerns regarding reproducibility (W2), the theoretical contributions are strong on their own (S1, S2). Therefore, reasons to accept currently outweigh reasons to reject. The submission can be made significantly stronger still by addressing W1 and W2, and I will be happy to update my final score accordingly.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mn4NkWUvny", "reviewer": "Reviewer_TnF2", "summary": "This paper proposes FlexMDM, a masked diffusion model that jointly models the unmasking posterior and the insertion expectation. Unlike traditional Masked Diffusion Models (MDMs), which begin with a fixed number of masked tokens and progressively unmask them, FlexMDM introduces an additional insertion operation that dynamically inserts new mask tokens during the diffusion process. This design enables variable-length generation while preserving any-order generation capabilities.\n\nFlexMDM can be adapted from an existing MDM with minimal changes—specifically, by adding a scalar output head at each position to predict the insertion expectation.\n\nThe authors validate their method across multiple language and planning benchmarks, demonstrating that FlexMDM:\n\n1. Effectively learns the true length distribution of the training data,\n2. Achieves comparable generative perplexity to standard MDMs,\n3. Outperforms MDMs on planning and reasoning tasks, and\n4. Can be efficiently scaled to larger models.", "questions": "Please see weakness.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. Writing is very clear and easy to follow\n2. The framework is mathematically sound and theoretically grounded and solid\n3. The idea is novel and the connection to traditional MDMs are very clear", "weaknesses": "I find the experiments relatively weak compared to the strong theoretical foundations of the paper, and several aspects require clarification:\n\nModel comparison: It is unclear what specific model is referred to as “MDM” in the comparisons. For example, is it equivalent to D3PM or another variant?\n\nFigure 4c: The figure (which should include an overall label) shows an interesting trend in perplexity as the number of sampling steps increases. Since neither curve appears to converge by 4096 steps, it would be informative to see what happens beyond this point.\n\nMaze task: While it is reasonable that FlexMDM outperforms traditional MDMs in the maze task, this setup alone seems insufficient to substantiate the claim that it “supports FlexMDM as a principled approach for subgoal-based planning.”\n\nScaling up experiment: The statement “Surprisingly, we observe rapid transfer: within three days on 16 H100 GPUs, the model generates variable-length sentences” requires clarification. MDMs can already generate variable-length sequences to some extent using padding tokens, as the authors note. It would be useful to clarify in what sense this transfer is “rapid.” Also, while it's an 8B model, the authors only trains a 400M LoRA adapter and claim FlexMDM is \"scalable\". Can the authors clarify on this?\n\nAblation study: Including an ablation study on modeling the insertion process would strengthen the empirical evidence. Alternatively, can the authors confirm if the comparison between MDM and FlexMDM already serves this purpose.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)