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.
Discrete diffusion models have demonstrated strong performance across a range of datasets, including natural language data and graph-structured data. Among many variants, score-entropy discrete diffusion (SEDD) has achieved particularly strong empirical results. In SEDD, new samples are generated by iteratively evaluating a sequence of concrete score functions, which are learned by minimizing a score-entropy loss. While much of the prior theoretical literature on discrete diffusion has focused on the sampling efficiency of SEDD under the assumption of small score estimation error, recent work has begun to investigate the finite-sample properties of score estimation itself. In this work, we take a different route by investigating the fundamental statistical limits of concrete score estimation. We focus on uniform and masking discrete diffusions, two of the most widely adopted discrete diffusion models. We establish a minimax lower bound under the score-entropy loss, and propose an MLE-based thresholding estimator that matches this lower bound up to constant and polylogarithmic factors that depend on neighboring density ratios. We further show that, for any target distribution, this density ratio is naturally controlled under both uniform and masking discrete diffusion models, yielding nearly matching minimax lower and upper bounds for the aggregated score estimation error. Our results imply that, with appropriate initialization and discretization, SEDD can achieve nearly optimal minimax sample complexity, as measured by the KL divergence between the target and generated distributions.
Chenghan Xie, Jose Blanchet, Renyuan Xustat.ML cs.LG math.OC
Estimating the difference of two Stein's score functions is a fundamental problem in generative modeling. In particular, score differences arise naturally in transfer learning, where the score difference provides the mechanism for adapting a pre-trained model to a new target distribution, and in diffusion model-based post-training methods such as discriminator guidance. Existing estimators for score differences in these settings either lack of statistical consistency or are difficult to scale up in high-dimensions. We propose a statistically consistent and scalable estimator for score differences based on Sobolev regularization, which plays a crucial role in ensuring consistency and stablizing the training in the small-sample regime. Mathematically, we establish a convergence rate of $O(n^{-\frac{s-1}{d+2s-2}})$ where $d$ is the dimension and $s$ denotes the smoothness of the underlying densities, and provide a minimax lower bound of $\tildeΩ(n^{-\frac{2(s-1)}{d+2s}})$ (in mean-squared error). Empirically, our estimator exhibits significantly improved stability in small-sample regimes compared to existing methods. We demonstrate its effectiveness on real-world tasks, including transfer learning for ECG signal generation, where it substantially outperforms non-regularized score difference estimators in downstream classification performance.
Recently, rectified flow has emerged as a fundamental framework for large-scale image generation, powering state-of-the-art systems such as FLUX.1 and Stable Diffusion 3. Despite its remarkable empirical success, the computational and statistical guarantees of iterative rectified flow have remained largely unexplored. We address this problem by studying \textit{c}-rectified flow, a cost-aware class of rectified flow that projects velocity fields onto a gradient class while preserving endpoint marginals. The ordinary rectified flow can fail to recover the optimal transport coupling: in a Gaussian case study, the iteration converges to the optimal coupling if and only if the source and target covariance matrices commute. In contrast, under suitable compactness and uniform-integrability assumptions, iterative \textit{c}-rectified flow always converges to the optimal transport coupling. We further establish quantitative one-step contraction and exponential convergence guarantees under projection-stability assumptions for both quadratic and strongly convex displacement costs. Finally, under a Hölder ball assumption, we develop new minimax-optimal score estimation rates and show that, when combined with iterative \textit{c}-rectified flow, they yield a rate-optimal estimator of the optimal transport for the dimension \(d \ge 3\) and a nearly parametric rate for \(d=1,2\).
Nicolas Brosse, Arnak S. Dalalyanmath.PR cs.LG math.ST
We study the small-noise asymptotics of Euclidean heat regularizations of probability measures supported on manifolds with corners. Near a boundary or corner stratum, the relevant regime is a conical boundary layer in which the observation point approaches the stratum at the same scale as the Gaussian smoothing parameter. After rescaling this layer, the support is replaced to leading order by its inward tangent cone. We prove a two-term expansion for the heat-regularized density in this regime. The leading coefficient is the Gaussian mass of the linearized cone, weighted by the density on the support and by the adapted corner Jacobian; the first correction records the variation of the density, the Jacobian, and the quadratic geometry of the embedding. A localization argument then yields the corresponding expansion for the full heat regularization, with the nonlocal contribution exponentially small. From this density expansion we derive logarithmic asymptotics and uniform expansions for the score, the log-Hessian, and the scale derivative of the score. These formulas show how lower-dimensional support, boundary faces, corners, and curvature are encoded in the singular differential structure of small-noise Gaussian regularizations.
Score-based generative models and Langevin samplers rely on estimating the score function $\nabla_x\log p_t(x)$ of a forward diffusion. Classically this is tractable when the drift is linear: the marginal density is Gaussian and the score is a global conditional expectation. For a general nonlinear, state-dependent drift the marginal density has no closed form, and existing methods--denoising score matching and global Fokker--Planck residual penalties--resort to global averaging that inflates estimation error in low-density regions precisely where accuracy is most critical. We address this by developing a local Fokker--Planck geometric framework that replaces global conditioning with local parabolic averaging. Our approach rests on three contributions. First, a time change to the cumulative-variance coordinate reduces the variable-coefficient Fokker--Planck equation to a standard inhomogeneous heat equation, on which we extend Evans' classical heat-ball monotonicity method to derive exact local mean-value representations for the score $\nabla_x\log p$ together with the density, log-density, and entropy density; local well-posedness is established under an explicit dimension-dependent drift budget. Second, for high-dimensional Monte Carlo evaluation of the resulting heat-ball integrals, we introduce the $κ$-measure and derive its exact factorized sampler with unit per-sample weight, $χ^2_2$ radial concentration. Third, the $r\to0$ limit of the heat-ball residual recovers the pointwise Fokker--Planck residual, showing that the local framework is a one-parameter generalization of global FP-residual methods, and that the DSM population minimizer is feasible for the heat-ball constraint at every scale. We validate the framework on 2D structured data on 256-dimensional MNIST, and on a dedicated sampler study confirming the concentration laws.
Sampling from an unnormalized target density by reversing an Ornstein-Uhlenbeck diffusion requires the score of each noise-perturbed marginal law. Two exact identities are available: Tweedie's identity and a target-score identity, each yielding unbiased finite-reference score estimators for the OU-marginal score. Score estimators induced by scalar blends of Tweedie and TSI score estimators can reduce variance, but they are too rigid for singular or strongly anisotropic targets. We formulate blended score estimation as a conditional risk-minimization problem over matrix valued blending coefficients, referred to as gates. Our central result is to show the optimal matrix valued gate for blended score estimation is given \[ G_\star(y,t) = α_t^2 \left(α_t^2 I_d + γ_t\, \mathbb{E}[H_0(X_0)\mid Y_t=y] \right)^{-1}, \qquad H_0=-\nabla^2\log p_0 .\] Here $α_t = e^{-t}$ and $γ_t = 1-e^{-2t}$ are the OU coefficients, and the conditional expectation is under the OU posterior of $X_0$ given $Y_t=y$. We call this formula the \emph{Laplace-Fisher Gate Identity} (\LFGI{}). Because the Tweedie-TSI disagreement has conditional mean zero, the gate changes the score-estimator variance but not its expected value. We derive the variance-optimal matrix gate, record the Gaussian special case, and establish finite-reference consistency and stability bounds for estimating the gate from weighted reference samples. We then use the finite-reference LFGI score estimator for normalized density evaluation in Bayesian inverse problems. In regimes where MCMC pilot samples and derivative information are already available, LFGI uses those byproducts to construct a normalized surrogate for the posterior density. The resulting surrogate supplies information that the MCMC samples alone do not provide: posterior-energy evaluation, model-evidence estimation, and downstream density-based diagnostics. On a PDE-constrained inverse-problem benchmark, the LFGI surrogate improves posterior-density calibration and sampling diagnostics relative to the other tested score-estimator classes. Experiments using LFGI with known model evidence check absolute evidence calibration in both Gaussian and non-Gaussian settings.
The task of compositional generation involves using a conditional generative model, trained only on a subset of the possible conditions, to produce samples from compositionally-defined target distributions such as a geometric combination of the source distributions. In this work, we argue that this task is often infeasible for vanilla conditional diffusion models: we conjecture that no inference-time technique can efficiently produce samples from the target distribution in certain well-motivated settings. This idea is supported by theory-guided generalization arguments and carefully-designed experiments on both synthetic and realistic data. In particular, while recent methods such as Feynman-Kac correction reduce inference-time approximation error, our results show that score estimation error has a more catastrophic effect on performance when the target distribution is out-of-distribution with respect to the sources, highlighting the need for a different approach to this task.
Abbas Mammadov, Ozgur Kara, Kaan Oktay +5cs.LG cs.CV stat.ML
Diffusion and flow-based models learn powerful data priors by training a denoiser to reverse Gaussian corruption. To use this prior to solve a linear inverse problem, one needs to sample from the posterior, but the score that the prior provides is the unconditional score, not the posterior score. Existing methods either steer a fixed pretrained denoiser with approximate measurement-matching corrections, or train a conditional restoration model that abandons the denoising structure of the prior. We derive the exact posterior score in closed form for linear Gaussian inverse problems under general Gaussian interpolants, and show that posterior sampling reduces to a denoising problem at an operator-dependent shifted pivot under an anisotropic noise covariance. We turn this identity into Exact Posterior Score (EPS), a denoising training objective that preserves the input/output structure of standard pretraining and can therefore be trained from scratch or fine-tuned from a pretrained denoiser. At inference, EPS uses the same sampler as the underlying backbone, with no likelihood gradients or projections. We evaluate EPS on five linear inverse problems across FFHQ and ImageNet, where it outperforms training-free and training-based baselines on fidelity, perceptual, and distributional metrics, while using roughly an order of magnitude fewer denoiser evaluations than gradient-based posterior samplers.
Continuous causal discovery typically couples representation learning with structural optimization via non-convex acyclicity penalties, which subjects solvers to local optima and restricts scalability in high-dimensional regimes. We propose a decoupled paradigm that shifts the causal discovery bottleneck from non-convex optimization to statistical score estimation. We introduce the Score-Schur Topological Sort (SSTS), an algorithm that extracts topological order directly from unconstrained generative models, bypassing constrained structure optimization. We establish that the causal hierarchy leaves a geometric signature within the score function: iterative graph marginalization is mathematically equivalent to computing the Schur complement of the Score-Jacobian Information Matrix (SJIM) under linear conditions. This translates the acyclicity constraint into an algebraic procedure with a dominant cost of O(d^3) operations. For non-linear systems, we formulate the expectation gap of Schur marginalization and introduce Block-SSTS to compress extraction depth, bounding structural error. Empirically, SSTS allows causal structural analysis on non-linear graphs up to d=1000. At this scale, our framework indicates that once the non-convex optimization bottleneck is mathematically bypassed, the structural fidelity of continuous causal discovery is bounded by the finite-sample estimation variance of the global score geometry. By reducing graph extraction to matrix operations, this work reframes scalable causal discovery from a constrained optimization problem to a statistical estimation challenge.