PaperlayerView exact revision on arXiv

Source-linked AI summary

Provably adaptive sampling with uniform and remasking discrete diffusion models

Daniil Dmitriev, Zhihan Huang, Yuting Wei

arXiv:2608.23554v1cs.LGcs.ITmath.STstat.ML

TL;DR

The paper asks whether dimension-dependent sampling costs for uniform discrete diffusion are intrinsic to the forward process and develops adaptive samplers for uniform and remasking processes. Its guarantees instead control sampling complexity through target-distribution dependence, while experiments illustrate the predicted dimension-adaptive behavior.

  • Problem

    Existing theory leaves open whether the linear ambient-dimension dependence for uniform diffusion reflects the forward process or the sampling algorithm, while remasking lacks comparable sampling-efficiency guarantees.

  • Method

    The paper analyzes a first-order leave-one-out sampler for uniform and remasking processes, using a Bayes-optimal auxiliary sampler to separate discretization and score-estimation errors.

  • Results

    Sampling guarantees are controlled by the target distribution’s dual total correlation rather than explicitly by ambient dimension, for both uniform and remasking processes.

  • Takeaways & Limitations

    Sampling efficiency depends critically on how reverse dynamics are approximated and discretized, not only on the choice of forward process.

  • Takeaways & Limitations

    Explaining why uniform and remasking processes outperform masking in the numerical experiments remains an open direction.

Abstract

from arXiv · show

Discrete diffusion models offer a promising alternative to autoregressive generation by enabling parallel updates, but their sampling efficiency can depend strongly on the choice of the forward process and the sampler. For the uniform forward process, existing lower bounds for the standard $τ$-leaping sampler scale linearly with the ambient dimension $d$, raising the question of whether this dependence is intrinsic to the forward process. We answer this question in the negative. We consider a first-order sampler based on the leave-one-out denoiser for uniform and remasking processes whose coordinate updates can be performed in parallel. In both cases, the sampler can correct denoising mistakes during the sampling process, which becomes necessary when many coordinates are updated together. Our main result establishes an adaptive sampling guarantee: up to logarithmic factors, $N = O(\mathrm{DTC}(X_0) / \varepsilon)$ discretization steps suffice to achieve sampling error $O(\varepsilon_{\mathrm{score}}+\varepsilon)$, where $\varepsilon_{\mathrm{score}}$ is the error in score estimation. Thus, the sampling complexity is governed by the intrinsic dependence structure of the target distribution, as measured by its dual total correlation $\mathrm{DTC}(X_0)$, rather than directly by the ambient dimension $d$. Our analysis proceeds through a Bayes-optimal auxiliary sampler that separates discretization error from score-estimation error. We also derive an exact information-theoretic representation of the discretization error in terms of the mutual information between different coordinates of the forward process at different times. This representation applies to general forward processes and, in the uniform and remasking cases, can be controlled by $\mathrm{DTC}(X_0)$. Numerical experiments on structured synthetic distributions illustrate the predicted dimension-adaptive behavior.

1 Introduction

The paper asks whether unfavorable dimension dependence in uniform discrete diffusion is intrinsic to the forward process or caused by the sampler. It introduces a leave-one-out first-order sampler for uniform and remasking processes and derives adaptive guarantees governed by target dependence.

  • Remasking and uniform diffusion can revise intermediate token decisions, unlike monotone masking diffusion, which cannot revisit unmasked positions.
  • The paper investigates whether uniform diffusion’s unfavorable dimension dependence is intrinsic to the forward process or merely a consequence of τ-leaping.
  • The proposed leave-one-out sampler uses independent one-dimensional CTMC updates that can be executed for all coordinates in parallel.
  • Up to logarithmic factors, N = O(DTC(X_0)/ε) discretization steps suffice for sampling error O(ε_score + ε).
  • A Bayes-optimal auxiliary sampler separates discretization error from score-estimation error and represents discretization error through mutual information across forward-process coordinates and times.
  • The leave-one-out denoiser is used both to construct the sampler and to analyze its approximation and discretization errors.

2 Problem setup

The setup models discrete diffusion as coordinate-wise CTMC forward and reverse processes on a discrete state space. In practice, learned scores at discrete times are converted into approximate reverse rates, with leave-one-out denoisers providing an equivalent representation.

  • Discrete diffusion models use forward and reverse continuous-time Markov chains on a discrete state space.
  • The forward process changes at most one coordinate at a time and acts on coordinates independently.
  • The uniform process transitions within the vocabulary, while remasking adds MASK and REMASK states; masking is the special case pM = 1.
  • Choosing pM < 1 lets the reverse process pass through REMASK and correct previously unmasked coordinates.
  • The exact reverse process requires score functions over all times, so discrete-time score estimates define approximate reverse rate matrices.
  • Score-estimation error is controlled through a discretization-weighted score entropy loss bounded by ε_score.
  • The leave-one-out formulation estimates conditional probabilities given all other coordinates, yielding a cleaner reverse-dynamics characterization and facilitating sampling-error control.
  • For uniform diffusion, the conditional probability combines the original leave-one-out probability with the uniform distribution through e^-t and 1 - e^-t.

3 Discrete diffusion sampling with leave-one-out sampler

The leave-one-out sampler approximates reverse dynamics using conditional probabilities that exclude the updated coordinate, enabling parallel coordinate-wise CTMC simulation for uniform and remasking processes. A Bayes-optimal auxiliary sampler separates discretization error from learned-denoiser error.

  • Comparison with τ-leaping: Unlike τ-leaping, the proposed sampler and its Bayes-optimal counterpart account for the time index within each interval and avoid relying only on the interval’s initial context.The paper also notes that τ-leaping can have suboptimal sampling complexity for the uniform process.
  • Leave-one-out sampler: For each discretization interval, the approximate reverse CTMC decomposes into d independent one-dimensional CTMCs, so coordinate updates can run in parallel.The transition rates do not depend on the evolving coordinates outside the updated coordinate.
  • Leave-one-out sampler: Algorithm 1 initializes from p0, advances through N discretization intervals, and updates all coordinates in parallel using leave-one-out denoisers.Its output is a sample in V^d, with process-specific transition expressions for uniform, masking, and remasking cases.
  • Leave-one-out sampler: Proposition 1 shows that Algorithm 1 simulates the full CTMC dynamics associated with its approximate reverse rate matrices.The proposition identifies each interval update with the corresponding CTMC transition.
  • Bayes-optimal sampler: The Bayes-optimal sampler conditions on information available at the interval start, whereas the practical sampler further substitutes estimated conditional probabilities for exact ones.This construction isolates time-discretization effects from score-estimation effects.
  • Bayes-optimal sampler: The leave-one-out sampler replaces full-context reverse scores with conditional probabilities based on the other coordinates at the discretization interval’s start.The practical sampler additionally replaces population conditional probabilities with learned estimates.

4 Main results

The main results show that leave-one-out sampling achieves complexity controlled by the target’s dual total correlation rather than ambient dimension for uniform and remasking processes. The analysis represents discretization error information-theoretically and separates it from score-estimation error.

  • Main sampling guarantee: N = rO(DTC(X0)/ε) discretization steps suffice for sampling error O(ε_score + ε) under the stated discretization and initialization conditions.The guarantee applies to both uniform and remasking processes.
  • Main sampling guarantee: The resulting step complexity scales with DTC(X0) rather than explicitly with ambient dimension d, benefiting structured high-dimensional targets with DTC(X0) ≪ d.The error decomposition includes initialization, score-estimation, and intrinsic discretization components.
  • Main sampling guarantee: For uniform diffusion, unfavorable dimension dependence is not intrinsic to the forward process but can arise from the standard τ-leaping sampler.The same analysis supplies an adaptive sampling guarantee for remasking diffusion, where comparable theoretical guarantees had been lacking.
  • Error decomposition: The sampling-error decomposition separates sampler-independent discretization error from sampler-specific approximation error, allowing the discretization analysis to extend beyond Algorithm 1.The approximation term measures discrepancy from the Bayes-optimal sampler.
  • Discretization error: The exact discretization-error representation depends on mutual information between one coordinate and the remaining coordinates at different times, and applies to general forward processes.It does not explicitly depend on ambient dimension or vocabulary size, but instead on information-theoretic properties of the data distribution and forward process.
  • Discretization error: For uniform and remasking processes, telescoping bounds control the discretization error by κDTC(X0) under a geometric discretization grid.Early stopping near time zero is required to control approximation error rather than discretization error.

5 Numerical examples

Synthetic experiments illustrate the predicted adaptive sampling behavior across dimension, discretization, and target structure. Uniform and remasking often outperform masking, while the proposed uniform sampler performs best among the compared uniform-process samplers.

  • The experiments vary ambient dimension, time discretization, and target structural complexity to test adaptive sampling behavior.
  • Uniform and remasking processes consistently outperform masking because they can correct early denoising mistakes.
  • Geometric grids benefit uniform and remasking processes, whereas masking performs much better on a constant grid.
  • The proposed sampler performs best among τ-leaping, truncated τ-leaping, and proposed samplers for the uniform process.The competing samplers incur additional errors from approximation of the Bayes-optimal sampler.
  • For sparse mixtures, fixed N = 20 yields KL divergence that grows logarithmically with k, consistent with Theorem 1.Uniform and remasking have smaller sampling errors than masking in this setting.
  • The mixture experiments estimate KL divergence by combining all generated samples outside the k binary strings into one bin, producing only a lower bound on true KL divergence.The approximation is reported as accurate in low dimensions and scalable to higher dimensions.

6 Discussion

The discussion argues that adaptive guarantees depend on target structure and reverse-process approximation rather than ambient dimension alone. It identifies extensions and unresolved theoretical questions for future work.

  • Adaptive guarantees for uniform and remasking processes show that uniform diffusion’s linear dimension dependence is not intrinsic to its forward process.The leave-one-out sampler instead yields guarantees controlled by an information-theoretic measure of the target distribution.
  • Sampling efficiency depends critically on how reverse dynamics are approximated and discretized, not only on the forward process.
  • The analysis requires an unstructured forward process whose post-jump value is independent of its initial value.
  • Future work: A theoretical explanation for the experiments’ consistent advantage of uniform and remasking over masking remains open.
  • Future work: Extending the techniques to discrete Gaussian or semantically dependent forward processes could clarify connections between discrete and continuous diffusion models.
  • Future work: Higher-order samplers may further improve dependence on accuracy or intrinsic complexity.
  • Future work: Characterizing score-estimation sample complexity remains important for a unified theory of discrete-diffusion sampling.

B Technical preparations

The technical preparations define information-theoretic quantities and establish forward-process identities used to analyze uniform and remasking diffusion. Key bounds relate conditional transition behavior to a common function F(t,b).

  • The section recalls total correlation and dual total correlation for quantifying dependence in a random vector.
  • F(t,b) represents probability mass moved to b after a jump, including the uniform process’s possible self-loop jump.
  • For the uniform process, F(t,b) = 1/S, while remasking has an explicit process-dependent expression.
  • For both processes, ν(t,b) is at least a constant fraction of F(t,b), with ν(t,b)/F(t,b) ≥ 1/2.
  • Conditional on an initial coordinate value or context, the relevant transition probabilities equal F(t,b) for uniform and remasking processes.

C.1 Proof of Theorem 1

The proof decomposes output error into noise, discretization, and score-estimation terms, then controls the discretization contribution through changes in dual total correlation. This yields the theorem’s adaptive step-count bound.

  • The auxiliary continuous-time Markov chain has the same distribution as Algorithm 1, enabling analysis of its sampling error.
  • The proof combines noise, discretization, and score-estimation bounds to establish the final sampling guarantee.
  • The discretization error is bounded by time-step factors multiplied by changes in DTC across adjacent forward-process times.
  • The score-estimation contribution is bounded by ε_score under the stated score-estimation assumption.
  • The forward-process noise error admits a dimension-dependent bound e^-T d log S.

C.2 Proof of Proposition 1

The proof reduces the parallel sampler to independent one-dimensional reverse processes on each discretization interval. Each coordinate update matches its corresponding continuous-time Markov chain transition, so the full transition kernel agrees with the approximate CTMC.

  • Coordinate decomposition: On each interval, the approximate CTMC decomposes into d independent one-dimensional CTMCs because each coordinate rate depends on the frozen context, not evolving other coordinates.This reduces verification of the parallel update to checking each coordinate separately.
  • Auxiliary process: An auxiliary one-dimensional forward process is introduced for each coordinate and its conditional marginal is identified with the relevant process marginal.The auxiliary process provides the object whose reverse dynamics are compared with the sampler.
  • Reverse-process identification: The score used in Eqn. (10) is exactly the score of the auxiliary one-dimensional process, and time reversal identifies its reverse generator with the coordinate restriction of the full reverse process.This connects the denoiser-based update to the reverse CTMC dynamics.
  • Coordinate update: The update of coordinate i in Algorithm 1 has the same law as the corresponding coordinate of the CTMC generated by Qt.The proof uses the observed endpoint initialization of the reverse auxiliary process.
  • Kernel equality: Because coordinates are sampled independently in parallel, Algorithm 1 and the CTMC have the same transition kernel on every interval, and hence over the full discretization.The interval-wise equality is applied successively across all discretization intervals.

C.3 Proof of Proposition 2

This proof decomposes the relevant divergence using a scalar identity and conditional expectations, then applies Girsanov’s theorem to bound the output divergence by an auxiliary noise-process term.

  • Divergence decomposition: The scalar identity γD(α,γ) = βD(α,β) + γD(β,γ) + (β − γ) log α decomposes the divergence into two comparison terms and a correction term.The identity is applied pointwise to transition-rate quantities.
  • Conditional averaging: The transition-rate terms depend on the current coordinate and context variables, allowing the proof to use the law of total expectation over the omitted coordinates.The conditioning structure is used to compare the original and auxiliary rates.
  • Error separation: The resulting decomposition separates the reverse-process divergence from the discrepancy between the true and auxiliary transition rates.The proof collects these terms before invoking a change of measure.
  • Final bound: Girsanov’s change-of-measure theorem yields KL(qT−tN || poutput) ≤ KL(qT || qnoise) + … .The displayed bound controls the output divergence through a noise-process divergence and the remaining derived term.

C.4.1 Proof of Theorem 2

The theorem proof represents the discretization error through transition-rate divergences and controls it for uniform and remasking processes using conditional information quantities and DTC-based bounds.

  • DTC control: The resulting integrated bound is related to dual total correlation through the identity involving d/dv DTC(Xv).This links the accumulated discretization error to the dependence structure of the forward process.
  • Process-specific control: Proposition 6 bounds the conditional rate-divergence contributions for the uniform and remasking forward processes.The proof combines the proposition with auxiliary inequalities and a Grönwall-type argument.
  • Step-size condition: Under the discretization condition v − u ≤ κ min(1,u), the proof controls the time-local approximation terms needed for the theorem.This condition appears in the bounds used to propagate the error across an interval.
  • Error representation: The discretization error is expressed as an integral of divergences between true reverse rates and their approximations over coordinates and time.The proof expands the rate divergence and separates diagonal and off-diagonal contributions.
  • Conclusion: Corollary 1 supplies the interval-wise estimate that is used to conclude Theorem 3 for the uniform and remasking processes.The corollary is applied at consecutive discretization times.
  • Zero diagonal terms: Terms corresponding to unchanged coordinates contribute zero because the associated transition distributions both assign probability one to the current state.This removes diagonal contributions from the rate-divergence sum.

D Proofs of results in Appendix B

The appendix establishes identities connecting conditional probabilities, entropy, mutual information, total correlation, and dual total correlation, then applies them to reverse-process rate calculations.

  • Information identities: The appendix recalls that mutual information, total correlation, and dual total correlation satisfy basic identities used throughout the proof.These identities are combined to rewrite sums of conditional information terms.
  • DTC derivative: Summing the conditional-coordinate expressions yields an identity involving d/dv DTC(Xv).The calculation aggregates coordinate-wise conditional probabilities across dimensions.
  • Entropy reformulation: The proof rewrites conditional log-probability expressions into entropy and information quantities using the definitions of the relevant divergences.This provides the information-theoretic form needed for the main theorem’s error analysis.
  • Reverse-rate identities: The appendix derives reverse-rate identities by combining the forward kernel, conditional probabilities, and the definition of the reverse transition rate.These identities support the rate-divergence representation used in the theorem proof.
  • Remasking process: For the remasking process, the reverse-rate calculation is handled by separating transitions according to whether the destination is a regular symbol, REMASK, or MASK.The three cases determine which transition probabilities can be nonzero.

E Connection to the effective total correlation

The paper connects the analyzed quantity to effective total correlation and recovers the best known masking bound. It also obtains the first adaptive guarantees for uniform and remasking noising processes.

  • Proposition 8 connects the analyzed quantity with the presented results for the masking noising process.
  • For masking diffusion, the relevant integral equals the data dual total correlation, DTC(X).
  • The proof relates both expressions to a common information-theoretic quantity.
  • The argument uses coordinate-wise derivative identities involving conditional mutual information terms.
  • The results recover the best previously known masking bound and provide the first adaptive bound for uniform and remasking processes.
Loading 2608.23554v1…