Skip to content
Preprint

Provably adaptive sampling with uniform and remasking discrete diffusion models

Aug 2026 · 1 citation · 36 references
Computer Science Mathematics

TL;DR

This work considers a first-order sampler based on the leave-one-out denoiser for uniform and remasking processes whose coordinate updates can be performed in parallel, and derives 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.

Abstract

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 $\tau$-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.

View source

Similar papers

Preprint Aug 2026

Forward-Evolution Error Analysis and Adaptive Design for Matrix-Valued Diffusion Models

This analysis transfers reverse-time discretization errors to the forward corruption law and treats two numerical schemes within a common framework, which identifies a local error, accumulates it through the forward evolution, and inserts the result into a common KL decomposition.

T. Pang, Zuowei Shen, Ruitong Zhang · 1 citation
#machine learning Preprint Aug 2026

Exact Global MCMC with Denoising Diffusion

This work shows that diffusion models learned with standard denoising loss can provide effective global MCMC proposals for complex high-dimensional target densities and offers preliminary evidence that the established scaling behavior of standard diffusion training transfers directly to exact sampling from high-dimensional unnormalized densities.

Mitch Hill · 0 citations
#artificial intelligence Preprint Sep 2026

From Truncation to Commitment: Persistent Context in Uniform Discrete Diffusion

Uniform-state discrete diffusion models update all tokens in parallel while keeping every position revisable. Even when the commonly used top-$p$ rule leaves only one candidate at a position, that choice affects only the current reverse step and can be revised at the next sampling step. We ask what changes when selected hypotheses instead become persistent context for later predictions. We therefore propose committed reveal sampling (CRS), a training-free sampler that stores selected argmax tokens and inserts them into subsequent model inputs. Our analysis gives a rationale for selecting later and for keeping selected tokens visible. Under the exact forward process, the Bayes error of selecting a clean token cannot increase as noise decreases, while in a simple latent-mode model, keeping the selected token visible helps later parallel predictions agree on the same sequence-level choice. Empirically, paired experiments on Duo-distilled then separate this persistent effect from single-step top-$p$ restriction and scalar temperature scaling. Under the same finalization rule, CRS without top-$p$ truncation reaches lower generative perplexity (GenPPL) than fixed $p=0.95$ and $p=0.9$ baselines across budgets of 8--64 function evaluations (NFE). At 64 NFE, the comparison at matched unigram entropy also gives lower GenPPL for CRS, yielding a more favorable GenPPL--entropy tradeoff. Base Duo shows the same direction in a descriptive comparison, while other diversity and continuation metrics can rank these operating points differently. These results identify support restriction and persistent context as distinct controls of that tradeoff.

Satoshi Hayakawa · 0 citations
Preprint Aug 2026

When Denoising Hurts: Rethinking the Terminal Step of Diffusion Time Series Forecasters -- Extended Version

Diffusion models offer a natural way to model uncertainty in time series forecasting, yet their iterative sampling process is often treated as a uniformly beneficial refinement procedure. Our study challenges this view by examining how forecast quality evolves throughout reverse diffusion. We find that general temporal structure is often recovered at relatively high noise levels, whereas continued low-noise refinement can introduce statistical drift and degrade the final forecast. Our analysis further suggests that this behavior explains why prior methods often favor relatively narrow diffusion architecture and schedule design. Building on this observation, we propose a label-free global stopping criterion that detects the optimal termination point, eventually speeding up inference and improving predictive accuracy. Additionally, since early stopping terminates inference in high-noise regions, we propose a Bernoulli timestep sampler that concentrates training on this region while preserving coverage of the full diffusion process. Extensive experiments conducted across eight real-world datasets demonstrate the superior performance of our method compared to existing approaches.

Dat Nguyen-Cong, L. Tran, T. Kieu · 0 citations
Preprint Sep 2026

ASSERT: Adaptive Stochastic Sampling for Robust Diffusion Models on Analog Compute-in-Memory Hardware

Diffusion models achieve strong image generation quality but incur high iterative denoising costs. Analog compute-in-memory (CIM) can accelerate matrix-vector multiplications, yet spatial memory variations perturb weights and accumulate during sampling. Unlike conventional neural networks, diffusion models'temporal sensitivity to hardware noise remains underexplored. We investigate diffusion inference using a noise model calibrated and validated against measurements collected from multiple physical CIM chips. Our results show that the early, high-noise denoising stage is substantially more vulnerable than the final refinement stage. A first-order trajectory analysis attributes this behavior to the repeated propagation of correlated prediction errors induced by a fixed hardware mapping. Based on this observation, we propose ASSERT, a training-free sampler that uses higher stochasticity early and smoothly transitions to deterministic denoising. The injected stochasticity changes subsequent activation trajectories and thereby reduces their alignment with persistent spatial errors. Across the evaluated settings, ASSERT achieves up to 2.58$\times$ lower FID than deterministic DDIM on high-resolution datasets and 7.68$\times$ lower FID in the CIFAR-10 step-count study, without changing model parameters or the number of network evaluations.

Yuan-Nuo Feng, Yizhe Chen, Wen-Shuai Yao et al. · 0 citations
Preprint Sep 2026

On the computational cost of Stochastic Gradient Langevin Dynamics

Stochastic Gradient Langevin Dynamics (SGLD) reduces the cost of Langevin-based sampling by replacing full-dataset drift evaluations with mini-batch approximations, but the resulting subsampling error may offset this computational saving. We study this trade-off for stochastic differential equations with finite-sum drifts and compare the computational cost of SGLD with that of the Euler-Maruyama (EM) method. For a prescribed mean-square accuracy $\varepsilon^2$, we derive complexity estimates that explicitly track the dependence on the dataset size $m$, mini-batch size $s$, and accuracy parameter $\varepsilon$. The resulting comparison reveals distinct parameter regimes in which either method is preferable. In particular, EM can have lower leading-order cost only in a small-data, aggressive-subsampling regime, whereas SGLD is favoured over most of the remaining parameter space. In the practically relevant regime $s \ll m$, the transition between the two methods occurs at the scale $m \asymp \varepsilon^{-1}$. We complement the theoretical analysis with numerical experiments based on a Gaussian Bayesian inference model, which examine the predicted cost regimes together with the underlying discretisation error and variance estimates.

Mateusz B. Majka, Tigran Nagapetyan, Łukasz Szpruch et al. · 0 citations

We use cookies to run the site and, with your consent, for analytics and to show ads. See our Cookie Policy.