The data geometry of masking diffusion: Certified-optimal schedules via unmasking growth complexity
2026-08-13 • Machine Learning
Machine LearningArtificial IntelligenceInformation Theory
AI summaryⓘ
The authors investigate a method called masking diffusion used for sampling discrete data and propose a new way to measure the complexity of data geometry, which they name unmasking growth complexity (UGC). They show that this measure helps control the error in approximating the sampling process and enables the design of optimized sampling schedules that adapt to the data's structure. The authors also provide a practical way to estimate UGC from data samples, leading to samplers that reliably meet error targets with near-optimal efficiency. Finally, they connect their new complexity measure to existing concepts and demonstrate significant improvements in certain high-dimensional settings.
masking diffusiondiscrete samplingunmasking growth complexityKullback--Leibler errorBernoulli-subset schemesadaptive samplinglog-reveal-oddsEuler discretizationmultivariate dependencehigh-dimensional statistics
Authors
Martin J. Wainwright
Abstract
We study masking diffusion for discrete sampling and introduce a path-resolved measure of data geometry called the \emph{unmasking growth complexity} ({\textsf{UGC}\xspace}). Its local increments directly control Kullback--Leibler (KL) discretization error, yielding a unified analysis of Bernoulli-subset and fixed-cardinality unmasking schemes. In log-reveal-odds coordinates, this structure yields optimized single-block and multi-block schedules, and quantifies the gains from adapting computational effort to data geometry. Crucially, we show how {\textsf{UGC}\xspace} increments can be estimated from samples via KL increments along coupled reveal trajectories. This leads to \emph{certified-optimal} samplers that achieve a prescribed KL error with high probability and iteration complexity within a constant factor of the corresponding oracle procedure. Collapsing the \ugc path yields the aggregate {\textsf{UGC}\xspace} mass, which connects to classical multivariate dependence measures and complexity measures from previous analyses of discrete diffusion. In the fine-partition limit, the squared integral of the square-root {\textsf{UGC}\xspace} density determines the sharp leading-order optimal Euler discretization error. Examples exhibit substantial dimension-dependent gains over coarse schedules, including $\widetildeΩ(\sqrt{d})$ improvements achievable with a constant number of adaptively placed blocks.