Forward-Learned Discrete Diffusion: Learning how to noise to denoise faster
Authors: Grigory Bartosh
Teodora Pandeva
Sushrut Karmalkar
Javier Zazo
Venue: ICLR 2026 Poster (2026)
Keywords: diffusion
generative models
variational inference
Primary area: generative models

Abstract: Discrete diffusion models are a powerful class of generative models that demonstrate strong performance across many domains. However, for efficiency, discrete diffusion typically parameterizes the generative (reverse) process with factorized distributions, which makes it difficult for the model to learn a target process in a small number of steps and necessitates a long, computationally expensive sampling procedure. To reduce the gap between the target and model distributions and enable few-step generation, we introduce a learnable noising (forward) process for discrete diffusion. Instead of fixing a Markovian forward chain, we adopt a non-Markovian formulation and introduce learnable marginal and posterior distributions. This allows the generative process to remain factorized while matching the target defined by the noising process. We train all parameters end-to-end under the standard variational objective.

Published as a conference paper at ICLR 2026
FORWARD-LEARNER DISCRETE DIFFUSION: LEARNING HOW TO NOISE TO DENOISE FASTER
Grigory Bartosh University of Amsterdam g. ba www.rrouvaa...... n1
Teodora Pandeva, Sushrut Karmalkar & Javier Zazo
Microsoft Research, Cambridge (tpandeva, ska rma kar, javierzazo) cmiooftt..com
ABSTRACT
Discrete diffusion models are a powerful class of generative models with strong performance across many domains. For efficiency, however, discrete diffusion typically parameterizes the generative (reverse) process with factorized distribu- tions, which makes it diffieult for the model to learn the target process in small number of steps and necessitates a long, computationally expensive sampling pro- cedure. To reduce the gap between the target and model distributions and enable few-step generation, we propose Forward-Learned Discrete Diffusion (FLDD), which introduces discrete diffusion with learnable forward (noising) process. Rather than fixing Markovian forward chain, we adopt non-Markovian formu- lation with learnable marginal and posterior distributions. This allows the gener- ative process to remain factorized while matching the target defined by the nois- ing process. We train all parameters end-to-end under the standard variational objective. Experiments sampling steps, our approach produces a higher quality samples than conventional discrete diffusion models using the same reverse parameterization.
various benchmarks show that, for given number of
INTRODUCTION
In the last years, diffusion models have demonstrated strong performance across many continuous (Hoogeboom et al. 2024) and discrete (Lou et al.) domains. Recent work has shown that distillation approaches and advanced training techniques allow learning a few-step (Salimans et al.) 2024), or sometimes
single-step, generative (Xu et al., (025) procedure in the continuous domain. However, discrete diffusion still requires a sigmificant number of consecutive generative steps (often comparable the data dimensionality). which makes inference computationally and time expensive.
The fundamental difference between continuous and discrete diffusion is the existence of a genera- tive ordinary differential equation (ODE) in the continuous case. Such ODE defines continuous trajectories that connect the noise and data distributions. Therefore, by learning the trajectories of an ODE, we can sample noise and directly map it to a data point. In discrete space, however, such continuous deterministic trajectories that would uniquely map noise to data do not exist, so these techniques are not applicable. At the same time, direct training of few-step conventional diffusion model fails due to a mismatch between the complicated target distribution and the limited param- eterization of the generative process. However, recent work on more flexible forward processes in diffusion models demonstrates that learning the noising process can improve generative performance in both continuous (Bartosh et al. 2024) and discrete Shi et al. settings.
Inspired by these results, propose Forward-Learned Discrete Diffusion (FLDD) framework for discrete diffusion models with a highly flexible parameterization of the forward process. We show how to parameterize the forward process that training iteration remains computationally efficient, while the way we destroy information at each individual coordinate depends on the entire datapoint and not just the time step, in conventional diffusion. Moreover, we show how train this model efficiently in an end-to-end manner by minimizing the variational bound on the model's likelihood. The flexibility of FLDD enables better alignment of the two processes with small number of steps, leading to improved generative performance.
Work done while interning Microsoft Research Cambridge.

---

Published as a conference paper at ICLR 2026
We summarize contributions follows:
We introduce Forward-Learner Discrete Diffusion (FLDD), generalizing discrete diffusion
models and providing efficient parameterization of the model.
We provide an efficient, end-to-end, simulation-free training procedure.
We demonstrate the robustness of our method on several discrete domains, including im- ages, molecules, and natural language text.
2 BACKGROUND
2.1 AUTOREGRESSIVE MODELS
We aim to model a distribution over discrete sequences ( where x (I....; ( K the number of discrete values, and D is the length of the sequence The standard approach for this is the autoregressive (AR) model, where each next token is predicted sequentially, conditioned on previous ones. Formally, the density of the AR model can be written as:
po(x) = III (x S5)
where x" represent coordinate number i of and < represents prefix of first coordinates.
An important drawback of the AR approach its slow sequential inference. To sample a sequence of D elements, must sample D times in sequence from the conditional distribution po xx **).
2.2 DIFFUSION MODELS
In contrast, diffusion models (DMs) provide fully parallel sampling procedure. A DM consists of main components: the forward and reverse processes. The forward, or noising, process is a Markov process that takes a data point and progressively corrupts it, producing sequence of latent variables Zu. ,
T
(Zo.rl x) = g(Zulxx xx) II (2z.22))
1=1
where = = 1.
The reverse process is also a Markov chain. starts from a sample zr and proceeds backward in time, inverting the corruptions applied by the forward process:
I
Po (Zo.T, x) p(x(20)p(ZT) II (2..21))
where p(x|zo) is usually simple, non non-parametric distribution.
The training objective of diffusion models is variational bound on the log- likeliiid the reverse process Ho et al. (2020):
(x()) < (u(())te,))) (DK.. (g(2.)2n (e,, x) x))lpo( 24
(4)
fuin
log p(x|2o))
+ DEL (g(er (x))p(()))
(5)
Cree
Sprior
where u(t) is uniform discrete distribution the time steps 1, 1,.... T. The three terms corre- spond to the reconstruction loss Lree. the diffusion term Laift, and the prior term Cprior.
Typically, the forward process is designed such that Zu and q(zz (x) zs p(zr), where p(zr) is a known prior. In practice, the reconstruction and prior terms are often negligible, and attention is usually focused on the diffusion term Laitt
2

---

Published as a conference paper at ICLR 2026
Although single step in DMs, following (z, (z)), may have computational cost comparable to sampling from px(x) in the AR approach, DMs allow this process to be parallelized, with all coordinates updated simultaneously each step. Thus, instead of D sequential steps in AR models, DMs require 7 sequential steps.
2.3 LIMITATIONS OF DIFFUSION MODELS
Despite this appealing formulation, in practice the number of steps T required for DMs can be comparable to sequence length D, undermining the speed advantage. To see why DMs perform poorly when T' is small, we need to examine what they actually learn.
It can be shown that, for given forward process, the target distribution for the reverse process is the marginalization of the posterior:
q(z,\z)) = f q Ze (z) ze, x)dx = (x(() [[(2)|2;,))
(2,
(6)
Thus, the forward process implicitly defines the target for the reverse process. If the reverse model pz(z,/Zz) is not expressive enough to match the target distribution (zz,/z,), it will fail to accurately approximate the data distribution g(x).
At the same time, efficient sampling from (z. lz)) is essential for DM performance. The most straightforward way enable fast, parallel sampling is parameterize the reverse process as fac- torized:
Po(z, (2) = II [o(zz where po(z)|zz)) Cat (2): v2(2z,)))
where z denotes the i-th coordinate of the discrete sequence 2s-
This factorization allows all elements to be sampled in parallel, but it also greatly limits the model's flexibility. However, in general, the target distribution q(z,, 2.) in Equation (6) is not factorized.
This difference is especially notable when the number of steps T small.
Importantly, the factorized parameterization in Equation 7 does not imply that each coordinate is treated independently. At step s, the distribution of cach z depends all coordinates of the previous state z; = [z!.... =!", not only on zi. Thus, even though sampling is performed through factorized distributions, the resulting distribution (x(x) can still be complex and non-factorieed.
3 FORWARD-LEARNED DISCRETE DIFFUSION
To construct a discrete DM with only a few generative steps, we need to reduce the gap between the
target distribution g(z, l2.) (Equation (6)) and the model approximation pe (z. (2,) (Equation ( is unclear how to make the reverse process more flexible without compromising efficiency. Instead, propose make the forward process more flexible. Then, the target distribution induced by the forward process (Equation (6)) can be adapted to the constraints of the generative process and enable few-step generation. To build intuition for the desired structure of the model, let us consider two motivational examples.
First, consider a mixture of two isotropic Gaussian distributions in D-dimensional space: g(x) wx;(x; #1, [) + (1 - w)((x; μ2, I). In the general case, we cannot sample from this distribution in a single step by drawing a sample from factorized distribution, since the coordinates are correlated. However, we can do it in two steps: first sample the index of the Gaussian, and then sample from the chosen Gaussian. In this case, both samples come from factorized conditional distributions.
A less trivial example is discrete random walk. Consider process that starts at x' 0 and takes steps (cither +1 or -1): x f+1 x + 1. This process generates sequences of elements and can be naturally modeled in D steps with autoregressive approach. Interestingly, can also be modeled in just two steps. First, sample D independent +1 values, which define the individual steps. Then, by taking prefix sums, we can construct the full sandom-walk sequence.
These examples demonstrate that, in contrast to conventional discrete diffusion with uniform injection or absorption (masking), some forward processes allow generation of complex data in
noise noise
3

---

Published as a conference paper at ICLR 2026
small number of steps. The question is: what should this process be? Unfortunately, in the gen- eral case, it is unclear how the forward process should look. This question is especially challenging because the structure of the forward process must depend on the underlying data distribution. There- fore, this work we propose learning it from data.
In this section, present Forward- Learned Discrete Diffusion (FLDD) discrete diffusion frame- work that enables end-to-end training of both the forward and reverse processes, which leads to better likelihood estimation and allows a significant reduction in the number of generative steps. At the same time, our approach does not change the reverse process and, as a result, the generative procedure, so introduces no inference overhead. FLDD also preserves important properties of conventional diffusion models, such as the variational objective and a simulation-fri optimization procedure.
3.1 LEARNABLE FORWARD PROCESS
Since we want to modify the target (Equation 66 ) we must change the forward process that induces Let us begin by generalizing the forward process.
The forward process is only required during training. From the objective in Equation we observe that the forward process must satisfy two conditions: efficient sampling from the marginals q(z, x) and tractable posteriors q( /zz x) for evaluating the KL divergence. With this in mind, we can reformulate the forward process from its Markovian definition in Equation to the following non- Markovian form:
(Zoo x)) = 9 (x) II qZz Ze, x).
(8)
This non-Markovian formulation enables more flexible and efficient parameterization of the for- ward process. Specifically, we suggest making the forward-process distributions qu (Zz(x) and 9, (z, Zz, x) learnable. the parameterization sufficiently flexible, we can expect find pa- rameters such that the resulting target distribution 2pp (2, (Equation (6)) is factorized.
To learn the parameters s,, we optimize the same variational objective as Equation ( In this setup, just as the reverse process adapts to the forward process, the forward process also adapts to the reverse one. Since the generative process is factorized by design (Equation ( the forward process is encouraged to produce a factorized target distribution that the reverse process can effectively match.
3.2 PARAMETERIZATION OF THE FORWARD PROCESS
Forward Marginals. For the forward marginal distributions 9, (2, propose using the same form as the generative model (Equation (7)) and parameterizing them as factorized distributions:
L (Zx (x) = II (z)lx, where (x = Cat Czal22/ ux (x)
This parameterization allows us to sample 2, given x efficiently during training. However, unlike in conventional discrete diffusion models, each may come from complex distribution that depends on the entire data sample x = x' "'....,"" x", not just the i-th element x!.
In addition, must enforee the boundary conditions q, (zolx) 8(20 -1)o-x) and qu (z(zz(x) (x) p(zr). Both properties can be ensured through an appropriate reparameterization u,
Forward Posteriors. In addition to being tractable, the posterior distribution (p Lz, x) must be consistent with the marginals q,p (z(()))
(z. = qu do x)q (Za|Ze. bafexx)) x)dzr.
(10)
other words, the posteriors qu (z, Z₁, x) must define valid transport plan for the probability mass across consecutive time steps. We propose using a simple non-parametrie technique, Maximum Coupling, for posterior parameterization.
4

---

Published as a conference paper at ICLR 2026
Algorithm Training procedure with REINFORCE method.
Require: Dataset, T time-steps, U. (x, t) forward parameterization, vo(zi, t) reverse param-
eterization
for learning iterations do
x ~ Dataset, ~ u(I), 1-1
> Sampling data point x and two consecutive time steps
u, = ux s), u, (x,,t)
Parameters of qu 9. |x) and qu (Zr(x) (Equation (9
Ze Cat (zri up)
Sampling 2, from categorical distribution with parameters u
Construct Uglt by ug. u, and x
Calculating parameters of 9,p (z, |z,, (Section 3.2
Va/t (z, log Ti
Parameters of the reverse process pe (zz) (Equation
L = Ei= 1
D Objective function is KL divergence (Bquation (5)
Vi
Gradient step on and w.r.t.
(me.ze)
REINFORCE functional (Section 3.3
[
end for
Suppose we want to construct transport plan between two 1-dimensional categorical distributions with parameters Us u K- Intuitively, when moving probability mass from u, to us. Max- imum Coupling attempts Z = k. if u 2 u, we keep 2, = Zz. Otherwise, if u < u, we redistribute the excess mass u u from bin & to bins with deficit.
minimize the amount of mass that must be moved. Specifically, for
Formally, we can write it as:
where u min(u u
21 k j,
(Za.22)) = Cat (2s; Unje),
min(0.u.
m 2/00
2t = k #
min(0, u, - -u,) min (0, u, uz)l
and
maje
(11)
Here, maje represents the parameters of a categorical distribution corresponding to the deficit of probability mass at timestep compared to timestep t. We redistribute the extra probability mass from u, to u, according to maje- Importantly, the computation of the posterior parameters requires only simple vector operations, it can be performed efficiently.
To construct the full D-dimensional distribution (1.22, Ize x), we apply the Maximum Coupling procedure independently each coordinate, using the parameters u, up(x,s) and u, up (x, t) (Equation (
We emphasize that, with this parameterization, the distributions of forward trajectories for each co- ordinate are conditionally independent given the data point However, in contrast to conventional discrete diffusion, the trajectory of each coordinate z; may have a complex, non-linear dependence on the entire datapoint not just on x'. Moreover, marginal distributions can still be complex, and the unconditional distribution of forward trajectories remains expressive.
The parameterization of the forward process described here is not unique. For instance, we may trade off sampling efficiency of the marginals for additional flexibility by making them partially autoregressive. For the posterior, chosen metric or introduce dependencies between coordinates or additional learnable parameters. We believe there may be better ways to parameterize the forward process, but we leave this question for future research.
could use element- -iiee optimal transport with respect to
3.3 OPTIMIZATION
As mentioned in Section train the model by minimizing the same variational objective in conventional diffusion (Equation (5)) with respect to the parameters 0 and sp While the objective itself is fully tractable, its gradients with respect &o difficult compute. Let take a closer look at the gradients of the diffusion term:
VeLain VPBu(t)as (z)o
(12)
5

---

Published as a conference paper at ICLR 2026
Algorithm 2 Sampling procedure.
Require: T- -imme-stess, vo(zi, reverse parameterisation
2T Cat Car: Vr) for t T.... do
Sampling 2r from prior distribution p(2r)
sa (z, 2g Cal (25 Vole)
Parameters of the reverse process pz (zz) (Equation D. Sampling 2, from categorical distribution with parameters Vaje
end for
The expression inside the expectation fully differentiable. However, the distribution lG (z, /x) and the samples Ze are discrete, which prevents us from using the reparameterization trick to efficiently estimate gradients.
Unbiased Optimization. The standard approach for computing gradients of this form is the REIN- FORCE method Williams (1992). We can rewrite the gradients of the diffusion loss as follows:
90 (z(x)( DKLL (go /2, x))lpo (2.)) [g x)
V V Caif (2.1))
V
(13)
where L-1sg sg is stop-grad operation.
In this way, we can use the Monte Carlo method to build an unbiased gradient estimator and train the model end-to-end. Importantly, in conventional diffusion, we optimize variational bound the model's likelihood. We summarize the training algorithm in Algorithm [
Relaxed Warm-Up. Unfortunately, REINFORCE known to produce high-variance gradients, so training a model from scratch with REINFORCE alone often leads to instability, To obtain better initialization, we adopt different strategy.
We begin training with a continuous relaxation of the categorical distribution (z, using the Concrete distribution Jang et al. (2016); Maddison et al. (2016). Specifically, instead of drawing hard samples z, according to Equation
we use the same parameters (x, t) to define Concrete distribution ( (7, x) (2 x) with Temperature
construct a posterior for a relaxed sample 2,, we combine the posteriors for discrete samples 2, according to the components of 2,
Tr.p (z.(22
Σ
z, 9 (z.12)) x,
(14)
where z, denotes the component corresponding to discrete value of coordinate in the relaxed sample 2.. Thanks to the simple structure of the posterior distribution (Section 3.2). this weighted average does not add computational overhead.
Relaxed samples allow with a temperature т = 1, which we exponentially decrease to r = 10 over 10* to 102 opti- mization steps. After this warm-up phase, switch to training the model with the REINFORCE approach described above.
the reparameterization trick for gradient estimation. Training starts
3.4 GENERATIVE PROCESS
We use the standard parameterization of the generative process as discussed in Section 2.2 (Equa- tion (7)). Importantly, the sampling procedure does not require the forward process, do not add any computational overhead for inference. We summarize the generative procedure Algorithm
It is important to note that another common parameterization, where the model samples an approx- imate data point & and then resamples an intermediate step (z, (2 x). is not applicable in our formulation. Although it may be possible make 9p (z,(z)) (Equation (6) factorized, the distri- bution 9,- (x|z,) will generally not be factorizable. As a result, the reverse process cannot learn accurately.

---

Published as a conference paper at ICLR 2026
Step
Step
Step 2
Figure I: Training data distribution and learned dynamies for a mixture of two Gaussians.
Step
Step
Step
Step
3
A
3
8
2
4
2
2
3
8
8
2
(a) Binarized MNIST generation with Masked Diffusion Model in steps. Gray color represents masks.
Step
Step
Step
Step
Step
8
7
6
8
0
4
5
2 9 3
0
0
9
6 0
4
g
.00
3
3
2
9
1
2
(b) Binarized MNIST generation with learned forward process in 4 steps.
Figure 2: Learned dynamics for Binarized MNIST dataset. Generative process starts from prior distribution p(2r) at time steps t = T and goes backwards in time.
EXPERIMENTS
In this section, we provide experiments both synthetic and real-world data. We present FLDD as a general framework for reducing the number sampling steps. Therefore, the goal this section to demonstrate that FLDD indeed improves sampling efficiency across different settings. Our objective is not to outperform all other approaches every domain. Instead, FLDD is compatible with existing extensions. We also assume there is potential room for optimizations and better design choices in these domains.
In all our experiments, eterize both the generative process Pa (z(2,lz)) and the forward marginals g,- (x) We emphasize that, while this doubles the number of learnable parameters, it does affect the capacity of the generative process. Additional details on parameterization, training, and evaluation are provided in Section A
the neural network architectures and hyperparameters param-
4.1 ToY DATA
First, we consider a toy setting with x (1,...,50)" We train a diffusion model to generate samples in just two steps. In Figure [ show results for mixture of two Gaussian distributions. As we discussed in Section 3 each individual Gaussian can be generated in a single step, since the coordinates independent. However, generating the mixture requires at least two steps. In this case, the model learns first to generate an intermediate factorized structure and then to produce the final data.

---

Published as a conference paper at ICLR 2026
Table I: Results on ROCStories dataset.
MAUVE
PPL
Div t
GPT2
Radford al. 2019)
0.789
20.5
0.252
GPT Neo (Black
202
0.720
19.9
0.258
AR- AR-Diffusion (Wu el al. 2023
0.066
0.101
DiffiuSeq (Gong ei al. 2022
0.086
50.5
0.124
SeqDiffuSeg Yuan al. 2022
0.103
29.3
0.137
TESS (Mahabadı
2023
0.061
22.4
0.163
SEDD Lou et al. [2233
0.598
70.8
0.336
LD4LG Lovelace ct al.
2023)
0.716
30.6
0,331
COSMOS Meshchammon et
2025
0,940
26.3
0.346
FLDD, 100
0.538
55.2
0,280
FLDD. T 10
0,511
60.5
0.285
Table 2: Molecular generation results.
QM99
ZINC250K
Valid
Unique f
FCD 4
Valid f
Unique
FCD 4
MoFlow (Zang & Wang 2020
91.36
98.65
4.467
63.11
99.99
20.931
EDP-GNN NNu et al. 2020)
47.52
99.25
2.680
82.97
99.79
16.737
GraphEBM LIy et al.
202T
8.22
97.90
6.143
5.29
98.79
35,471
GDSS Jo 2022
95.72
98.46
2.900
97.01
99.64
14.656
Digress Vignac al.
2022)
99.00
96.20
Flow Matching (1 Jppmann et al.
2023)
94.10
98.20
5.155
94.01
96.68
18.764
Dirichlet FM (Stark et al. 2024)
99.10
98.15
0,888
97.52
99.20
14.222
CadFlow (Bijkelboom
2024
99.81
99.95
0.441
99.21
100.00
13.211
LO-ARM (Wang
2025
99.85
98.85
0.240
96.70
100.00
3.229
FLDD, T
100
99.67
98.93
0.328
97.79
100.00
8.487
FLDD, T =
99.08
99.95
0,385
96.77
100.00
10.414
4.2 REAL WORLD DATA
We evaluate FLDD on the text dataset ROCStories (Mostafazadeh et al. 2016) and two molecular generation benchmarks: QM9 (Ramakrishnan al. 2014) and ZINC250K (Irwin et al., 2012). The results are summarized in Table T For each problem, we evaluate FLDD in two settings: first, with 100. which is comparable what other models use; and second, with a significantly reduced T = 10 steps.
In all cases, FLDD demonstrates results comparable other baselines with 100. Some base- lines may outperform FLDD, but we attribute this to suboptimal parameterization and hyperparam- eters that we do tune. As we discuss in Section 3.4 important for FLDD parameterize the distribution 4 (2, 2,) directly, which is known to be suboptimal in conventional diffusion. We believe there are ways to reparameterize properties. However, we leave this for future research.
generative process efficiently while preserving the same
At the same time, FLDD shows only a slight drop in performance with T 10 steps, whereas other diffusion models do not generate realistic-looking data with such small number of steps. This demonstrates that FLDD allows significantly better trade-off between sample quality and inference speed.
4.3
LEARNABLE MASK DIFFUSION MODELS
A popular variation of discrete diffusion is mask discrete models (MDM), where the forward process randomly replaces tokens from the sequence x with masks, and the reverse process replaces mask tokens with tokens from the vocabulary. For MDM, it is clear that if there is correlation between coordinates and the number of generative steps T is significantly smaller than the sequence length D. the generative model will fail, since it will have unmask tokens independently without taking
8

---

Published as a conference paper at ICLR 2026
correlations into account. In the general case, this problem is unavoidable in conventional MDM. However, with FLDD, we can improve performance by learning to unmask at each step only tokens that are less correlated (conditionally
current observations).
To demonstrate how this works, we restrict the forward process to masking. Specifically, qp x) defines only the probability of masking each token conditionally the full data point Then, we expect the model to learn masking schedules such that the generative process corresponds to unmasking less correlated tokens.
We conduct an experiment on the MNIST dataset Salakhutdinov & Hinton 2009). Figure shows images generated by a conventional MDM and by FLDD with learnable masking schedules. Since an MDM learns to unmask pixels uniformly, it struggles to generate natural natural-looking images. In contrast, our model learns more convenient intermediate distributions and produces more realistic samples. We emphasize that the generative processes in both models share the same parameterization.
5 RELATED WORK
Diffusion models Ho et al. (2020) have recently been adapted to discrete spaces for language, graphs. and molecules. Text diffusion includes sequence-to-sequuenee formulations Gong et al. (2022); Yuan
al. (2022), hybrid aatoregressively-ccndiiiooo condiiiond variants |Wu al. (2023), simplex- or ratio-
processes with improved schedules and theory Shi et al.) For molecular and graph data, flow-/score-bases and diffusion-style approaches continue to advance sampling quality and validity Lipman et al. (2023);
based objectives Mahabadi et al. ) Lou et al. (2023), and masked/absorbing--tee
Stark et al. (2024): [Eijkelboom et al.] (2024); Jo et al. (2022); Vignac et al. (2022); Wang et al. (2025). Reducing the number of reverse steps has been pursued via distillation and consistency-style accelerations in the continuous domain Salimans et al. (2024); Xu et al. (2025), but discrete models often retain long sampling chains due to factorized reverse parameterizations. Our work targets this limitation by learning forward process that induces a reverse target distribution compatible with the same efficient, factorized generative family, thereby enabling few-step sampling without changing the sampler itself.
Most diffusion methods fix the forward (noising) dynamics, which implicitly defines the target that the reverse model must match. Recent work in the continuous setting shows that learning the for- ward process can tighten likelihood bounds and improve generation Bartosh et al. (2024); Shi al. We build on this thread but in the discrete regime: we introduce a non-Markovian yet tractable parameterization of the forward marginals and posteriors, and train end-to-end together with the standard factorized reverse model. This design preserves the usual variational training objective and the parallel, stepwise sampler, while adapting the training target to better match the model's inductive biases.
CONCLUSION
We presented Forward-Learned Discrete Diffusion (FLDD), framework that learns the forward (noising) process-viz a non-Markovian factorization of forward marginals and tractable posteri- ors- while keeping the standard factorized reverse sampler unchanged and preserving the usual variational training objective. This alignment between learned target and the fixed reverse family enables few-step generation without inference overhead.
Empirically, FL.DD matches strong baselines when 7=100 and maintains competitive quality even 7=10 across text (ROCStories) and molecule benchmarks (QM9, ZINC250K), and improves masked-diffusion generation images by learning data-aware masking schedules-ap with the same reverse parameterization. This yields a better quality-latency trade-off than conventional dis- crete diffusion.
Limitations and future work include exploring richer forward parameterizations, reducing reliance on REINFORCE via lower-variance estimators. We also note the extra parameters of the forward network (though not affecting reverse capacity) and leave broader design optimizations to future research.

---

Published as a conference paper at ICLR 2026
REFERENCES
Grigory Bartosh, Dmitry P Vetrov, and Christian Andersson Naesseth, Neural flow diffusion models: Learnable forward process for improved diffusion modelling. Advances Neural Information Processing Systems, 37:73952-73985, 2024.
Sid Black, Leo Gao, Phil Wang, Connor Leahy, and Stella Biderman. Gpt-neo: Large scale autoregressive language modeling with mesh-tensorflow, march 2021. URL https://do.. org/10.5281/zenon 5297715(5):3, 2021.
Floor Eijkelboom, Grigory Bartosh, Christian Andersson Naesseth, Max Welling, and Jan Willem van de Meent. Variational flow matching for graph generation. Advances in Neural Information Processing Systems, 37:1 1735-11764, 2024.
Shansan Gong, Mukai Li, Jiangtao Feng, Zhiyong Wu, and LingPeng Kong. Diffuseq: Sequence to sequence text generation with diffusion models. arXiv preprint arXiv:2210.08933, 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-688 0851, 2020.
Emiel Hoogeboom, Thomas Mensink, Jonathan Heek, Kay Lamerigts, Ruiqi Gao, and Tim Sali- mans. Simpler diffusion (sid2): 1.5 fid on imagenet512 with pixel-space diffusion. arXiv preprint arXiv:2410.19322 2024.
John J Irwin, Teague Sterling, Michael Mysinger, Erin Bolstad, and Ryan G Coleman. Zinc: free tool to discover chemistry for biology. Journal of chemical information and modeling, 52 (7):1757-1768, 2012.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizzatio with gumbel-softmax, arXiv preprint arXiv: 1.0/ 144, 2016.
Jachyeong Jo, Seul Lee, and Sung Ju Hwang. Score-based Generative Modeling of Graphs via the System of Stochastic Differential Equations. In Proceedings of the 39th International Conference on Machine Learning. pp. 10362-10383. PMLR, June 2022.
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Repre- sentations, 2023. URL
Meng Liu, Keqiang Yan, Bora Oztekin, and Shuiwang Ji. GraphEBM: Molecular graph generation with energy models. In Energy Based Models Workshop - ICLR 2021, 2021. URL https:
/ //openreview. wwt. fttt://formmmi:cc51 dc 1_
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution, 2024. URL htlps://uxxi. org/abs/2310. 16834.
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. arXiv preprint arXiv:2310. 16834, 2023.
Justin Lovelace, Varsha Kishore, Chao Wan, Eliot Shekhtman, and Kilian Weinberger. Latent dif- fusion for language generation. Advances in Neural Information Processing Systems, 36:56998- 57025, 2023.
Youzhi Luo, Keqiang Yan, and Shuiwang Ji. Graphdf: A discrete flow model for molecular graph generation. International conference on machine learning, pp. 7192-7203. PMLR, 2021.
Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. arXiv preprint arXiv:1611.00712 2016.
Rabeeh Karimi Mahabadi, Hamish Ivison, Jaesung Tae, James Henderson, Iz Beltagy. Matthew E Peters, and Arman Cohan. Tess: Text-to-text self-conditioned simplex diffusion. arXiv preprint arXiv:2305.08377 083799 2023.
10

---

Published as a conference paper at ICLR 2026
Viacheslav Meshchaninov, Egor Chimbulatov, Alexander Shabalin, Aleksandr Abramov, and Dmitry Vetrov. Compressed and smooth latent space for text diffusion modeling. arXiv preprint arXiv:2506.21170, 2025.
Nasrin Mostafazadeh, Nathanael Chambers, Xiaodong He, Devi Parikh, Dhruv Batra, Lucy Vander- wende, Pushmeet Kohli, and James Allen. A corpus and cloze evaluation for deeper understanding of commonsense stories. In Proceedings the 2016 Conference of the North American Chupter of the Association for Computational Linguistics: Human Language Technologies, pp. 839-849, 2016.
Chenhao Niu, Yang Song, Jiaming Song, Shengjia Zhao, Aditya Grover, and Stefano Ermon. Per- mutation invariant graph generation via score-based generative modeling. In International Con- ference on Artificial Intelligence and Statistics, pp. 4474-4484. PMLR, 2020.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp. and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
Ruslan Salakhutdinov and Geoffrey Hinton. Deep boltzmann machines. In Proceedings of the twelfth international conference on artificial intelligence and statistics, pp. 448-455, PMLR, 2009.
Tim Salimans, Thomas Mensink, Jonathan Heek, and Emiel Hoogeboom. Multistep distillation of diffusion models via moment matching, Advances in Neural Information Processing Systems, 37: 16046-36070 2024.
Chence Shi, Minkai Xu, Zhaocheng Zhu, Weinan Zhang, Ming Zhang, and Jian Tang. Graphaf: flow-based autoregressive model for molecular graph generation. In International Conference on Learning Representations, 2019.
Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, and Michalis K Titsias. Simplified and gener- alized masked diffusion for discrete data, 2024. URL https://arxiv. org/abs/2406.043329 (Cited on page 10).
Hannes Stark. Bowen Jing, Chenyu Wang, Gabriele Corso, Bonnie Berger, Regina Barzilay, and Tommi Jaakkola. Dirichlet flow matching with applications to DNA sequence design. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 4 Ist International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 46495-465111 PMLR, 21-27 Jul 2024. URL https:tpp//cmmm httppss//ppoccee.nnggeee mlr prwss v235 htw.ees///2///tttt22bbbooo
Hugo Touvron, Louis Martin, Kevin A. Stone, Peter Albert, Amjad Almahairi, al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023. URL
https:://artae org/ abs 220//2377..0 2200//9277308 09288
Clement Vignac, Igor Krawezuk, Antoine Siraudin, Bohan Wang, Volkan Cevher, and Pascal Frossard. DiGress: Discrete Denoising diffusion for graph generation. The Eleventh Inter-
national Conference Learning Representations, September 2022.
Zhe Wang, Jiaxin Shi, Nicolas Heess, Arthur Gretton, and Michalis K Titsias. Learning- order autoregressive models with application to molecular graph generation. arXiv preprint arXiv:2503.05979, 2025.
Ronald J. Williams. Simple statistical
algorithms for connectionnsttreifforreeee. learning. Machine Learning. 19229-256, 1992. doi: 10.1007/BF00992666
Tong WU, Zhihao Fan, Xiao Liu, Hai-Tac Zheng, Yeyun Gong, Jian Jiao, Juntao Li, Jian Guo, Nan Duan, Weizhu Chen, al. Ar-diffusion: Auto-regressive diffusion model for text generation. Advances in Neural Information Processing Systems, 36:39957-39974, 2023.
Yilun Xu, Weili Nie, and Arash Vahdat. One-step diffusion models with f-divergence distribution matching. arXiv preprint arXiv:2502. 2025.

---

Published as a conference paper at ICLR 2026
Hongyi Yuan, Zheng Yuan, Chuanqi Tan, Fei Huang, and Songfang Huang. Seqdiffuseq= Text diffusion with secoder-decoder transformers. arXiv preprint arXiv:2212.10325, 2022.
Chengxi Zang and Fei Wang. Moflow: an invertible flow model for generating molecular graphs. In Proceedings of the 26th ACM SIGKDD international conference on knowledge discovery & data mining, pp. 617-626, 2020.
A IMPLEMENTATION DETAILS
In all our experiments, the generative process po po(2//zz) and the forward marginals de (x 24). First, we warm up the model for 10 steps as discussed in Section 3.3 then optimize with REINFORCE for 100k iterations. We use the AdamW optimizer with learning rate of 2 10-14
identical neural network architectures and hyperparameters for both
For evaluation on the molecular datasets QM9 (Ramakrishnan et al. 2014) and ZINC250K (Irwin al. 2012), we follow the standard setup Shi al. (2019); Luo (2021): Vignac al. (2022);
Jo al. (022). We use the same hyperparameters for model parameterization as in Eijkelboom et al. (2024).
For the experiment on the ROCStories (Mostatazadeh et al. 2016) dataset, we use LLaMA 2 Tou- vron al.aal22222 (2023) with layers and heads and an embedding size of 512 parameterize the model, and pretrained GPT-2 for PPL calculation.
12


Reviews:
[{"review_id": "Aa7Jo0irqg", "reviewer": "Reviewer_mvjP", "summary": "This paper introduces Forward-Learned Discrete Diffusion, a method designed to accelerate the reverse process while supporting end-to-end training.\n\nUnlike previous works, FLDD adopts a non-Markovian formulation with learnable marginal and posterior distributions, which allows the generative process to remain factorized while matching the target distribution defined by the noising process. The proposed method enables a better alignment between the forward and reverse processes even with a small number of steps, thereby improving its generation performances.", "questions": "1. In the forward process, what are the key differences between FLDD and Star-shaped Diffusion? Could the authors elaborate on the unique design choices of FLDD in this regard?\n\n2. Compared with continuous diffusion models that adopt similar training strategies on continuous datasets (e.g., MNIST), what specific advantages does FLDD offer? For example, in terms of generation quality, training efficiency, or adaptability to data characteristics.\n\n3. In the experimental results (Table 1), FLDD performs poorly on real-world data, particularly with a PPL of 60, which indicates an inability to generate semantically meaningful sentences. Could the authors provide a more specific analysis of the underlying reasons? For instance, whether it is due to mismatches between the forward process design and real-world data distributions or limitations in the learnable marginal/posterior functions.\n\n4. While FLDD demonstrates strong performance in few-step diffusion, what are the precise time and computational costs compared to existing methods? A smaller number of steps does not always translate to lower overhead, so quantitative comparisons would help verify its efficiency advantages.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- A novel method that integrates insights from SS-DDPM and DPS, achieving excellent performance in few-step generation under specific scenarios.\n\n- The proposed method significantly reduces the requirement for the number of timesteps in the backward process, which is a notable advantage for practical applications.\n\n- This paper is well-structured and logically coherent, with smooth transitions between sections. The examples are appropriately designed, making the core ideas easy to follow.", "weaknesses": "- The definition of \"end-to-end\" in the context of FLDD remains unclear. It is necessary to clarify whether this refers to training without separate stages or other specific characteristics.\n\n- In Section 4.2, FLDD exhibits poor performance on the ROCStories dataset, but the authors' explanation for this phenomenon lacks sufficient persuasiveness and requires further elaboration.", "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": "vguzWoIDkm", "reviewer": "Reviewer_ATG6", "summary": "This paper introduces Forward-Learned Discrete Diffusion (FLDD), a new framework for discrete diffusion models aimed at improving sampling efficiency and enabling few-step generation. The core idea is to replace the standard fixed, Markovian noising process with a learnable, non-Markovian forward process. By co-training this flexible forward process alongside the reverse (denoising) process, the authors argue that the model can learn an easier-to-invert corruption path. This, in turn, allows a simple, factorized reverse sampler to generate high-quality samples in a significantly reduced number of steps. The method is trained end-to-end using the standard variational objective and is evaluated on synthetic data, binarized MNIST, text generation (ROCStories), and molecular generation (QM9, ZINC250k).", "questions": "N/A", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The paper does a good job of grounding the proposed method in the established variational inference framework for diffusion models. The formulation for end-to-end training of both the forward and reverse processes while preserving the standard variational objective is principled and clearly presented.\n2. The paper is generally well-written and easy to follow.", "weaknesses": "1. The primary qualitative results are presented on a 2D Gaussian mixture and binarized MNIST. These are considered solved or overly simplistic problems in the current deep generative modeling landscape. Demonstrating success on these toy tasks provides very little signal about the method's effectiveness on complex, high-dimensional discrete data that is of interest to the community.\n2. A significant weakness of this paper is its failure to demonstrate a clear performance advantage over existing, strong baselines. The paper's own results on the ROCStories dataset (Table 1) show that SEDD (Lou et al., 2023), a recent and relevant discrete diffusion model, achieves a higher MAUVE score (0.598 vs. 0.538) and a better (lower) PPL than the proposed FLDD, even when FLDD uses a large number of steps (T=100). This directly undermines the paper's central claims regarding improved sampling efficiency and achieving a better quality-latency trade-off. Given that a method from 2023 already shows superior performance on the same benchmark, the practical contribution of the proposed, more complex FLDD framework is called into question.\n3. The paper fails to cite, discuss, or compare against the seminal work: LLaDA[1], a landmark paper that demonstrated how to successfully apply diffusion models to high-quality text generation and is arguably the most famous work in this specific area.\n\n[1] Nie, Shen, 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.09992 (2025).", "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": "a1HN7KIqY2", "reviewer": "Reviewer_vLZq", "summary": "This paper proposes forward-learned discrete diffusion (FLDD). It’s a discrete diffusion with a learnable forward process. The idea is to keep the generative process factorized while matching the target, so it can work well for small numbers of diffusion steps. Experiments show that the performance of FLDD drops a little when the number of diffusion steps goes down from 100 to 10.", "questions": "- What's the batch size used in training and how does it affect the performance?\n- What are the additional costs of training, in terms of memory (chips) and duration (time)?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Novel approach for training discrete diffusion models", "weaknesses": "- FLDD with T=100 doesn’t outperform some previous work. The authors argue that FLDD performs better for small Ts. However, there are methods to accelerate sampling of discrete diffusion models, e.g. distillation. For small-T experiments, FLDD needs to be compared with methods focusing on accelerated sampling", "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": "DNOH61I8ww", "reviewer": "Reviewer_wcR4", "summary": "The paper learns the forward noising process in a non-Markovian way so that the training targets match what a factorized reverse model can learn well in just a few steps. It keeps the reverse sampler unchanged. Via learned forward, the paper achieves much better quality–inference latency trade-offs. Across text (ROCStories) and molecular (QM9, ZINC250k) data, the paper with only $T=10$ reverse steps delivers quality close to its $T=100$ counterpart—unlike conventional discrete diffusion.", "questions": "The paper does not compare REINFORCE against other standard gradient estimators for discrete latents. It does not discuss control variates such as Rao–Blackwellization tricks. These would reveal which knobs actually make training computationally stable/cheap. Could you report some of these results?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 4, "strengths": "FLDD’s main idea is both theoretically sound and practically useful: learn the forward noising process in a manner so that the task for the reverse process becomes easier. In addition, the algorithm uses the standard variational objective and does not change the reverse sampler. The posterior is implemented with a cheap Maximum Coupling transport between consecutive marginals, which keeps training tractable while letting each coordinate depend on the whole input. Empirically, the paper shows few-step generation benefits illustrating a better quality–latency trade-off than conventional discrete diffusion.", "weaknesses": "FLDD adds a learned forward network, increasing training compute relative to standard discrete diffusion. The paper does not report quality vs total training and inference budget (e.g., GPU-hours or FLOPs), leaving unclear whether the few-step gains persist under fixed total compute budget.", "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)