Informed Correctors for Discrete Diffusion Models
Authors: Yixiu Zhao
Jiaxin Shi
Feng Chen
Shaul Druckmann
Lester Mackey
Scott Linderman
Venue: NeurIPS 2025 poster ()
Keywords: Discrete diffusion
diffusion models
generative models
transformers
predictor corrector sampling
Gibbs sampling
Primary area: deep_learning
TL;DR: We propose an informed corrector for masked discrete diffusion that reduces approximation errors, enabling faster sampling and better sample quality in both synthetic and large-scale settings.

Abstract: Discrete diffusion has emerged as a powerful framework for generative modeling in discrete domains, yet efficiently sampling from these models remains challenging. Existing sampling strategies often struggle to balance computation and sample quality when the number of sampling steps is reduced, even when the model has learned the data distribution well. To address these limitations, we propose a predictor-corrector sampling scheme where the corrector is informed by the diffusion model to more reliably counter the accumulating approximation errors. To further enhance the effectiveness of our informed corrector, we introduce complementary architectural modifications based on hollow transformers and a simple tailored training objective that leverages more training signal. We use a synthetic example to illustrate the failure modes of existing samplers and show how informed correctors alleviate these problems. On the Text8 dataset, the informed corrector improves sample quality by generating text with significantly fewer errors than the baselines. On tokenized ImageNet 256x256, this approach consistently produces superior samples with fewer steps, achieving improved FID scores for discrete diffusion models. These results underscore the potential of informed correctors for fast and high-fidelity generation using discrete diffusion.

Informed Correctors for Discrete Diffusion Models
Yixiu Zhao
Jiaxin Shi Google DeepMind ishi jiaxiiigmaii.ccom com
Stanford University
yi rixiuzestanf ord. edu
Feng Chen
Shaul Druckmann Stanford University shauld@stanford. edu
Stanford University Fengc@stanforr ord. edu
Lester Mackey Microsoft Research New England 1mackey@microsoft.
Scott Linderman Stanford University scott Soott.LLnderran@@aatffoo.ecc
Abstract
Discrete diffusion has emerged as powerful framework for generative modeling discrete domains, yet efficiently sampling from these models remains challenging. Existing sampling strategies often struggle to balance computation and sample quality when the number of sampling steps is reduced. even when the model has learned the data distribution well. To address these limitations, we propose a predictor-correcto sampling scheme where the corrector is informed by the diffusion model to more reliably counter the accumulating approximation errors. To further enhance the effectiveness of our informed corrector, we introduce complementary architectural modifications based on hollow transformers and a simple tailored training objective that leverages more training signal. We use synthetic example to illustrate failure modes of existing samplers and show how informed correctors alleviate these problems. On the text8 and tokenized ImageNet 256 x 256 datasets, our informed corrector consistently produces superior samples with fewer eirors improved FID scores for discrete diffusion models. These results underscore the potential informed correctors for fast and high-fidelity generation using diserete diffusion. Our code is available https: h/pp::///t/a thub. com/ com/linde http//nndrrmalaa.... inf wwweooomddooooo grmeddcorrectors.
Introduction
Denoising diffusion models are powerful generative models for high-dimensiona data [1-3]. The central idea of diffusion models to gradually corrupt data into noise using "noising" or "forward" process and then to train parameterized model (usually a deep neural network) to learn its time reversal, commonly known as the "denoising" or "backward" process. In continuous domains like image generation, the forward process is usually defined by gradual injection Gaussian noise and scaling that transform the data distribution into standard normal. The backward process is then approximated by learning the gradient of the log density (also known as the "score function") of the marginal distribution. To draw samples from trained model, one first generates Gaussian noise and then simulates the backward process using the score information. Diffuusio generative models are currently the dominant framework for image generation, and they can generate high- regh-resolution images with stunning details given prompts from a user [4, 5].
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Given the success of diffusion models for continuous data, recent work has explored diffusion modeling in discrete domains [6-12]. Discrete diffusion models can be applied language [11]. protein sequences [13, 14), graphs [15], and more. Notably, Campbell et al. [8] developed general framework for discrete denoising diffusion in continuous time. They formulated the forward and backward processes as continuous time Markov chains (CTMCs), and they learned to predict the distribution over denoised data via an evidence lower bound objective. Concurrent work by Shi et al. [12], Sahoo et al. [16] and Ou et al. [17] focused on discrete diffusion with the absorbing state forward process (a.k.a., masked discrete diffusion), deriving simple formulation of the ELBO that reduces to weighted cross-entropy. Shi et al. [12] also noted that this simple objective provides unifying perspective to previous works [8, 11, 18]. all of which are parameterizations and modifications the same ELBO objective.
Despite conceptual advances in discrete diffusion, the efficiency-accuracyy trade-off in simulating the continuous-time backward process still limit their effectiveness. Song et al. [19] proposed using approximate MCMC correctors to fix the discretization error in simulating the backward process. Campbell et al. [8] extended the predictor-correcto sampling to discrete diffusion models and discovered that the corrector steps significantly impact generative performance. In practice however, clear understanding of how corrector steps contribute to sample quality is still missing, and recent works [11, 12, 17, 16] still use predictor-only samplers for generation.
We show that predictor- reedictor-corrcctor schemes can be improved by leveraging different sampling methods. Analyzing the
corrector used by Campbell al. [8], we identify its failure mode in masked diffusion and illustrate it with simple Markov chain modeling task. To address this shortcoming, we propose the informed corrector - an alternative corrector scheme inspired by adaptive Gibbs sampling that fixes the issues of the forward-backward corrector and explicitly targets low-probability dimensions. Sampling with the informed corrector drastically decreases error rate in synthetic example, as well as in Text8, where the model makes much fewer spelling errors than baselines. On tokenized ImageNet 256 x 256, we show that the informed corrector yield diverse, high quality samples competitive with existing diffusion models.
2 Background
2.1 Continuous-Time Discrete Diffusion Models
Consider a data vector 2g E S sampled from data distribution Pdata where the base space for each component. Denoising diffusson models [1-3, 19] model this unknown data distribution by progressively adding noise to data samples until their distribution converges to a simple target, such as uniform distribution. Then, a deep neural network is trained to invert this process and remove the noise. This reverse process forms a generative model from which new samples can be drawn. The noising and generative processes are also called the forward and backward processes, respectively. We will adopt this convention for the rest of the paper.
Continuous-time diffusion models define a forward process with noising distribution qe (= and marginals qe (x;)) f qe (zre2o))o0 (0(())) (ro)) with qo Paata and limiting distribution qr ss æ. The forward process is designed that is simple. Then, parameterized process with marginals pf is learned to match the true backward process.
When the data is discrete (e.g., text, protein sequences, neural spike counts), the base space becomes a "vocabulary"" S (1....,S) S} of size S, and the forward process can be described as a CTMC with
initial distribution go and transition rate matrix R; s x The off-diagonal terms R2 (2: y) of the matrix represent the rate of state & transitioning to state at time t, and the diagonals Ri (x, x) are defined the negative of all other elements in the same row: Re x) = = -E Re (z. y). Assuming that R., Re commute for all and I, the finite-time transition probabilities satisfy
A+
}
444q0ly 1))
exp
R. ds
(x,v) [(x,v) +
o(Al).

---

where I represents the identity matrix, and in the square brackets we compute matrix exponential. The time reversal of this CTMC also a CTMC, and its transition rate is,
Re(y,x) x) Re(x, y) (x) q: (:(y)
Q0jo( zo)
Rí R1 v)
Σ
2: (y mo) 20)o y)-
(I)
Campbell et al. [8] approximated the backward rate by learning parameterized denoising model (ore(Tr y) 4o ( yy. Alternatively, Meng al. [20] and Lou et al. [11] observed that the ratio st(v)= - (v)))/a.(() plays similar role the continuous score function V log qe(y), up to additive constant. Therefore, they proposed to learn an approximate score function instead, of (v).
The training objective is negative evidence lower bound (ELBO), which is function of the approximate backward rate Re ()) (x y))( (v)
((0, - 1/" Eggio(lso))
[R (v. - Rs (xx,1) ( zo) log( (v, x))} di + C, (2)
Σ #F
moly mo)
where C' is a constant independent of A.
Given the intractable sums over the exponentially large state space S above, it is natural to constrain the forward process to be composition of identical independent processes on each dimension, where the rate over the entire state R2 can be written as
Re(x,y) = 3(t Ro (
(3)
d=1
for a base rate matrix RG E RSxS absorbing the time dependence into the scalar coefficient 3(t). Here, v represents all components of x apart from the d-th dimension. Backward rates are also factorized with this forward rate factorization, allowing the objective in (2) to be nicely decomposed into a sum over dimensions, as detailed by Campbell et al. [8].
2.2 Absorbing Discrete Diffusion
For generic discrete data, two forward processes are most commonly used: the uniform process and the absorbing state process [7]. We focus on the absorbing state process in this paper, for which the base rate matrix for each dimension can be written
yd MASK Idd + MASK
Rabsorb Rig (xd yd)
yd xd + MASK
0
otherwise,
where we augment the vocabulary set by introducing the MASK token, which represents the absorbing state. The absorbing state diffusion is important to consider because it is very commonly used model sequence data without ordinal structure (e.g., text and protein sequences). is also intimately connected to masked language models [7, 22].
We use special notation for the absorbing state process. For a sequence use the mask operator y (x) to denote the sequence obtained by changing the d component of x: to MASK. Finally, we denote the masked indices of x by M(xx) = (d: rdd = MASK). The set of non-mask indices is the complement, M(x).
{MASK, 1, .,..., S).
Recently, Shi al. [12] introduced the MD4 diffusion model, which outperformed previous discrete diffusion models on both text and pixel-level image generation benchmarks. Furthermore, they noted that several continuous time denoising diffusion objectives (including (2) ) are mathematically equivalent and can be greatly simplified in the case of masked diffusion. The simplified MD4 loss is:
A L/" (eppo(
Lm(0,
Σ dEM((T)
log Dog. (ag x:) dt,
(5)
1-at
where ar A exp (-S6 ((s ds) is the survival probability of the token each position and af indicates its time derivative. We denote this loss as L. since it sums over the masked dimensions. As we will show later, this loss has counterpart, L that can be combined with the MD4 loss to yield better results for training diffusion model.
3

---

Informed corrector
Informed
Predicted
1
1
1 2
2
2
2
3
2
Uninformed corrector
Corrected
1
2
0.010
No corrector
Errer rate
Uninformed corrector
0.006
Predicted
2
2
Corrected
1
2
2 2
2
2
0.002
Figure Demonstration of the failure modes of the uninformed corrector the Markov chain example. The predictor made a mistake (red) in sampling an impossible transition, which is immediately corrected by the informed corrector. The uninformed corrector, however, masks out tokens at random, while ummasking other tokens at the same time. Changes made by the correctors are highlighted in orange.
0.000
128 Number Function Evaluations (NFEs)
The informed cor- rector substantially lowers er- rates the Markov chain synthetic experiment.
Figure 2:
2.3 Sampling From the Backward Process
Camphell et [8] proposed to use tau-leaping, an approximate sampling method that allows simultaneous updates multiple dimensions for single model evaluation. Tau-leaping assumes that the backward process transition rate constant between [t- T,,,, f). regardless of whether states change during that interval. Under this assumption, multiple transition events in the interval can be applied simultaneously. Leverging the structure of masked diffusion, MD4 [12] used a simpler approach. Instead of trying to simulate the continuous time process, they directly sampled from the finite-time backward conditional distribution peje 21). They call this approach ancestral sampling.
Corrector steps. Both the above approaches treat the one-step backward update as independent across dimensions, which accumulates simulation errors the course of the backward process. To combat this problem, additional corrector steps can be used to ensure the marginal distribution of samples 2, time matches qe. Campbell et al. [8] show that corrector process with rates,
R5 y) = Re(I, y) + 2)
(6)
leaves the stationary distribution ge invariant. In practice, one must the learned backward rates Ieς approximation for the true backward rates ft, This can be thought of as "uninformed" corrector, as it is simply a summation of the forward and backward rates, neither of which explicitly targets problematic tokens in the masked diffusion setting. Nevertheless, Campbell et al. [8] show that this uninformed corrector increases sample quality in practice.
Correcting for the absorbing diffusion. The necessity of corrector steps is particularly apparent when tau-leaping is used in conjunction with an absorbing forward process. the reversal of the absorbing process, only transitions from the absorbing state to other states are allowed, which is reflected by the structure of R,, Thus, once token has been generated. it cannot be erased changed, rendering it impossible for the predictor to fix its own errors. wewward-backwee mitigate this issue by introducing the forward term RG, allowing transitions in the forward direction. resulting in generated tokens being "masked out" again with uniform probability.
corrector steps
Though forward- backkard corrector steps solve the issue in theory, they are far from optimal in practice. The steps can only correct mistakes by masking tokens at random. Ideally, we would like the corrector to make informed transitions based on which states are more likely under the model. corrector.
This observation motivates us to find better alternative to the
Methods
To address the limitation of existing uninformed correctors, we design corrector scheme that uses the model inform which tokens to correct, prioritizing tokens that are more likely have errors. Next, we introduce hollow wransoormr-bbasee architecture to parameterize this informed corrector, along with tailored training objective to enable efficient training of the corrector. We describe these methods in detail below.

---

3.1 Informed Correctors
To design corrector scheme for our sampler, we first identify Markov chain that has the marginal distribution as a stationary distribution. For discrete distributions, a straightforward option is to use Gibbs sampling, i.e., we iteratively select an index from d 1...., D and resample the d-th dimension according to the conditional distribution
xd ge
There are two common ways of selecting the index d. known systematic scan and random scan. Systematic sean enumerates the dimensions following a fixed permutation, whereas random sean selects the dimensions uniformly each resampling step. In both cases, it can take many steps for the Markov chain to "correct" simulation error in one the dimensions.
Informed selection. To improve efficiency, we use an adaptive random-scan Gibbs sampler with a non-uniform selection distribution [25-28]. We choose this distribution so that the corrector prioritizes updating those dimensions that are more likely to have simulation errors. To achieve this, we let Ca be a confidence score indicating the likelihood correctness the d-th dimension of the current sample. We select the indices to update in the next steps dj...., dy by sampling exp(-ea/T)
without replacement from the distribution Cat( Ee exp(=).//7) where T temperature parameter.
This sample can be generated using the Gumbel-top-k trick [29, 30]: First draw /) independent Gumbel random variates, Id Gumbel(O, 1), one for each dimension. Then rank the dimensions in descending order of - + 9u and select the top k dimensions.
We consider two choices of the confidence score. The first option is ca log (xx (i)), the log- bog-probability of the d-th dimension of the current sample given all other dimensions, where:
rdd MASK
gr (x
(8)
x (x=
xd * MASK
and 9o (xx (x) the denoising conditional distribution. An alternative confidence score is max log ge (ad = 1
margin d
log ge (xr (x)))
Similar to in [31], the advantage of this definition is that it maintains low confidence for di- mension when there is an alternative token that the model assigns comparable high probabilities (i.e., max log qe (ar'd x xx)d) is equally large). practice, we often use this definition of the confidence we found it consistently outperforms the former.
Parallel updates. Rather than updating the & dimensions sequentially, we update them in parallel, as in Hogwild Gibbs sampling [32]. In practice, we want to keep & small to avoid errors in the parallel resampling, while keeping large enough so that errors can be fixed reliably. We study the effects of this hyperparameter in section 5.3.
Conditioning on the mask configuration. In masked diffusion, sample x, qe(-) contains both mask and non-mask tokens. The Gibbs-inspired update above might turn a mask into a non-mask token or vice versa. This is nonoptimal generate new tokens but rather to fix errors introduced by the predictor steps; on the other, re-masking tokens that the predictor will later unmask creates redundancy. To improve corrector efficiency, instead of the marginal distribution ge( we target the mask-conditional distribution,
on the one hand, the corrector steps are introduced to
q M(z)) qe( (MM)))M(M))-MM
(10)
In other words, starting from the current sample we fix the mask configuration M(z) and only allow the corrector to change non-mask tokens into other non-mask tokens. To update position d E M(x), we now simply need sample from the denoising conditional distribution:
dd Gojo d 1d)
(11)
'Other Markov chains could be employed instead. For example, the barkkardd corrector proposed by Campbell et al. [8] leverages a birth-death process that has qı as the stationary distribution [23]. We also tested -balanced Markov processes [24], but their performance inferior Gibbs.
5

---

Intuitively, this is true because a non-mask token at time t must be consistent with the token from the true token at time given the masking forward process. This is convenient for us, since the only distribution required by the informed corrector scheme is Goja (xx xx))). For derivation and more implementation details, see appendix B.
Next, will discuss how implement this corrector scheme efficiently, that only network evaluation needed to select and update positions d1....,dy simultaneously.
3.2 Parameterizing the Conditional Distribution Using Hollow Transformers
Since the MASK token provides no information about the original data, we make the further observa- tion that goje (xt xx)) quja (xd (z)))) which is the familiar denoising distribution that we can approximate using P
(x)).. Therefore, we can implement the corrector on dimension d by (2* (ar
calling our denoiser on (z). However, correcting & dimensions in this case requires forward passes through the denoiser. In order to simultancously obtain Goje( .d x x'd) for multiple d, we need an architecture where information about zd does not propagate to the output at the dth position.
We solve this problem with hollow transformer [33], two-stream self-attention architecture where each stream runs causal attention in a different direction. The embeddings are then offset by one position so that the embedding position does not attend the corresponding input position. The two streams are then summed together and sent through multilayer perceptron layers to the output. By construction, the hollow transformer satisfies the property that for all (f, 0),
Ad yy - (x, 1 (y, ()ds
(12)
where f (xx f)a represent network output at sequence position d with inputs (2), t) and parameter In other words, the output at dimension d agnostic the input at the same dimension We expand upon architectural details of the hollow transformer in appendix C.
Once we have the hollow transformer, we can use it to parameterize the denoiser, which is learned through regular masked diffusion training. This denoiser can then be used for both the predictor and corrector steps in sampling. Nonetheless, we observe that switching from standard transformer to hollow transformer can sometimes degrade the predictor performance. thereby diminishing the overall benefit of correctors. these situations, we propose to keep the original denoiser architecture and train a separate hollow transformer for the corrector. We will show this in section 5.3.
3.3 Learning the Hollow Transformer with the ELBO
The preceding section demonstrated that the corrector can be trained using the same ELBO objective employed for training the denoiser. Interestingly, we reveal that for hollow transformers, there an alternative expression of the ELBO that offers training signals that complementary to the widely used simplified objective (20). This new expression is presented in the following proposition.
Proposition 3.3.1 When R₃ Raghorb, the continuous-timg time function (2) simplifies to
= (0, = =f/ Bano(ol a
log (2 M²)))) dt C.
(13)
dEM(()
We include a full derivation of this result in appendix A and show that is related to L. via simple label-switching trick. Notably, Campbell et al. [8] and Shi et al. [12] both avoid writing the loss in this form, since evaluating ,0,("" (xd (x)) for each different d requires a separate call to the denoising network, which is prohibitively inefficient in practice. The hollow transformer fixes this issue, as now the inner sum over d can be evaluated with one forward pass of the network.
Furthermore, observe that this different form of the ELBO complements the MD4 loss in (5), in the sense that it sums over the mank dimensions M(x), while MD4 sums over the masked dimensions M(x) instead. Both losses ignore half of the learning signal from the sample ~ qooo xo), which seems unideal. This observation suggests that a more informative objective can be obtained by averaging the two losses:
LHD 1/2 (C + L.).
(14)

---

Since both terms are different expressions of the same ELBO, LHD can be viewed as applying variance reduction trick for the sample estimator. We observe experiments that this objective indeed outperforms its individual components and makes training converge faster (see appendix D.2).
We term our hollow transformer masked diffusion model trained with loss (14) HollowDiff.
4 Related Work
Alongside the development of discrete diffusion models, other classes of non- autoregressive models employing similar iterative generation principles have emerged [34, 35, 33]. Among these, MaskGIT [35] is notably similar to masked diffusion models [12, 36], particularly as both are trained using denoising cross-entropy losses. However, key differences distinguish them: MaskGIT utilizes non-likelihoodd-aasee weighting for its loss function and employs confidence- based decoding order. In contrast, masked diffusion models typically unmask tokens in a random order that aligns with their theoretical backward process. Inspired by MaskGIT, large language diffusion models such as LLaDA [37] adopt similar heuristics for remasking tokens at each prediction step. This effectively alters the predictor's unmasking order based on confidence, which can introduce bias into the sampling process, causing it to deviate from the diffusion model's 's theoretical reverse process. In contrast, our confidence metric is only used to define the selection distribution of random-scan Gibbs sampler, and the informed corrector preserves the marginal distributions of the original diffusion framework first order.
MaskGIT.
Discrete diffusion correctors. Campbell et al. [8] and Gat al. [38] introduced the forward- backward and DFM correctors for discrete diffusion. However, these "uninformed" correctors can only address errors through inefficient random re-masking. ReMDM [39] proposed the remasking diffusion process that includes forward-backwer combining the predictor and corrector steps into one. Relatedly, DPC [40] proposed a corrector scheme for MaskGIT sampling. Although their method can be seen as "informed", they reportedly require hundreds of steps to achieve good sample quality on ImageNet 256x256.
and DFM as special cases, while effectively
Experiments
In section 5. we test the informed corrector against other sampling methods in a hidden Markov model setting where the true denoising conditional distribution is known. We find that informed correctors consistently reduce error rates across wide range of function evaluation counts. This observation is confirmed in section 5.2 on real-world data, where informed corrector also greatly reduce spelling errors in sampled text. In section 5.3, we informed correctors for tokenized image synthesis, achieving competitive FID scores on ImageNet 256 x 256 with a 230M parameter model and only small number of function evaluations.
5.1 Hidden Markoy Modeling
To check if informed correctors improve sample quality, we propose a simple setting where the true denoising conditional distribution is easily accessible. Consider sequence data x E {O..... 1}1 sampled from Markov chain: (u(x) qu(x') ITo qu(xx zd-1), where qo (x d-
j) Ai and z! is sampled uniformly. At time of the forward process, the denoising model observes a masked version of this sequence:
( mо) = II qe d=1
ol
= II fa az ( d=1
+ + az) (at -ASK))
(15)
). This can be seen performing inference in a hidden Markov model (HMM), where the prior is the Markov chain go. and the observation model is 2rjo- Given this, we can perform standard message-passing inference [41] to
and tries to predict the conditional demoising likelihood (=] )))
find the posterior (2) 22)))
7

---

Table 1: A comparison of guidance-free discrete diffusion and MaskGIT models ImageNet 256 x 256. We report two informed corrector approaches (see section 3.2): "HollowDiff + informed corrector" uses hollow transformer for both predictor and corrector, while "MD4 informed corrector" uses a separately trained hollow transformer model to correct predictions made by the MD4 model. We use 8 predictor steps, corrector steps, and 1 final denoising step.
Method
Params
# Steps
IS f
Precision
Recall
MaskGIT (42)
227M
6.56
203.6
0.79
0.48
DPC [40)
391M
180
4.45
244.8
0.78
0.52
VQ-Diffusion [43]
518M
100
11.89
ReMDM [19
227M
16
7.40
145.27
MD4
230M
17
6.28
175.32
0.79
0.48
Uninformed corrector
230M
17
7.41
174.49
0.80
0.44
+ ReMDM sampler
230M
246.34
0.85
0.40
Informed Corrector (Ours)
HollowDiff informed corrector
230M
6.26
188.18
0.78
0.45
MD4 + informed corrector
400M
17
5.78
212.52
0.82
0.43
To demonstrate the problems with existing samplers, we choose a structured transition matrix A:
j j=i
A.j 1 p j= (i+11) mod S
(16)
otherwise.
At each position, the token xd has probability of staying fixed and probability 1 - of changing into the next token. To evaluate sample quality, one important metric the number of errors in the samples, i.c., occurrences of impossible transitions within the sampled sequence. The overall error rate is computed by:
D-1 Σ
Tem ((Tm) N(D 1)
Σ
(d * 1 (ad #(25 mod S}*
(17)
Since the samplers have access the true denoising distribution, the error rate will be if no more than one token is generated in a single step. However, if more than one token is generated during parallel update, there a risk of generating impossible transitions.
We show error rates (mean and standard deviation over 5 random seeds) in Fig. 2. We find that ancestral sampling with the informed corrector yields the least error for all examined NFE settings. Meanwhile, the uninformed corrector consistently performs worse than using no corrector at all, showing that uninformed corrector step is not as valuable setting (see fig. 1). More experimental details can be found in appendix D.
additional predictor step this
5.2 Character Level Text Modeling
We next evaluate our method on the character-level modeling task using the text8 dataset [44]. Follow- ing standard practice, we use the provided dataset splits and train on text chunks of length 256 for I million steps using a batch size of 512. As a simple measure of sample quality, we compute the error rate, defined as the fraction of total characters in words that are absent from the training set vocabulary.
Informed
MD4 (uninformed
MD4 (no correctari
MD4 (RaMDM)
40
Error Rate (%)
30
20
10
We benchmark against MD4 [12], which employs standard transformer architecture and has a compara- ble parameter count (fig. 3). HollowDiff significantly outperforms the standard MD4 baseline in MD4 en- hanced with the uninformed corrector and the ReMDM sampler, demonstrating the utility of our proposed sam- pling framework. We direct the readers to table for examples of generated text.
01
16
32
64
128
256
Number of Function Evaluations (NFEs)
Figure Comparison of word error rate across different sampling methods and models the text8 dataset.

---

5.3 Class-oodditinaal Image Modeling
We evaluate method tokenized ImageNet 256x256 using the VQ tokenizer from [43]. All models are trained on the standard training split and evaluated without classifier-free guidance. Table presents a comparison of generative models in the guidance-free setting. HollowDiff with the informed corrector achieves an FID of 6.26, using only 17 total steps (8 predictor, 8 corrector and final denoising step). This outperforms several baselines with similar or larger model sizes and substantially more sampling steps. VO-Diffusion [43] requires 100 steps to reach FID 1 11.89, while ReMDM [39] achieves FID 7.40 16 steps. We also note that the MD4 baseline achieves similarly strong results without any correctors, and the uninformed corrector and ReMDM sampler fail to improve performance further. Generated samples from our method are shown in appendix D.3.5.
Transformer architectures. Without using correctors, Hollow DoflowDiff significantly underperforms MD4 in FID (10.24 vs 6.28 in table 2), indicating Hollow transformer's subpar denoising capabilities compared to standard transformer image data. This deficiency is offset by its capability to "self-correct" via informed correctors.
Table 2: FID scores (ImageNet 256) for MD4 vs HollowDiff under various corrector configurations.
Model Corrector
FID
To further demonstrate the value of informed correctors, also test a hybrid setup where MD4 serves as the predictor and HillowDiff improves FID to 5.78, exceeding the performance of both individual models (see table 1 for full set of metrics). Despite the hybrid setup using more total parameters (400M) than the individual HollowDiff (230M), it highlights the modularity and compatibility of our corrector: it can be applied to stronger pretrained predictors without need for joint training. This finding confirms that our corrector provides complementary gains outside of architectural choice, and is broadly applicable to discrete diffusion models in general.
Hollow (no corrector)
10.24
Hollow uninformed corrector
10.93
DiflowDiff informed corrector
6.26
the corrector. This combination
MD4 (no corrector)
6.28
MID4 informed corrector
5.78
MD4 model
6.5
Score
6.0
5.5
5.0
investigate how the num- ber of parallel updates affect sample quality, conduct an ablation test on with the hybrid MD4 + informed corrector setting. Figure 4 shows the FID scores function of k. across different total number of update steps (evenly split between predictors and correctors). For each number of up- date steps (16, 32, 64, 128), we find that one temperature parameter 7 {1.0, 2.0, 5.0, 10.0} respectively works best across all values of k.
Number of parallel updates.
4.5
2
32
64
Figure 4: FID number of in- formed corrector updates under var- ious sampling budgets. Increasing k initially improves efficiency but de- grades sample quality for large k.
We observe that for small NFE budgets, increasing & significantly improves sample quality up to point. For higher NFE budgets, performance is much less sensitive to k, and a small value (e.g., k=1 or k= already yields near-optimal FID. This aligns with the intuition that predictor steps dominate performance at high compute budgets. When too many dimensions are updated simultancously. their mutual dependencies are ignored, leading to update collisions where corrections can contradict or undo each other. This is observed empirically as in all cases, too large values of i (e.g., k 2 32) lead to worse FID.
Discussion
We proposed informed correctors and an accompanying training framework for learning and sampling from masked discrete diffusion models. We showed that the uninformed forward- backwrd corrector not well-suited to the masked diffusion setting, since does not leverage information about the conditional likelihood of the generated tokens. Our approach, HollowDiff, produces superior sample quality with fewer model evaluations than alternative discrete diffusion samplers both in the synthetic setting and real-world datasets. Overall, work serves to further enhance the capabilities

---

of diffusion models in discrete regimes and to bridge the sampling quality gap between discrete diffusion models and other types of generative models.
Limitations and future work. Using the informed correctors for absorbing diffusion requires the of the hollow transformer architecture. As result, cannot apply the correctors directly to arbitrary pretrained models. Learning the informed corrector also presents tradcoff's between quality of the denoiser and that of the corrector. While the HillowDiff scheme represents a compact algorithm for learning to generate and correct at the same time, practically, better performance is achieved by learning separate predictor and corrector models and combining them. Future work could investigate how this scheme can be made more efficient, potentially by architectural or algorithmic improvements on hollow transformers.
Impact Statement. This paper aims to improve the efficiency and quality of discrete generative models, which have broad applications in areas such image synthesis, text generation, and molecular design. Like many generative modeling techniques, the methods developed here could potentially be misused for generating misleading or harmful content (e.g., deepfakes or disinformation). We encourage responsible use and stress the importance of deploying generative models with appropriate safeguards and monitoring.
Acknowledgements
This research was supported with Cloud TPUs from Google's TPU Research Cloud (TRC). Y.Z. was partially supported by the Stanford Interdisciplinary Graduate Fellowship (SIGF). S.D. and F.C. were partially supported by the McKnight Foundation, the Simons Foundation and an NIH CRCNS grant RO1DC020874.
References
[1] Jascha Dohh-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamies. In International Conference on Machine Learning, pages 2256-2265. PMLR, 2015.
[2] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 32, 2019.
[3] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-6851, 2020.
[4] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International Conference on Machine Learning, pages 8821-8831. PMlr, 2021.
[5] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10684-10695, 2022.
[6] 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: 14:12454-12465 2021.
[7] Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34:
2021.
[8] Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnad Doucet. continuous time framework for discrete denoising models. Advances in Neural Information Processing Systems, 35:28266-2827,, 2022.
[9] Javier Santos, Zachary R Fox, Nicholas Lubbers, and Yen Ting Lin. Blackout diffusion: Generative diffusion models in discrete-state spaces. arXiv preprint 230::2305.11089, 2023.
10

---

[10] Lin Zheng, Jianbo Yuan, Lei Yu, and Lingpeng Kong. A reparameterized discrete diffusion model for text generation. arXiv preprint arXiv:2302.057.3,, 2023.
[11] Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion language modeling by estimating the ratios of the data distribution. arXiv preprint arXiv:2310.1633 2023.
[12] Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, and Michalis K Titsias. Simplified and generalized masked diffusion for discrete data, arXiv preprint arXiv:2406.04433 2024.
[13] Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian Trippe, Jason Yim, Helen Eisenach, Woody Ahern, Andrew J Borst, Robert J Ragotte, Lukas F Milles, et al. De novo design of protein structure and function with REdiffusion. Nature, 620(7976):1089-10
2023. [14] Sarah Alamdari, Nitya Thakkar, Rianne van den Berg, Alex Lu, Nicolo Fusi, Ava Amini, and Kevin Yang. Protein generation with evolutionary diffusion: sequence is all you need. In NeurlPS Generative Al and Biology (GenBio) Workshop. 2023.
[15] Clement Vignac, Igor Krawezuk, Antoine Siraudin, Bohan Wang, Volkan Cevher, and Pas- cal Frossard. DiGress: Discrete denoising diffusion for graph generation. In The Eleventh International Conference on Learning Representations, 2023.
[16] Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin T Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. arXiv preprint arXiv:2406.07552, 2024.
[17] Jingyang Ou, Shen Nie, Kaiwen Xue, Fengqi Zhu, Jiacheng Sun, Zhenguo Li, and Chongxuan Li. Your absorbing discrete diffusion secretly models the conditional distributions of clean data. arXiv preprint arXiv:2406.0377 033366 2024.
[18] Joe Benton, Yuyang Shi, Valentin De Bortoli, George Deligiannidis, and Arnaud Doucet. From denoising diffusions to denoising Markov models. Journal of the Royal Statistical Society Series B: Statistical Methodology, 86(2):286-301, 2024.
[19] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
[20] Chenlin Meng, Kristy Choi, Jiaming Song, and Stefano Ermon. Concrete score matching: Generalized score matching for discrete data. Advances in Neural Information Processing Systems, 35:34532-345555 2022.
[21] Jacob Devlin, Ming-Woo Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 2018.
[22] Emiel Hoogeboom, Alexey A Gritsenko, Jasmijn Bastings, Ben Poole, Rianne van den Berg, and Tim Salimans. Autoregressive diffusion models. arXiv preprint arXhv:2110.02037, 2021.
[23] Jiaxin Shi, Yuhao Zhou, Jessica Hwang, Michalis K Titsias, and Lester Mackey, Gradient estimation with discrete stein operators. arXiv preprint wwxiv:2202.004999 2022.
[24] Giacomo Zanella. Informed proposals for local MCMC in discrete spaces. Journal of the American Statistical Association, 2019.
[25] Richard A Levine and George Casella. Optimizing random scan Gibbs samplers. Journal of Multivariate Analysis, (7110)2071-2000 2006.
[26] Richard A Levine, Zhaoxia Yu, William G Hanley, and John Nitao. Implementing random scan Gibbs samplers. Computational Statisties, 20 2005.
[27] Krzysztof Latuszyáski, Gareth O Roberts, and Jeffrey S Rosenthal. Adaptive Gibbs samplers and related MCMC methods. The Annals of Applied Probability, pages 66-98, 2013.
[28] loannis Mitliagkas and Lester Mackey. Improving Gibbs sampler sean quality with DoGS. In International Conference Machine Learning, pages 2469-2477. PMLR, 2017.

---

[29] Chris J Maddison, Daniel Tarlow, and Tom Minka. A sampling, Advances in Neural Informa- tion Processing Systems, 27, 2014.
[30] Wouter Kool, Herke Van Hoof, and Max Welling. Stochastic beams and where to find them: The Gumbel-top-k trick for sampling sequences without replacement. In International Conference on Machine Learning, pages 3499-3500 PMLR, 2019.
[31] 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.06763 2025.
[32] Matthew J Johnson, James Saunderson, and Alan Willsky. Analyzing Hogwild parallel Gaussian Gibbs sampling. Advances in Neural Information Processing Systems, 26, 2013.
[33] Haoran Sun, Lijun Yu, Bo Dai, Dale Schuurmans, and Hanjun Dai. Score-based continuous-time discrete diffusion models. arXiv preprint arXiv:2211.16750, 2022.
[34] Nikolay Savinov, Junyoung Chung, Mikolaj Binkowski, Erich Elsen, and Aaron van den Oord. Step-unrolled denoising autoencoders for text generation. In International Conference on Learning Representations, 2021.
[35] 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, pages 11315-11325, 2022.
[36] Kaiwen Zheng. Yongxin Chen, Hanzi Mao, Ming- Yu Liu, Jun Zhu, and Qinsheng Zhang. Masked diffusion models are secretly time-agnostic masked models and exploit inaccurate categorical sampling. In The Thirteenth International Conference on Learning Representations, 2025.
[37] 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.09992, 2025.
[38] 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:133345-133385 2024.
[39] Guanghan Wang, Yair Schiff, Subham Sekhar Sahoo, and Volodymyr Kuleshov. Remasking discrete diffusion models with inference-time scaling. arXiv preprint (XXi:25003.00007 2025.
[40] Jose Lezama, Tim Salimans, Lu Jiang, Huiwen Chang, Jonathan Ho, and Irfan Essa. Discrete predictor-correctn diffusion models for image synthesis. In International Conference on Learning Representations, 2022.
[41] Kevin P. Murphy. Machine Learning: A Probabilistic Perspective. MIT Press, 2012.
[42] José Lezama, Huiwen Chang, Lu Jiang, and Irfan Essa. Improved masked image generation with token-critic. In European Conference on Computer Vision, pages 70-86. Springer, 2022.
[43] Shuyang Gu, Dong Chen, Jianmin Bao, Fang Wen, Bo Zhang, Dongdong Chen, Lu Yuan, and Baining Guo. Vector quantized diffusion model for text-to-image synthesis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10696-10706, 2022.
[44] Matt Mahoney. Text8 dataset. http://tttaaaooeye.eeeeee/artaaa html, 2006.
[45] Zhilin Yang. XLNet: Generalized autoregressive pretraining for language understanding. arXiv preprint 17Xiv:1906.08233 2019.
[46] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, 2019.
12

---

A Diffusion Objectives
A.1 Equivalence of objectives
For a complete picture, we include the derivation of the simplified objective in Shi et al. [12], Sahoo al. [16]. Ou et al. [17] and show how connects to the Campbell et al. [8] objective expressed with CTMC transition rates. A similar treatment can be found in Shi et al. [12, App. H.I]. The derivation highlights that while the two expressions are mathematical equivalent, they nonetheless exhibit important differences in when it comes to estimation cost.
We start from Campbell et al. [8]'s continuous time discrete diffusion ELBO objective. For clarity, we use variable names v. y. and z in a way that indicates temporal order:
[())= / (w)rr {E x"yy ((((vx'))
(Eryy Re(v.2)) log( (2, dt
/ (n(urr(a (V( Rv( log( RR. (2. w))] dt,
(18)
where (y) & R) = =E Ru z²). Here r=ly)) represents the probability of transition- ing from state y to state 2, given that we know a transition happens at time t:
Be( R.(
リ 2
r2(2 v)
=
0
==2.
We can rewrite the objective with a label switching trick:
(()) = / E (ego (i((() (w) (v) log( 12) (2, w)) di
E (x0(
(2) Σ
Ry log( v)) dt
-f Ego (ee)eee (20)
Σ zz4
R2 (v)
Σ
Rely: =) 1o (2, di
=f (
ΣΣ (alv ao) 2) ( za) Re) 2) 2og Rog(RR R% (2.8))
dt
= -f (
ΣΣ ェデツ
EE(a(u (: 20) Rg ge(= Re log RC (u. (V.(
dt. (19)
Note that we switched the labels y and z on the last line, which is the key trick that simplified our derivation. Proceeding from line (19) becomes straightforward:
(())= (eo
Σ zΦリ
Ri R2(2, v)
d:(z ao) q: xo)
log( log(1
di
= = - (
Repoys x) log log (ig( (M. Σ [((()))) 241
dt, (20)
where Rejo x) = Re y) 30l0 is the true reverse transition rate conditioned on No. and we renamed to to preserve alphabetical order of the variables for readability. Now we have simple Monte Carlo estimator for this objective: we just need to sample y from the marginal distribution and then evaluate the parameterized transition model R2 (y, x)) for neighbors x of y, which can be done using only one model evaluation.
13

---

A.2 Simplifying the Score Function for Masked Diffusion
The relationship between the (concrete) score function and denoising distribution can be simplified in the case of masked diffusion. We first start with the relation between the denoising distribution and score function in [8]:
(x(x) q: (:(1))
sc(y)
90 ( I v).
(21)
x0))
For a state z where xt 4 MASK, consider the score from Md (z) to æt
gnj(2 zo)
(M = E
Poje (2o (x)
o
(22)
(e(((M(x)) FO
010o(2" (2))
(M 2g)) Poje (0)
(2 (x))
(23)
ueooo | 2²
(MaK andy Polc(a" Mx)x))
(24)
a- Po)r(0d (x))
0z
(25)
where O; represents the survival probability of any non-mask token at time t, which assumed to be constant for all dimensions and token values. Going from (23) to (24) we used the fact that x. = x if x f MASK. This formula tells us that in the case of masking diffusion, the score function between mask and non-mask tokens is just constant factor away from the denoising distribution.
In this work, we use a hollow transformer f that outputs a x S array, and we use (z, t)d.e to represent P (zd = (x)). We can also omit the dependence of the network here because in absorbing state diffusion the true 20 (a: MMdx))
does not depend on t (proof omitted).
A.3 Deriving the HollowDiff Objective
The hollow transformer provides us with way to optimize the optimize the Campbell et al. [8] objective (18) directly:
- log p(zo) & L(0) = /0 E (ospoollooore(all))
R. (v) - Re( log (2, w))} dt + C,
(26)
E aayo (mo)
Rr(x.. u)s (v)) Σ
Σ
Rv(u,)) log( R2 (2,
di C
(27)
1-- T
E g (yl
Σ
Re u)s (y)= - >
Ro 2) 200(())((())))) dt + C.
(28)
Notice that we can further simplify the second term by leveraging the sparsity structure of the forward process:
Σ RR(y, log(s = Σ
Rely, log (og((
(29)
2
dEM(2)
Σ
Raly, My())) log( I- a
8 (V)lag))
(30)
dEM(s)
0/
Σ Hog(f((u)a,) C,
(31)
a
dEM(y)
where of represent the time derivative of 0g. Note that here we used the hollow property (y)2 (MM which allows to evaluate the sum over E M(y) using only one function evaluation
14

---

of f. Similarly we can write out the first term, which conveniently reduces to constant due to the structure of the forward process:
Σ Re( v)s (v),
Σ
Σ
Re(2, v) (v)
(32)
1 - at deM((y) a:x" MASK,
1i =y
of 1-a
Σ Σ
(V)a
[M(3)]-
(33)
04
dEM(U) F.E AMASK,
Here we are using the fact that tokens and therefore sums to Putting these back together, we have:
1 (y)a is vector representing the distribution over the non-mask
()) /'" d E
oo mo)
Σ
log(s" (v)((())) de C,
(34)
dEM((
which is very similar the MD4 loss (5) of Shi et al. [12], except that the sum over all mask positions is replaced with the sum over all non-mask positions and the coefficients are different.
In practice, we want to maximize learning signals for the model. Since the objective (34) only uses the non-mask dimensions y, we propose to combine it with the MD4 objective which uses the mask dimensions instead. These two simplified objectives are mathematically equivalent and have the same optima, but can suggest different stochastic gradient estimates. The combined objective is then
((0) ! E 2
O/
log(f (og(f(())) (w)d.ag) + 1 dEM(N)
Σ
Σ
log(f (v) dt + +
Late
Ote deM(U)
(35)
where we have changed the index in the first term from vd to zg for consistency, since y²= ag when is an unmasked dimension. Note that the HollowDiff objective leads to a lower-variance stochastic gradient estimator, as it leverages more learning signal for single sample of We confirm this empirically via loss curves in appendix D.2.
B Informed Corrector
We explore informed corrector updates that leverages the hollow transformer architecture. Specifically, consider the Gibbs operator:
ge (axd
Σ
(0)( 10)001(00 x//d) = ar90 x xd) (a (Mss)) + (1 (zr)(((AM
(36)
To approximate this operator, we can use the learned network (x)a, in the place of 20 (ojo i v')),, since the hollow transformer architecture blocks out information of ad on the dth output dimension.
In order to reach the stationary distribution (z),, we need to repeat the Gibbs update, iterating over all dimensions d. This is very inefficient in realistic settings. Instead, we choose prioritize updates on the "most unconfident" dimensions, where Id is most likely to be affected by the compounding errors in the simulation of the backward process. To achieve this, rank the dimensions according to the confidence score Ca and perform k updates at the same time.
One observation is that changing the mask configuration does not actually do anything useful for us, since the mask configuration does not contain any information. This suggests that we should instead consider the distribution conditioning on the mask configuration:
(x x id pd # MASK)=
Σ (x (
22 xd * MASK) 2dd, xd * MASK)
(37)
Σ
E (ad)) I 110d=ag) (20 wi/))
(38)
= I [ed (x xld).
(39)
15

---

Repeatedly applying this operator on the non-masked dimensions creates Markov chain with the stationary distribution equal to (z (Ma) _ gr(æ æ = MASK) where Mr is the set of masked indices after the predictor update.
This translates to a straightforward procedure: find the non-mask dimensions with the lowest scores, and then sample from the model predictions Poj (ad 22)d) of that dimension. For real-world datasets, we found that doing an argmax update in the end, where all dimensions are updated to the token that maximizes the conditional likelihood can be helpful for sample quality. We set the number of correctors C for the majority of our experiments. The algorithm is outlined in algorithm 1 below.
Algorithm 1 Backward process with informed corrector steps
Require:
fmin: 0, temperature
Require: number of updates k number of predictor steps P, number of corrector steps C.
1: Initialize time t +- 1
2: Initialize sample x MASK
3: Set predictor step size At + (t - tmin))))
4: for step = 1 to do
5: Compute denoising probability pld, = f (d)
6: Apply predictor step x <- PREDICTORUPDAATE(C..P
7: Update time +-1 - At
8:
for cstep = to do
9:
Find most likely alternative tokens x² = arg max,1s log Pole (a'd (2)))
10:
Compute confidence score c|d] log (ogpp)((t (x _ (xv - log ( (2'd æ/d), vel M(x)
11:
Sample Gumbel noise g[d] ~ Gumbel 1)
12:
Compute ranking criterion r[d] = - c[d] + Tg(4]
13:
// Update the top-k dimensions
14:
for ent = to do
Get dimension of transition d' SORTED(T) ont]
Update state zd Cateeoriicall&y (x" | |x)dd)) end for 18: end for
16:
end for
20: // Deterministic update at the last step
21: Find most likely values for each mask dimension ad + argmax, og Polo(ag x)) for d E M(x)
22: Return x
C Hollow Transformer Architecture
C.1 Hollow Transformer Architecture Details
We build upon the hollow transformer in [33] and design class of hollow transformers that is more parameter efficient and expressive. First, we divide the computation into two streams, the content (causal) stream and the query (mixing) stream. The content stream consists of two transformers with causal self-attention in opposite directions, whereas the query stream attend to the content stream and outputs the final result. By design, non of the attention mechanisms allow a token at position d to attend to a token that has access to information of zd
We differ from the original hollow transformer in that introduce third stream combine information of the forward/backkward causal streams. Specifically, we have the attention layers F, Br, Mj that implement multi-head attention with different causal structures. The updates of these attention streams are respectively (ignoring multilayer perceptron layers):
Fri f- Attention F F KV FT")
(40)
Bi + Attention(Q Q = B KV = B²d
(41)
Md t Attention(( Q = h(M,, Film Bid KV (FFS. B2d))
(42)
16

---

where we insert a mixing layer for every m forward/backward layers. Notably, the new mixing layer M₁ like "query stream" " while the forward and backward layers Fi and B, act like "content streams" [45]. To prevent information leak and preserve the "hollowness," we offset the inputs by one position when initializing the forward/backward layers:
F Embed (Pad(z )), By Embed(Pad(z**1)))
(43)
and we initialize the mixing layer to zero. Due to the attention structure, each position M of the mixing layer have access to all the other positions z in the input æ, and therefore it cannot attend to itself at other positions to preserve hollowness. We double the feature dimension of the mixing layer and set the update function h to:
hM F. Jem Bkm) = Mi + Concat(FFem B
(44)
Finally, we use the output from the last mixing layer to compute the output logits.
Weight tying. Since the forward and backward content streams carry out the same type of com- putation, we use weight-tied network where all weights are shared between Fi and B,, which significantly decreases the parameter count of our network and makes training and evaluation faster. In practice, we have not observed significant downgrade in performance of the weight-tied network on ImageNet.
D Experiments
The main experiments in this paper are performed on v3-128 TPU pod and a v3-8 TPU machine with Google cloud.
D.1 Hidden Markov Modeling
D.1.1 Hyperparameter Selection
We select the hyperparameters for the correctors in the Markov chain experiments with grid search. For the informed corrector, the number of parallel updates & is swept over the values {1, 2, 4, 8, 16), and the temperature hyperparameter is swept over {0.01, 1, 5, 1., 2., 4. J. For the uninformed corrector, the corrector step size swept over (0.01, 0.1, 0.5, 1.0, 1.5, 2.0, 3.0, 5.0.
D.2 Text8
The experiments on the Text8 dataset were performed on machine with 8 NVIDIA H100 GPUs.
Training loss
Test loss
Training Objective
PM
1.6
104
105
106
10a
106
Steps
Steps
Figure 5: Loss curve on Text8 dataset [44]
We followed the standard dataset split and trained our models on text chunks of length 256 for million steps with batch size 512. We train two models: one with the masked loss LM and the other
with HollowDiff loss LuD on the text8 dataset of Mahoney [44]. The model consists 12 standard
transformer layers, and we add one mixing layer at the end to combine the forward and backward streams. We empirically confirmed that the model trained with our proposed Hollow Diff' loss function LHD is more efficient (see fig. 5).
17

---

D.2.1
Example samples
We show some generated samples taken using informed correctors and MD4 in table 3.
Table 3: Comparison of unconditional samples from informed corrector and MD4 trained on Text8. Samples generated with 256 NFEs. Words not contained in the training set vocabulary marked in red.
Informed corrector
MD4
n of european population cities in northern indonesia the german integration of the societies in the united states the expansion of the province with the public education reform one eight zero three the kingdom was held in one eight three five by the pe zero three zero two two four seven one zero zero three one three three five one one six four of the european union or trade union for it is the leniency of the netherlands five germany and three of the country s two five canadian publishing corp was created in one nine eight seven and helped days later it was designed or investigated according to the databases the date was also called out of the country in the one nine five zero s in the one nine six zero
a deco dsssom soylently taped several meter blue for time zone kutrata s mattom tte space gas one nine five six the celliba deitylic cylinder head designed by the longer gpu a thergameen ard glase nine five four nine four nine auto gav feature technology to harmonize anonymous aapproaches ragnar trifls livre by the word silva dam meaning small top or multi release good repodentiam defense genius multi pllo or hella spasca is derived from the verb meaning kill haltenere nickname me mean ccording to naldi and closere clinton are a lucid living scientist palger political views held into persist unlike high leg sacieties in sunghi s future which muscle meat low ly no chef is public catalysis of the frequent hebrid tlf vampires nympheo
two mine one zero
acturers of the developers rather than the starting elements of the day there are some useful models that restriction practice can be used the standard command to give the same instruction this practice will make them faithful with complete list of
ic other area two km tho west of shire mael nontionment home of the united planets ampstone six zero one km long in five zero five ft two zero five kh total ground opening gin ashers air force navy the nine seven ford mk one five two mk three fa mz four fi
D.3 Tokenized ImageNet
D.3.1 Training Details
Our model is adapted from that proposed by Sun al. [33] and consists 24 standard transformer layers. We use an embedding dimension of 784 and hidden dimension of 2048. We add mixture layers (1 after every 8 transformer layers) to the model to facilitate information exchange between the forward and backward streams. We train the model using AdamW [46] with an initial learning rate of le-4 and batch size of 512. We use linear learning rate warmup the first 100 steps. We adopt a stepwise learning rate schedule, dropping the learning rate to 3.3e-5 at 2. 2.11% x 10 steps and le-5 2.4 x 10 steps. We stop the training at 2.5 x 10t steps. The ImageNet experiments are done on v3- v3-128 128 TPU pod machines on Google cloud and training takes 100 hours of wall clock time.
18

---

D.3.2 Hyperparameter Selection
We select the hyperparameters for the samplers in the ImageNet experiments with grid search. For the informed corrector, the number of parallel updates & swept over the values {1, 2, 8, 16). and the temperature hyperparameter is swept over {.1, 1., 2., 10., 1000. For the uninformed corrector, corrector step size is swept over (.5, 1., 2,, 4.j. For maskgit, the temperature hyperparameter swept over (0.5, 1., 2., 4., 8., 10., 12., 16., 20., 40.).
D.3.3 Other Sampling Details
Final argmax update. We found empirically that applying final update of the form
20 argmax ( 2))))
(45)
of
for all d E (1,..., D} improves the results for the single network informed corrector experiments. Our intuition is that this final update helps eliminate remaining noise and local errors from all positions. When the base denoising model is stronger than the corrector model, is the case when the informed corrector is applied on top of MD4 predictors, this special update becomes unnecessary.
D.3.4 Ablation confidence metrics.
We compared confidence scoring schemes used in the informed corrector: the log-likelihood of the predicted token and the margin between the top-1 I and top-2 logits. As shown below (appendix D.3.4), the margin-base based confidence yields slightly better FID scores in both settings.
Model
Confidence Type
FID
HollowDiff + informed corrector
log-likeling -ogllike
6.45
HillowDiff informed corrector
margin
6.26
MD4 + informed corrector
log-likelihood
5.85
MD4 informed corrector
margin
5.78
Table Comparison of log-likelihood vs margin-based confidence scores.
D.3.5 Generated Image Examples
We show some ImageNet samples taken using HollowDiff with informed correctors in figs. to
19

---

MADIA
Figure 6: Example informed corrector generations for the classes Pomeranian (259) and jack-o'-
lantern (607).
20

---

Figure 7: Example informed corrector generations for the classes timber wolf (269) and home theater (598).
21

---

Figure 8: Example informed corrector generations for the classes seashore (978) and bolete (997).
22

---

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

---

Justification:
The assumptions underlying the Markov chain formulation, factorized transi- tion rates, and score function approximations are clearly stated. Proofs or derivations are complete and included in the appendices. Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: The paper describes the experimental setup in sufficient detail for reproducibil- ity. Appendix C outlines the hollow transformer architecture, while Algorithm gives step-by-step breakdown of the informed corrector sampling procedure. Hyperparameters are provided or referenced for the key experiments, enabling reproduction of results relevant to the main claims.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the partieular way they provide for reproducibiliity In the case of closed-source models, it may be that access to the model is limited in

---

some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We plan to release the full codebase, including training scripts, model weights, and instructions for reproducing all main experimental results, upon acceptance. The datasets used (Text8 and ImageNet) are publicly available.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
Senziides/CoeeSubmii www.snnnoniiyyiioooiiiiiiiiiiiil olicy) for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ //nips. cc/publiei
Seriidde//ooeuuui weiaiissssoopooooooooooo oniolicy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The paper describes the experimental settings for all evaluated tasks in suffi- cient detail. Details to the experiments provided in Appendices B-D.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: For the synthetic Markov chain experiments (Figure 2), the paper reports means and standard deviations. For large-scale ImageNet evaluations, we follow common practice in the generative modeling literature and report single-run FID, IS, and precision/recall scores. While we not include error bars for these metrics, this aligns with standard conventions in prior work.
25

---

Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Details on machine types and training time described in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: / hwwrrpsuu...///
inee?
Answer: [Yes]
Justification: We have reviewed the NeurIPS Code of Ethics and affirm that our research adheres to its principles.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
26

---

Answer: [Yes]
Justification: While the paper is primarily foundational nature, we acknowledge the broader context in which improvements to generative models operate.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
The paper does not release any pretrained models or datasets that pose high risk for misuse. The primary contribution is methodological.
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
The paper uses the Text8 dataset and the ImageNet 256x256 dataset, both of which are publicly available and widely used in the community.
Guidelines:
The answer NA means that the paper does not existing assets.
27

---

The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (c.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re- re-packaged, both the original license and the license of the derived asset (if has changed) should provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: he paper does not release any new datasets, pretrained models, or other assets at submission time.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
28

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non non-standard components.
Please refer to our LLM policy http:s://nnuripp Seeeurips.cc/Cooeeeeeeeeee /tt.:p/p....///treeeeee/e/////// for what should or should not be described.
29


Reviews:
[{"review_id": "SD3Fd8cLWL", "reviewer": "Reviewer_E3jH", "summary": "This paper proposes a novel informed corrector for discrete diffusion models, aimed at improving sample quality and efficiency during inference. While discrete diffusion has shown promise across a range of tasks—such as language and image generation—the sampling process remains computationally expensive and prone to errors when using fewer steps. Existing uninformed correctors make naive, often random, corrections without considering model uncertainty, leading to suboptimal performance. To address these issues, the authors introduce: 1. An informed corrector based on adaptive Gibbs sampling, which targets uncertain or likely incorrect dimensions using a confidence-based criterion. 2. A hollow transformer architecture that allows efficient multi-token updates without information leakage. 3. A complementary ELBO-based training objective, named HollowDiff, which leverages both masked and unmasked positions for more efficient learning. Experimental results validate the effectiveness of the proposed method.", "questions": "1. Could the authors provide a direct ablation comparing the two proposed confidence scores (log-likelihood vs. margin-based)? This would clarify their individual contributions.\n2. The corrector is only activated when t < 0.9. How sensitive is this threshold, and would more frequent corrector steps improve performance or introduce instability?\n3. While the hollow transformer enables efficient parallel updates, it underperforms compared to a standard transformer in denoising tasks (Table 2). Additionally, the proposed informed corrector demonstrates greater effectiveness when used with the hollow transformer. Could the authors elaborate on potential architectural improvements to reduce this performance gap while preserving the corrector’s efficiency benefits?", "limitations": "see above", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths:\n1.The adaptive corrector design is a clear advancement over uninformed methods and incorporates techniques (e.g., Gumbel-top-k) not previously applied in this context.\n2. The design of the corrector is both novel and practically effective, with consistent empirical improvements.\n3,The writing is clear and logically structured.\n\nWeakness:\n1. The requirements on separately trained predictor and corrector models for best performance adds complexity to the training pipeline, somewhat diminishing the elegance of the proposed unified approach.\n2. The hollow transformer, while clever, imposes a constraint that makes it incompatible with arbitrary pretrained denoisers. This limits plug-and-play applicability.\n3. Some architectural choices  could benefit from more ablation or sensitivity analysis.", "was_revised": true, "final_justification": "The reviewer addressed my concerns. Considering the contribution and novelty of the paper, I will maintain my original score of slightly positive.", "rebuttal": "We thank the reviewer for the constructive and thoughtful feedback. We are pleased that you found our method novel, effective, and clearly written. Below we respond to your questions and concerns.\n\n### **Ablation: log-likelihood vs. margin-based confidence scores**\n\nWe agree with the reviewer that the ablation on the confidence score is an important result to include. We give the results in the table below. All numbers are reported with 17 NFEs with 8 corrector steps.\n\n| Method | Confidence | FID |\n| --- | --- | --- |\n| HollowDiff + informed corrector | log-likelihood | 6.45 |\n| HollowDiff + informed corrector | margin | **6.26** |\n| MD4 + informed corrector | log-likelihood | 5.85 |\n| MD4 + informed corrector | margin | **5.78** |\n\nIt can be seen that margin-based confidence yields better performance. We will include this table in the appendix for the final version.\n\n### **Sensitivity of the corrector activation time**\n\nThank you for pointing this out. In earlier versions of our implementation, we applied a threshold to delay the start of corrector updates (e.g., t < 0.9) following [8]. However, in our final experiments, we found this unnecessary and removed the threshold entirely. We regret the oversight in leaving this outdated detail in the appendix and will correct it in the final version.\n\n### **Hollow Transformer underperformance and potential improvements**\n\nWe agree with the reviewer’s observation that the hollow transformer underperforms compared to a standard transformer in denoising tasks. We believe that hollow transformers are difficult to train because it is designed to simultaneously learn both the predictor and the corrector via a structural constraint. While there should be a fair amount of weight sharing between the two learning tasks, it is fundamentally a harder learning problem than learning the predictor (denoising) on its own.\n\nOne potential way to speed up training and also address the incompatibility of hollow transformers is by fine-tuning a “hollowed out” pretrained denoising transformer, using its learned weights to initialize the layers of a hollow transformer. We will leave this line of inquiry for future work.\n\n### **On unified vs. separate training**\n\nWe agree that requiring separate predictor and corrector models adds complexity. However, we view this not as a limitation but as flexibility. Our unified HollowDiff model (230M params) already performs competitively (FID 6.26). The hybrid setup (MD4 + hollow corrector) achieves even better performance (FID 5.78), but it uses more parameters (400M total). We do not claim this as a strictly fair comparison to unified approaches — rather, we include it to demonstrate that our informed corrector is broadly compatible and effective even when paired with stronger pretrained predictors. We will clarify this point in the final version."}, {"review_id": "84CclMICzJ", "reviewer": "Reviewer_PFik", "summary": "This paper proposes a novel predictor-corrector sampling scheme for discrete diffusion models to improve the efficiency of inference process. \nThe technique proposed is an informed corrector that leverages the diffusion model itself to identify and resample tokens that are most likely to be erroneous, drawing inspiration from adaptive Gibbs sampling.\nSpecifically, to quantify the model's certainty, it constructs a \nconfidence score for each generated token. \nThis score is defined as the margin between the logit of the current token and that of the next most likely alternative.\nTo efficiently enable the informed corrector, the authors employ a hollow transformer architecture for parallel processing with a single forward pass and modify the training objective that combines complementary loss terms to improve training efficiency.\nIn experiments, the authors conduct synthetic text and image generation tasks, the proposed informed corrector consistently outperforms baselines including predictor-only samplers, the uninformed corrector, and the ReMDM sampler.", "questions": "As shown in Figure 4, the FID score rises sharply when the number of parallel corrector updates $k$, becomes too large (e.g., k ≥ 32). Could the authors provide a more intuitive explanation for why excessive parallel updates harm sample quality.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "### Strength\n- The proposed method is natural and intuitive, introducing a well-motivated informed corrector that directly targets likely errors. \nThis marks a significant improvement over prior uninformed methods, which rely on inefficient random re-masking.\n- The use of the hollow transformer architecture is a natural way to efficiently implement the parallel updates required by the informed corrector.\n\n### Weaknesses\nThe core idea is easy to follow and insightful, several aspects could benefit from further clarification and investigation:\n- The paper compares performance based on the NFEs but does not report wall-clock sampling times and the memory usages. Since a corrector step includes confidence score calculation and resampling, a time-based comparison would offer a more complete picture of the method's practical efficiency.\n- In the Discussion, the authors state that \"better performance is achieved by learning separate predictor and corrector models\". While Table 2 shows the superior result of this informed corrector, could the authors provide a more direct comparison to quantify the performance difference between the jointly trained models and the separately trained models?", "was_revised": true, "final_justification": "The authors have addressed my previous concerns. I am now prone to accept the manuscript.", "rebuttal": "We thank the reviewer for the thoughtful summary and for recognizing the clarity, motivation, and practical significance of our proposed informed corrector and hollow transformer architecture. We respond below to the raised concerns and questions.\n\n### **Wall-clock time comparison**\n\nTo show that a corrector step does not significantly increase the wall-clock time of a sampling step, we compare the sampling time of informed corrector (HollowDiff) with a HollowDiff without corrector, we can see that additional time taken by the corrector step is negligible. However, we did notice that our Hollow Transformer implementation currently has a slightly slower forward pass than a regular Transformer, as shown in the table.  \n\nIn the table, wall clock time in seconds is measured for generating sequences of  256 tokens with a batch size 64 on a machine with 8 NVIDIA H100 SXM GPUs. The first row indicates the number of NFEs.\n\n|  | 16 | 32 | 64 | 128 | 256 |\n| --- | --- | --- | --- | --- | --- |\n| Informed Corrector | 0.1120 ± 0.0003 (k=16) | 0.2098 ± 0.0005 (k=8) | 0.4052 ± 0.0003 (k=4) | 0.7965 ± 0.0004 (k=1) | 1.5775 ± 0.0001 (k=1) |\n| Hollow Transformer (no corrector) | 0.1100 ± 0.0001 | 0.2092 ± 0.0001 | 0.4061 ± 0.0003 | 0.8003 ± 0.0002 | 1.5885 ± 0.0001 |\n| md4 (uninformed corrector) | 0.0831 ± 0.0004 | 0.1343 ± 0.0006 | 0.2349 ± 0.0008 | 0.4260 ± 0.0005 | 0.8242 ± 0.0007 |\n| md4 (ReMDM) | 0.0684 ± 0.0004 | 0.1214 ± 0.0005 | 0.2319 ± 0.0003 | 0.4489 ± 0.0003 | 0.8928 ± 0.0003 |\n| md4 (no corrector) | 0.0618 ± 0.0008 | 0.11160 ± 0.00008 | 0.2127 ± 0.0004 | 0.4154 ± 0.0002 | 0.8198 ± 0.0003 |\n\nWhile the hollow transformer does incur a slowdown, we note that we did not explicitly optimize the computation graph and the architecture for wall-clock time. We are confident that further optimizations will significantly decrease this gap in performance. \n\n### **Joint vs. separate predictor/corrector models**\n\nYou are correct that the best performance in Table 2 comes from the decoupled setup (MD4 predictor + hollow corrector, FID 5.78). However, this configuration uses more total parameters (400M vs. 230M for HollowDiff) and thus is not intended as a strictly fair comparison. Rather, we include it to demonstrate that the informed corrector generalizes well to stronger predictors and is not tied to a particular architecture.This is intended as a proof of concept, and we are excited to explore such hybrid modeling opportunities in future work.  \n\n### **Why excessive parallel updates hurt performance (Figure 4)**\n\nWe appreciate the reviewer’s insightful question. This is indeed an important point to clarify.\n\nThe informed corrector is theoretically grounded in Gibbs sampling, and the correspondence is exact in the limit of single-coordinate updates (k = 1) and many corrector steps. As k increases, we improve the chance of hitting generation errors in fewer steps, but we also move away from the Gibbs framework.\n\nIntuitively, each update assumes the rest of the sequence remains fixed during correction. When too many dimensions are updated in parallel, this assumption no longer holds. The result is a kind of “update collision” — simultaneous updates may become mutually inconsistent, or even undo each other’s corrections. This degrades the effectiveness of the corrector and ultimately harms sample quality. The ablation results on k tells us that there is a good middle ground where some number of parallel updates allow faster generation without hurting performance. We will clarify this point in the revised caption for Figure 4."}, {"review_id": "Zr2igbxaTH", "reviewer": "Reviewer_KG2Y", "summary": "- Discrete diffusion models are an increasingly popular class of generative models for text, proteins, and images. Since the computational cost of each forward pass is high, most methods sample tokens in parallel, introducing errors in generation and reducing sample quality. The authors of this paper propose a predictor-corrector sampling scheme by learning a corrector model.\n- The predictor-corrector scheme is based on adaptive Gibbs sampling which makes use of the corrector model proposed in the paper. Unlike typical Gibbs samplers, which sample one coordinate at a time, the authors make use of parallel updates which are then corrected using their corrector model.\n- The corrector model makes use of a hollow transformer such that the corrector model does not take into account the position whose correctness is being judged, reducing the number of forward calls required to 1.\n- The authors derive an ELBO for training the correct model which can be combined with the denoiser.\n- As baselines, the authors compare against uninformed correctors such ReMDM which mask out text randomly or based on weights computed from previous denoising time-steps.", "questions": "See weaknesses", "limitations": "yes", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Strengths\n\n- The proposed architecture for the corrector model allows efficient inference.\n- The proposed objective leads to faster training as shown in figure 5 in appendix D.2\n- The adaptive Gibbs sampler with parallel updates and then correction leads to faster convergence across various NFEs.\n- For MD4 table 1 does show improved performance with the informed corrector.\n\nWeaknesses\n\n- The experimental evaluations can be improved, for instance, the table 1 compares performance across different parameter counts. \n- Several large-scale masked diffusion language models have been released earlier this year, it would be great to see if a separate corrector model can be scaled to larger parameter counts.", "was_revised": true, "final_justification": "The authors have addressed my concerns and answered my questions. The work done here is original, interesting and timely.", "rebuttal": "We thank the reviewer for the positive and thoughtful assessment. We address the weaknesses identified by the reviewer below.\n\n### **Parameter counts in Table 1**\n\nWe appreciate the reviewer’s observation regarding parameter mismatches in Table 1. Given the breadth of existing work, it is computationally infeasible to reproduce and fairly re-train every baseline under a fixed parameter budget. As such, we chose to focus our evaluation on comparisons with MD4 to better isolate and demonstrate the benefits of the informed corrector. Other entries such as DPC and VQ-Diffusion are included to provide broader context on the landscape of discrete diffusion models.\n\nWe also note that the hybrid MD4 + informed corrector setup uses more total parameters (400M) than the unified HollowDiff model (230M). This comparison is not intended to be strictly fair, but rather to highlight the modularity and compatibility of our corrector: it can be applied to stronger pretrained predictors without need for joint training. We will clarify these distinctions in the final version of the paper.\n\n### **Scaling to larger models**\n\nWe agree that scaling to larger masked diffusion models is a promising and important direction. Our current experiments use relatively lightweight models in order to isolate and rigorously evaluate the contribution of the corrector scheme.\n\nThat said, we are optimistic about the scalability of informed correctors, particularly in the context of modern language models, where output sequences are much longer and generation speed is a key concern. As longer outputs are generated in fewer steps, more tokens are sampled in parallel, increasing the risk of errors. This makes the ability to perform efficient, parallel, and targeted corrections even more valuable. We believe informed correctors are well-positioned to address this need.\n\nWe leave this direction to future work."}, {"review_id": "RvQEkrXKKu", "reviewer": "Reviewer_88LA", "summary": "The paper identifies a key inefficiency in traditional predictor-corrector sampling within mask absorbing diffusion models, specifically that it fails to effectively target problematic tokens. To address this, the authors propose a Gibbs sampling-based mask absorbing diffusion framework. The forward process is defined conditionally for each dimension based on the values of the other dimensions, enabling training through a redefined ELBO term. They introduce the \"hollow transformer\" architecture, specifically designed to handle this training regime. During sampling, the authors employ a Hogwild Gibbs sampling approach, selecting k dimensions based on their proposed confidence heuristic. The authors report improvements in sample quality as a result.", "questions": "Q1: What specific theoretical advantages does the proposed confidence heuristic offer compared to heuristics in existing literature? Although the Gibbs sampling-based mask absorbing diffusion framework is theoretically sound, its comparative advantage over existing frameworks should be also theoretically justified.\n\nQ2: How does the proposed Gibbs sampling approach scale with respect to dimensionality and complexity in realistic datasets?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths**\nThe theoretical justification for the proposed Gibbs sampling-based diffusion framework, which builds upon existing CTMC frameworks, is robust. The idea of selectively correcting problematic dimensions to enhance sampling efficiency is intuitive and appears promising.\n\n**Weaknesses**\n\n1) Theoretical Perspective:  Although the Gibbs sampling-based mask absorbing diffusion framework is theoretically sound, its comparative advantage over existing frameworks remains unclear. The proposed confidence heuristic for informed correction seems lacking theoretical justification, raising doubts about its superiority over heuristics presented in previous methods such as LLADA.\n2) Experimental Perspective: The experiments provided are limited in scope and scale, focusing only on an HMM-based toy experiment, character-level text modeling, and ImageNet (256x256). Important domains mentioned in the introduction, like language modeling and protein generation, lack experimental validation. Even within the character-level text modeling experiments, key metrics such as perplexity, standard in the NLP community, are not reported.", "was_revised": true, "final_justification": "The authors' explanation addresses my questions, so the reviewer will maintain the original score, which is already positive.", "rebuttal": "We thank the reviewer for carefully considering our work and recognizing the theoretical soundness of our Gibbs sampling-based framework, as well as the promise of selectively correcting problematic dimensions. Below, we address the main concerns.\n\n### **Q1: Theoretical basis for the confidence heuristic**\n\nWe thank the reviewer for raising this important question. Although our adaptive Gibbs sampler depends on a confidence metric conceptually similar to those in literature [Nie et al. 2025, Kim et al. 2025], our method is unique and more principled. Other methods, like the one used in LLaDA, alter the predictor's unmasking order based on confidence. This can introduce bias into the sampling process, causing it to deviate from the diffusion model's theoretical reverse process. In contrast, our confidence metric is only used in corrector to define the selection distribution of a random-scan Gibbs sampler, which approximately preserves the marginal distributions of reverse process. Therefore, our sampling process continues to closely approximate the true reverse process of the diffusion model.  \n\nWe will clarify these motivations and connections and add the references in the final version.\n\nNie, Shen, et al. \"Large language diffusion models.\" *arXiv preprint arXiv:2502.09992* (2025).\n\nKim, Jaeyeon, et al. \"Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions.\" *Forty-second International Conference on Machine Learning*.\n\n### **Q2: Scaling to high-dimensional data**\n\nOur informed corrector is specifically designed to scale well in high-dimensional settings. The hollow transformer architecture enables evaluating confidence scores for all dimensions in a single forward pass, and the Gumbel-top-k trick allows us to select and update multiple uncertain positions efficiently in parallel. \n\nFor more discussion on this, refer to our response to Reviewer KG2Y on scaling to masked diffusion based large language models.\n\n### **Experimental scope and evaluation metrics**\n\nWe clarify that our informed corrector method focuses on improving discrete diffusion sampling, an aspect not captured by likelihood-based metrics such as perplexity. For this reason, our evaluation is based on sample quality metrics, namely word error rate (on text8) and FID (on ImageNet).\n\nGiven our computational constraints, we prioritized demonstrating the generality of our method across diverse evaluation settings, including HMMs, character-level text, and images, rather than pushing for scale within each domain. We believe our our method is readily applicable to other discrete domains including protein and large-scale language modeling and we're excited to explore these applications in future work."}]

Author rebuttal:

Meta-review: The submission introduces informed correctors for sampling discrete diffusion models. The approach leverages adaptive Markov chains to enable more efficient and accurate sampling. While it's unclear whether the stationary distribution remains the same for the given adaptive selection and hogwild update of indices, the method seems to perform well in practice.

The reviewers reached a consensus recommending acceptance based on a novel application of Gibbs sampling to diffusion models, improved efficiency, intuitive/timely approach. After considerations of the points I see no reason to recommend differently. Please make sure to include the additional experiments, clarifications and additional discussions/points raised by the reviewers in the revision.
Decision: Accept (poster)