PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Samson Gourevitch, Yazid Janati, Dario Shariatian, Umut Simsekli, Eric Moulines, Eric P. Xing, Alain Durmus

arXiv:2605.22765v1cs.LGstat.ML

TL;DR

The paper addresses a mismatch between denoising objectives and reverse-kernel parameterizations in Uniform Diffusion Models. It derives leave-one-out conversions and absorbing-state reformulations, finding improved UDM generation and masked-diffusion-comparable absorbing models. The authors conclude that parameterization and inference design matter more than corruption marginals for the observed practical gap.

  • Problem

    UDM plug-in bridge parameterization is not optimized by the standard denoising posterior, creating a mismatch with the usual cross-entropy objective.

  • Method

    The paper derives conversions among denoisers, leave-one-out predictors, and scores, and constructs absorbing-state UDMs that preserve the UDM joint law.

  • Results

    Leave-one-out parameterizations consistently improve UDM generative performance, while AUDM and ReAUDM match or surpass masked diffusion on several evaluations.

  • Takeaways & Limitations

    Parameterization and sampling design can transfer masked-diffusion-like structure and performance advantages to models with uniform-diffusion marginals.

  • Takeaways & Limitations

    Evidence for the leave-one-out advantage remains mainly empirical, and the best generative frontiers are still obtained by leave-one-out UDMs.

Abstract

from arXiv · show

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.

1 Introduction

The paper identifies a parameterization mismatch in Uniform Diffusion Models and introduces leave-one-out predictors and absorbing-state constructions to improve UDM inference and generation.

  • UDM plug-in bridge parameterization is optimized by a leave-one-out posterior rather than the standard denoising posterior.The leave-one-out predictor estimates each clean token using noisy tokens at all other positions.
  • The plug-in ELBO and usual cross-entropy objective target different posteriors in UDMs.The ELBO targets the leave-one-out posterior, while cross-entropy targets the plain denoising posterior.
  • Exact conversions between denoisers, leave-one-out predictors, and scores separate reverse-kernel parameterization from training objective.These conversions support training either parameterization with either objective.
  • Leave-one-out parameterizations consistently improve UDM generative performance across objectives and samplers.The improvement is reported across large-scale language modeling and a smaller Sudoku task.
  • Absorbing-state constructions preserve the UDM joint law while enabling masked-diffusion-like sampling operations.The constructions support simpler denoising posteriors, carry-over unmasking, and remasking mechanisms.
  • The paper suggests the practical gap between masked and uniform diffusion depends less on marginals than on parameterization and sampling design.The absorbing construction is reported as on par or slightly better than masked diffusion models.

2 Preliminaries on discrete diffusion models

Discrete diffusion models define a forward corruption process from data to a reference distribution and reverse it using learned approximations to conditional transitions.

  • Discrete diffusion transports a data distribution p0 through successive Markov transitions toward a reference distribution p1.The reference distribution factorizes into categorical distributions across token positions.
  • A monotone noise schedule controls the transition ratios, with α0 = 1 and α1 approximately 0.The marginal distribution at time t is denoted pt.
  • The generative model reverses the Markov process using transitions conditioned on the denoising posterior and bridge distribution.The bridge is the conditional distribution of an earlier state given the clean and noisy states.
  • Practical reverse transitions are approximated by a parameterized joint law and trained by minimizing an expected NELBO.For fixed data, the NELBO is defined up to a parameter-independent additive constant.
  • Masked and uniform diffusion differ in their reference distributions: MDM uses a mask token, whereas UDM uses the uniform distribution over tokens.The MDM bridge is restricted by the mask state, while the UDM bridge has an explicit form.

3 Elucidating the parameterizations in discrete diffusion models

The paper distinguishes reverse-transition parameterizations that coincide for masked diffusion but diverge for uniform diffusion, where the bridge is nonlinear in the clean-token prediction. It shows that bridge plug-in optimization targets a leave-one-out posterior, establishes exact conversions among representations, and uses them for improved training and sampling.

  • Parameterization choices: Reverse transitions can be parameterized by directly predicting transitions, marginalizing a clean-data prediction, or inserting that prediction into the bridge.The marginalization and bridge plug-in forms use a simplex-valued neural prediction, but apply it differently.
  • Parameterization choices: For masked diffusion, marginalization and bridge plug-in coincide because the bridge is token-wise affine; for uniform diffusion, a normalization term makes the bridge nonlinear and breaks this equivalence.This leaves the optimal object for uniform bridge plug-in parameterization non-obvious under the expected NELBO.
  • Leave-one-out parameterization: The bridge plug-in parameterization minimizes the expected ELBO at the leave-one-out posterior rather than the ordinary denoising posterior, with a unique minimizer for UDM.The leave-one-out predictor estimates each clean token using noisy observations at the other positions, while the denoising posterior uses the full noisy input.
  • Leave-one-out parameterization: For UDM, exact conversions among the denoising posterior, leave-one-out posterior, and score disentangle parameterization from training objective and remain available because the forward process has full support.The inverse conversion fails for masked diffusion on unmasked positions because the denoiser lacks enough information.
  • Practical consequences: Leave-one-out predictions support standard cross-entropy training and an informed predictor-corrector sampler without an auxiliary model; empirically, the resulting loss outperforms plain-denoiser cross-entropy.The corrector uses Gibbs updates that preserve the noisy marginal distribution, and trained or converted denoisers can supply the required leave-one-out posterior.

4 Bridging Masked and Uniform Diffusion

The paper gives two exact liftings that connect uniform diffusion to absorbing- or masked-diffusion-like processes while preserving the UDM law. These constructions simplify reverse sampling and clarify when denoising behavior remains time-dependent or requires remasking.

  • 4.1 Absorbing State Uniform Diffusion: Conditioning uniform diffusion on an independently sampled absorbing token at each position yields an absorbing-state process whose marginal law remains the original UDM law.The absorbing state is a random token rather than a distinguished mask token.
  • 4.1 Absorbing State Uniform Diffusion: The conditional bridge is token-wise, providing an analogue of masked-diffusion carry-over structure for the absorbing-state representation.The conditional process fixes the absorbing sequence U throughout its trajectory.
  • 4.1 Absorbing State Uniform Diffusion: Unlike masked diffusion, the noise-conditioned denoiser remains explicitly time-dependent and can require prediction from the full noisy sequence.The single-token case illustrates this remaining ambiguity when the observed token equals its absorbing value.
  • 4.2 Remasked AUDM: Resampling the absorbing sequence after each reverse step decomposes uniform diffusion into denoiser, bridge, and absorbing-state resampling operations.The resulting trajectory exactly matches the uniform UDM reverse-chain law.
  • 4.3 Masked Uniform Diffusion: Conditioning on latent transition times makes the UDM denoising problem precisely a masked-diffusion denoising problem and permits reuse of an MDM denoiser.Marginalizing the transition times recovers the UDM forward transition, while the lifted reverse chain can recover the full UDM law.
  • 4.3 Masked Uniform Diffusion: The transition-time construction formalizes a remasking mechanism for UDM that is absent from standard MDM.The full UDM joint distribution is recovered by resampling transition times after each reverse step.

5 Related Works

Prior work studies leave-one-out representations, predictor-corrector sampling, and augmented-state formulations of uniform diffusion. The paper positions its contribution as showing that standard UDM bridge training already targets a leave-one-out object and as proposing related sampling constructions.

  • Leave-one-out denoiser: Categorical ratio matching learns a concrete score, and its conditional distribution is leave-one-out; the Hollow Transformer is essential for recovering that target with the associated loss.Without the architectural constraint, the model can learn the identity function instead.
  • Empirical comparison: Figure 1 compares denoiser and leave-one-out parameterizations using validation perplexity and top-p sampling Gen-PPL frontiers.The top-p sweep covers p ∈[0.8, 1.0], with curves shown through entropy 5.6 and a vertical dataset-entropy reference.
  • Leave-one-out denoiser: The paper distinguishes its result from prior leave-one-out work by showing that standard UDM bridge parameterization with ELBO already targets the same leave-one-out object.It also proposes a simpler cross-entropy objective for this target.
  • Predictor-corrector samplers: Predictor-corrector methods alternate reverse transitions with marginal-preserving corrector kernels, including uninformed birth-death and informed Gibbs correctors.Remasking in masked diffusion can be interpreted as predictor steps interleaved with a single corrector step.
  • UDMs with augmented states: Related UDM augmentations use routing variables, corruption masks with planners, or latent transition times for different sampling or factorization purposes.These approaches augment the state space but pursue goals distinct from the constructions developed here.

6 Experiments

Experiments isolate the effects of UDM parameterization, training objective, sampling, and absorbing-state process design across language modeling and Sudoku. Leave-one-out methods improve UDM generation, while AUDM remains competitive with masked diffusion.

  • LOO denoiser vs denoiser: Matched continuous-time training shows that leave-one-out parameterization improves optimization and generative frontiers over the denoiser counterpart.The comparison covers both cross-entropy and ELBO training under matched architectures.
  • LOO denoiser vs denoiser: Leave-one-out cross-entropy is competitive with ELBO training and avoids the weakness of cross-entropy with the denoiser parameterization.The experiments support using the LOO denoiser as the default UDM prediction target, especially with cross-entropy.
  • Sampling: Converting a denoiser-trained model to the leave-one-out representation before top-p sampling improves the frontier at no additional training cost.The converted model closes a substantial part of the gap to a model trained directly with the leave-one-out target.
  • Sampling: Predictor-corrector sampling Pareto-dominates ancestral temperature sampling for both directly trained and converted leave-one-out denoisers.The sampler uses no auxiliary model and has the same runtime as regular ancestral sampling in the reported implementation.
  • Absorbing-state uniform diffusion: AUDM adds a 4.1% parameter overhead, improves likelihood over UDM, and is often slightly better than MDM along the generative frontier.AUDM is competitive across NFEs, while MDM remains slightly better on OWT validation perplexity.
  • Absorbing-state uniform diffusion: On Sudoku, leave-one-out UDMs achieve nearly perfect solve rates with both cross-entropy and ELBO, while AUDM and ReAUDM improve over MDM and AUDM.ReAUDM improves over AUDM without additional sampling cost or a learned component.

7 Conclusion

The paper separates UDM reverse-kernel parameterization from forward-process design, showing that leave-one-out targets and absorbing-state reformulations improve practical performance. It concludes that the masked–uniform gap is more associated with learned representation and inference than corruption marginals, while identifying open limitations.

  • Parameterization: The UDM bridge plug-in objective is optimized by a leave-one-out posterior rather than the standard denoising posterior.Exact conversions among denoiser, leave-one-out predictor, and score separate training targets from sampling parameterizations.
  • Parameterization: The leave-one-out viewpoint improves UDM optimization and generative frontiers, enables effective top-p or temperature application, and supplies predictor-corrector sampling without an auxiliary network.These benefits follow from the conversion formulas and the resulting representation choices.
  • Process design: Absorbing-state and resampled reformulations transfer masked-diffusion-like denoisers, carry-over structure, and remasking to uniform diffusion while preserving its joint law.The resampled version recovers the full UDM reverse chain, and masked-diffusion denoisers can be reused by conditioning on transition time.
  • Conclusion: Leave-one-out parameterizations consistently improve UDM performance, while AUDM and ReAUDM match or surpass masked diffusion on several frontier evaluations and improve likelihood within the UDM family.The reported results support a practical gap driven less by corruption marginals than by learned representation and inference scheme.
  • Limitations and future work: The advantage of the leave-one-out target remains mainly empirical, and the best generative frontiers still come from leave-one-out UDMs rather than AUDM or ReAUDM.Improving absorbing-state models through better parameterizations or sampling schemes remains future work.

A.1 Proof of Proposition 1

Under the simplex-extension assumption, the Bridge Plug-in ELBO is minimized by the leave-one-out posterior, uniquely for UDM. The proof establishes conversion identities and highlights support and injectivity conditions underlying the result.

  • Assumptions: The result depends on extending the one-hot bridge to simplex-valued inputs in a Bayes-compatible way.The canonical extension satisfies this condition, whereas arbitrary extensions need not preserve the proposition.
  • Proposition 1: Under (A1), the Bridge Plug-in minimizer equals the leave-one-out posterior on the support of the forward process.The proof reduces the objective to token-wise KL divergences and identifies the leave-one-out posterior as a minimizer.
  • Uniqueness: For UDM, this minimizer is unique because the parameter-to-transition map is injective under the stated conditions.The argument establishes strict monotonicity for the relevant ratio, which determines the simplex-valued prediction uniquely.
  • Assumptions: Naive simplex extensions can violate (A1), so the proposition may fail when the extension does not reproduce exact marginal reverse transitions.The appendix gives a counterexample in which positive leave-one-out mass is not mapped to the correct reverse marginal.
  • Conversion: The denoiser is obtained from the leave-one-out posterior by reinserting the local likelihood term, enabling score conversions as well.The score can be parameterized from either the denoiser or the leave-one-out denoiser through the derived identities.
  • Conversion: Denoiser-to-leave-one-out conversion is exact for UDM because its forward process has full support, but fails on unmasked MDM positions.On those MDM positions, the denoiser lacks enough information to reconstruct the leave-one-out posterior.
  • Empirical illustration: In the Sudoku visualization, the denoiser is more confident than the leave-one-out posterior near digit edges, with earlier-time blurring from marginal convolution.The difference varies with time and is observed earlier in sampling for larger t.

B.2 Resampled AUDM

Resampled AUDM augments uniform diffusion with an absorbing state and uses tractable resampling to preserve the UDM marginal path law. Its practical approximation factorizes the exact one-position posterior while retaining the exact bridge and resampling law.

  • Construction: AUDM augments each state with an absorbing variable and resamples that variable after drawing the clean and intermediate states.This avoids the intractable direct draw from the absorbing-state conditional distribution.
  • Exact lifted chain: The ideal lifted chain preserves the UDM marginal path law, but directly sampling its absorbing variable is intractable.The propagation identity holds one step at a time and extends across the time grid.
  • ReAUDM: ReAUDM replaces the exact one-position posterior marginal with a factorized approximation while keeping the UDM bridge and resampling law exact.The resulting token-wise reverse transition factorizes across positions.
  • Optimality: The factorized approximation is optimal when its denoiser matches the exact noise-conditioned posterior marginal at every position.The associated NELBO decomposes into backward-transition KL divergences minimized by matching these marginals.
  • Relation to planned denoising: AUDM conditions on the full absorbing vector and uses one denoiser network, unlike planned denoising’s planner-denoiser split.The construction recovers the UDM process through resampling without training a planner.

C Masked Uniform Diffusion

MUDM represents uniform diffusion through a latent transition-time variable and uses resampled lifted sampling to recover the UDM reverse-chain law. Its mask-conditioned posterior has carry-over unmasking and is time-independent.

  • Latent transition-time view: MUDM decomposes the exact UDM denoising posterior into a mixture over latent transition times.This provides a masked-diffusion view of uniform diffusion, although the posterior over transition times is not directly tractable.
  • Sampler: The resampled MUDM sampler initializes the lifted state with a sample from the terminal marginal and its conditional transition-time distribution.It then iterates reverse updates across the prescribed time grid.
  • Tractable resampling: The lifted transition conditions on the clean and intermediate states to make the resampling step tractable.The transition-time law reweights whether a position has transitioned or has remained unchanged by chance.
  • Path-law preservation: Algorithm 4 generates a trajectory with the same law as the UDM reverse chain.The equality follows from a one-step identity that iterates over the time grid.

D On the role of the bridge extension in the plug-in parameterization

The plug-in target depends on how the bridge is extended from one-hot clean tokens to simplex-valued predictions. Canonical nonlinear extension targets the leave-one-out posterior, while barycentric affine extension targets the denoiser.

  • Bridge extension: The Bridge Plug-in parameterization is undefined until the token-wise bridge is extended from one-hot inputs to the simplex.Different extensions agree on exact one-hot bridges but can induce different plug-in reverse dynamics.
  • Canonical extension: The canonical Bayes-ratio extension is nonlinear in the predicted probability vector and has the leave-one-out posterior as its ELBO minimizer.The normalization denominator creates the distinction from denoising-posterior optimization in UDM.
  • Barycentric extension: The barycentric extension averages exact bridges over clean tokens, is affine in the simplex argument, and returns the denoising posterior as the plug-in optimum.In MDM, this extension coincides with the canonical one because the bridge is already affine.
  • Interpretation: The difference between denoiser and leave-one-out targets comes from the bridge extension rather than from plug-in parameterization alone.A denoiser can be converted to the leave-one-out representation before insertion into the canonical bridge.
  • Objective duality: The linear-bridge marginalization ELBO and converted canonical plug-in ELBO are the same generalized-KL objective in denoiser coordinates.This links the two objectives through the exact UDM conversion formulas.
  • Sampling consequence: Informed Gibbs correction uses the leave-one-out object already learned by Bridge Plug-in, requiring no auxiliary leave-one-out model.The corrector focuses updates on positions with small confidence margins, and discrete token changes can substantially alter later reverse trajectories.

G.3 Maximum coupling bridge and ELBO for uniform diffusion

The maximal-coupling bridge preserves the uniform-diffusion forward marginals while reaching the maximal-coupling bound. Its reverse generator yields a corresponding continuous-time ELBO parameterized by a denoiser.

  • The token-wise bridge is a maximal coupling of forward marginals that preserves the original one-time marginals.The equality probability reaches the maximal-coupling upper bound.
  • The bridge factorizes across positions, so nonzero off-diagonal transitions change only one token at a time.The reverse-rate derivation considers states differing at one position and isolates the corresponding clean-token contributions.
  • Averaging the conditioned bridge against the denoising posterior gives the exact reverse generator.The diagonal term is determined by the row-sum-zero condition.
  • The maximal-coupling reverse generator can be parameterized by a denoiser model and associated with a continuous-time ELBO.Proposition 11 derives the ELBO up to a parameter-independent constant.

H Leave-one-out and Hollow Transformers

Leave-one-out predictions should be invariant to each position’s local noisy observation, but standard transformers do not automatically achieve this in high-dimensional settings. Hollow Transformers enforce the invariance structurally, at the cost of harder optimization and restricted information flow.

  • Leave-one-out invariance: The leave-one-out posterior at position ℓ is independent of the local observation xℓ_t, making local sensitivity a diagnostic for suboptimality.Residual dependence after conversion certifies that the model has not reached the leave-one-out optimum.
  • Sensitivity diagnostic: On MNIST, converted denoisers and directly trained leave-one-out predictors are nearly insensitive to xℓ_t, suggesting the conversion is effectively exact.Their sensitivity curves are almost identical and both models appear close to their respective optima.
  • Sensitivity diagnostic: On OWT, sensitivity remains nonzero, although it is lower for directly trained leave-one-out predictors than for converted denoisers.The diagnostic is aligned with downstream performance: models closer to the invariance tend to attain better optima.
  • Hollow Transformer: Hollow Transformers prevent each output from attending to its same-position input while retaining access to surrounding context.The architecture combines left-to-right and right-to-left autoregressive streams.
  • Hollow Transformer: Cross-group interactions in Hollow Transformers require deeper layers, weakening direct attention and making prediction harder.In LM1B experiments, this produced worse optimization and final performance than corresponding standard DiT models, so the architecture was excluded from main experiments.

I.2 Additional results

Additional experiments confirm that leave-one-out parameterization and absorbing-uniform designs improve or preserve sampling quality across inference settings. Predictor-corrector sampling becomes especially advantageous at higher NFE budgets.

  • Temperature and top-p sampling: Leave-one-out parameterization remains preferable in the low-entropy regime across temperature and top-p frontiers.The maximal-coupling variants show a smaller but similar parameterization trend.
  • Temperature and top-p sampling: Converting a denoiser to the leave-one-out representation before nucleus sampling improves the frontier without additional training.The converted model closes a substantial part of the gap to a directly trained leave-one-out model and remains better than applying top-p to the denoiser.
  • Predictor-corrector sampling: At 128 and 256 NFEs, predictor-corrector and filtered top-p sampling are broadly competitive.At 512 and 1024 NFEs, several predictor-corrector configurations move to a strictly better frontier region.
  • Performance comparison with AUDM: AUDM and resampled AUDM track each other closely and remain competitive throughout the NFE range.They outperform MDM and denoiser-parameterization variants in the reported frontier comparison.
Loading 2605.22765v1…