Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and Absorbing State Reformulation

Abstract: Discrete diffusion models are often trained through clean-data prediction, but the prediction can be used in different ways to define the reverse dynamics. In Masked Diffusion Models (MDM) these choices largely coincide, whereas in Uniform Diffusion Models (UDM) they do not. We show that the standard plug-in bridge parameterization for UDM is not optimized by the denoising posterior, but by a leave-one-out posterior that predicts each clean token without using its own noisy observation. This identifies a mismatch between the plug-in ELBO and the usual cross-entropy denoising objective. We characterize the leave-one-out target and derive exact conversions between the denoiser, the leave-one-out posterior, and the score. These conversions allow us to disentangle parameterization and training objective. Our results also lead to inference improvements without any additional training through an informed predictor-corrector sampler and improved temperature sampling based on the leave-one-out predictor. We further introduce an absorbing-state reformulation of uniform diffusion that preserves the UDM joint law while decomposing it into masked-diffusion-like sampling operations, with simpler denoising posteriors, carry-over unmasking, and a natural remasking mechanism. On language modeling, leave-one-out parameterizations consistently improve UDM generation, while the absorbing construction matches or surpasses masked diffusion. These results suggest that the empirical gap between masked and uniform diffusion is driven less by the choice of marginals themselves than by parameterization and sampling design. The code and models can be found at https://github.com/samsongourevitch/rev_udm.

• The paper shows that optimizing the plug-in ELBO in UDMs yields a leave-one-out denoiser rather than the standard denoising posterior.
• It establishes conversion formulas among the denoiser, leave-one-out denoiser, and score, enabling flexible inference and improved sampling.
• The absorbing-state reformulation bridges uniform and masked diffusion models, leading to enhanced generative frontiers and lower perplexity.

Uniform Discrete Diffusion: Leave-One-Out Denoisers and Absorbing State Reformulation

Introduction and Motivation

Uniform Diffusion Models (UDMs) are a prominent alternative to autoregressive generative models for structured discrete data, especially in language modeling. Unlike Masked Diffusion Models (MDMs), UDMs uniformly replace each token in the input with a vocabulary token during the forward (corruption) process, which yields distinct properties both mathematically and algorithmically. This paper rigorously re-examines the foundational choices in UDM parameterization and training objectives, introducing the role of the leave-one-out (LOO) denoiser and presenting a novel absorbing-state reformulation. The work addresses both theoretical and practical concerns, disambiguating how reverse dynamics should be parameterized and optimized and showing how to bridge uniform and masked diffusion using auxiliary variables.

Re-examining Reverse Process Parameterizations

A fundamental aspect of discrete diffusion models is the parameterization of the reverse process, i.e., the transition from a noisy (corrupted) state back toward the clean data distribution. There are several parameterizations:

• Score-based: A network produces a score vector parameterizing a rate matrix.
• Denoiser: The network predicts the posterior $p(x_0|x_t)$, assigning token probabilities for the clean sequence.
• Plug-in (Bridge) Parameterization: A network prediction replaces $x_0$ in the bridge kernel $\fw{s 0, t}{x_0, x_t}$, directly yielding the reverse kernel.

In MDM, these parameterizations are essentially equivalent due to the affine dependence of the bridge on $x_0$. In UDM, nonlinearity in the normalization introduces a discrepancy: optimizing the Evidence Lower Bound (ELBO) with a plug-in parameterization does not recover the standard denoising posterior, but instead produces a leave-one-out posterior, in which the prediction for each coordinate omits its own noisy observation and instead conditions only on the rest of $x_t$. This is a nontrivial divergence from standard intuition.

The Leave-One-Out Denoiser and Its Implications

This work rigorously proves the following assertions for UDMs:

• Plug-in ELBO is minimized by the LOO posterior, not the denoising posterior. This result is unique to UDM due to non-affinity.
• Conversion formulas among the denoiser, LOO denoiser, and score are established, providing an explicit way to interconvert these representations for UDMs.
• Structural invariance: The LOO denoiser prediction for token $\ell$ must be invariant to the value of $x_t^\ell$, a property that can be enforced architecturally (e.g., Hollow Transformers) but poses optimization challenges in language tasks.

The practical implications are substantial: the standard cross-entropy objective corresponds to denoiser-oriented learning, while the plug-in ELBO targets a LOO predictor; using the latter as the primitive consistently improves both training and generative metrics (perplexity and Gen-PPL), as shown in the following empirical results.

Figure 1: Comparison of training and generative performance for denoiser vs. leave-one-out parameterizations; LOO yields lower perplexity and improved sampling frontiers.

Conversion, Inference, and Predictor-Corrector Sampling

Exact conversion between denoiser, LOO denoiser, and score enables:

• Optimization flexibility: A network can be trained with any of these targets and converted at inference time to the required form.
• Enhanced sampling: Applying heuristic top-$p$ or temperature schemes on the LOO predictor rather than the denoiser yields strictly better generative frontiers for a fixed entropy level.

  Figure 2: Top-$p$ Gen-PPL frontiers; LOO-based sampling and post-inference conversion outperform direct denoiser-based sampling.

The LOO representation also yields a theoretically justified predictor-corrector sampler: the LOO conditional yields a Gibbs kernel that preserves the model's marginal at each step. Empirically, this approach outperforms ancestral sampling and can be realized without any auxiliary network, leveraging just the conversion formulas if only a denoiser is available.

Figure 3: Predictor-corrector sampling—both with trained LOO denoiser and a converted denoiser—improves Gen-PPL frontiers across entropies.

Absorbing-State Reformulation and Bridging to Masked Diffusion

A major conceptual contribution of the paper is the absorbing-state reformulation:

• The UDM forward process can be decomposed, conditioned on an auxiliary random variable $U$ (the "absorbing state" for each coordinate), into separate Markov processes for each position with a unique absorbing outcome.
• Conditioning and marginalization recover standard UDM marginals, but the reformulation exposes new architectural and algorithmic flexibility, resembling the "carry-over" property and remasking mechanism of MDMs.
• Joint state augmentation and resampling yields an exact lifted model: the absorbing-state sequence is resampled as part of the reverse chain, recovering the full UDM reverse trajectory.

Figure 4: Left: Sudoku accuracy versus NFE for different methods; Right: Gen-PPL frontiers for AUDM, ReAUDM, UDM, and MDM, indicating the practical benefits of absorbing-state formulations.

Additionally, through a transition-time randomization, the authors show that the UDM denoiser can be exactly mapped to a mask-conditioned MDM denoiser, enabling re-use of existing MDM checkpoints in the UDM context with matched joint law.

Empirical Evaluation

Empirical results substantiate all theoretical claims:

• LOO denoiser as primitive consistently yields lower validation perplexity during training and strictly Pareto-improves generative frontiers across entropy levels on large scale language modeling tasks such as LM1B and OpenWebText.
• AUDM and ReAUDM not only close the gap with MDM in terms of likelihood and generalization but, in some settings, surpass MDM in zero-shot evaluation and generative frontier quality.
• On structured non-text data such as Sudoku, UDMs with LOO parameterization outperform both MDM and absorbing-state models, demonstrating generality.

Theoretical and Practical Implications

• Principled separation of parameterization and training target: Decoupling these provides richer opportunities for model selection and architectural innovations in discrete diffusion.
• LOO parameterization yields a universal representational primitive for inference-time flexibility, better sample quality, and empirically superior likelihoods.
• Absorbing-state constructions rigorously connect uniform and masked diffusions, supporting transfer and hybridization between model classes.

From a practical AI systems perspective, these results call for prioritizing LOO parameterization in UDMs and suggest re-evaluating the previously perceived superiority of masked marginals solely on the basis of their marginal distributions. Rather, generative quality and likelihood appear more sensitive to parameterization, sampling heuristics, and architectural choices revealed by the theoretical analysis.

Conclusion

This paper disentangles the roles of diffusion process design, parameterization, and training objective in uniform discrete diffusion models, establishing that the optimal reverse process parameterization is leave-one-out rather than standard denoising. The derived conversion formulas enable architectural and training flexibility, and the absorbing-state reformulation opens new avenues for model hybridization. Empirically, these insights result in improved likelihoods, generative frontiers, and inference efficiency. Open questions include a theoretical understanding of the superiority of LOO objectives for sampling quality, and further development of parameterizations that leverage absorbing-state constructions for even stronger generative performance.