The information geometry of product-reference discrete diffusion: Interaction growth complexity and optimal scheduling
Martin J. Wainwright
stat.ML, cs.AI, cs.LG, math.ST
2026-08-29
IGC sets PRD sampling cost: equispaced log-SR-odds tracks total IGC mass (≤2 min{TC,DTC}); optimized grids scale with (∫√q)². A bad reference can cost Θ(√d) or Θ(log d).
Discrete diffusion corrupts a clean sample with a Markov kernel on a finite alphabet, then tries to reverse the corruption. The exact reverse kernel is exponential in dimension. In practice, coordinates are updated independently: the joint reverse is replaced by the product of one-coordinate reverse marginals. Lou et al. showed that this tokenwise-independent kernel is KL-optimal among τ-leaping strategies, for each input state.
The theoretical gap is adaptivity. For masking diffusion, Dmitriev, Huang, and Wei showed that sampling cost is controlled by a functional bounded by min{TC, DTC}, so the sampler tracks low-dimensional latent structure. For one non-masking CTMC scheme they analyzed, the upper bound scales with dimension, and they proved an algorithmic lower bound: that particular sampler does not adapt. They left open whether some other discrete diffusion sampler could.
This paper answers yes, for product-reference diffusion (PRD). The terminal noise is an arbitrary product measure. Uniform noise, absorbing-state noise, and marginal-matched noise all sit in this class.
The forward process is parameterized by squared reliability r in [0, 1]. Each coordinate is kept as-is with probability √r, otherwise redrawn from a product reference νi. The clean law sits at r=1; the product reference sits at r=0. The channel satisfies a semigroup identity, so {Xr} is a genuine continuous-parameter Markov process.
The exact reverse kernel needs the path marginals and is intractable. A PRD sampler does three things:
Complexity is not the mutual information itself. It is how fast that information, and the cross-coordinate part of it, is removed along the path. Give each coordinate its own reliability ti. IGC starts from the off-diagonal Hessian of Info(Z; Xt), summed and signed to get g(r). Weight by r(1-r) and integrate to get the IGC mass G(a,b) on an interval. In log-SR-odds time λ = log(r/(1-r)), the density is q(λ). In plain terms, q records how much coordinate interaction remains at that denoising layer, and how fast it is being taken apart.
Log-SR-odds is forced by Theorem 1. The one-step KL bound has a prefactor equal to the ratio of squared-reliability odds ψ(r)=r/(1-r), minus one. Equal steps in λ make that prefactor identical across steps, so a schedule can be written down without knowing the shape of q.
The true object is bivariate. The kernel QIGC(η,ξ) is a mixed partial of a cut mutual information; on the diagonal, QIGC(η,η)=q(η). The exact one-step KL error is the double integral of QIGC over the upper triangle, and G is the integral along the diagonal. Off the diagonal, Q is sandwiched by an exponential in |ξ-η|. That comparison yields G(0,1) ≤ 2 min{TC, DTC}.
If clean samples from the target are available, a Bregman divergence hidden in the same proof gives an unbiased estimator of the IGC increment on each interval, which is then used to allocate steps across blocks.
Theorem 1: on any grid, the output KL is at most 2 Σj (ψ(r{j+1})/ψ(rj) − 1) G(rj, r{j+1}), plus an initialization cost TC(X{r0}) and a termination cost TC(Z | XR). With symmetric endpoints rN=1-r0, those two terms are ≤ c r0 (d|A|)^3. Setting r0=(d|A|)^{-k} for k≥4 drives them to a polynomial; the extra steps are only logarithmic in r0.
For an equispaced log-SR-odds grid (Corollary 1, N≥2ℓd), KL ≤ (8ℓd / N) G(r0, 1-r0) plus the endpoint terms, with ℓd=log((1-r0)/r0). Reaching KL error ε takes on the order of (G/ε) log(d|A|/ε) iterations. Combined with G(0,1)≤2 min{TC(Z), DTC(Z)}, even this oblivious PRD schedule inherits the low-dimensional bounds already computed for hidden Markov models, stochastic block models, and quantized latents in the masking literature.
Theorem 3 turns the bound into a sharp leading term. On the equispaced grid, the sum of exact one-step KL errors equals CIGC/(2N)+o(1/N). The best N-step grid equals PIGC/(2N)+o(1/N). Here CIGC=2ℓd ∫q and PIGC=(∫√q)^2. Cauchy–Schwarz gives CIGC/PIGC ≥ 1, with equality iff q is constant. The more peaked q is, the more a refined schedule buys.
A K-block schedule keeps equal log-odds steps inside each block and allocates steps in proportion to √(Sk Gk). That explicit rule is within a factor of four of the best integer allocation for the same partition. Refining the partition monotonically lowers the complexity toward PIGC.
The reference ν reshapes q. Two paired-binary ensembles, with no language or graph sampling bake-off:
| Ensemble | Parameter | Uniform | Absorbing | Marginal-matched |
| A | εd=d^{-1/2} | aggregate IGC does not blow up in d | similar to uniform | aggregate IGC grows as Θ(√d) |
| B | 1-εd=d^{-1} | constant order | Θ(log d) | no asymptotic stated in the text |
There is no universal ranking. Marginal matching, often preferred in graph generation, can be worse than uniform by a √d factor on Ensemble A.
The density q itself traces data geometry. For a noisy repeated bit at d=128, q has a sharp peak at crossover η=0.01 and the peak dies as η→1/2; a product law makes q identically zero. For Curie–Weiss, q is unimodal below criticality, bimodal at β=1.20, and nearly unimodal again at β=1.60, matching a global choice of magnetized phase plus residual within-phase fluctuations. Hierarchical prototype mixtures grow one peak per tree level when Hamming gaps differ by orders of magnitude: two peaks at L=2, d=32776; three peaks at L=3, d=33288.
If the one-site posterior is learned rather than exact, the KL bound picks up an additive coordinatewise denoiser penalty Eden(a,b), booked separately from discretization error.
For anyone tuning discrete diffusion, noise type and the time grid stop being purely empirical knobs. They become computable geometric quantities. An equispaced log-SR-odds grid needs no clean samples and still adapts to TC/DTC structure. With clean samples one can estimate IGC and spend more steps where q is large, moving the leading cost from CIGC down to PIGC.
The paper also closes the question Dmitriev et al. left open: non-masking discrete diffusion is not doomed to scale with dimension. At least for oracle-posterior PRD, adaptivity is possible.
This is still a discretization theory, not a new sampler with NFE numbers on language or molecules. Putting IGC into a training loop still requires estimating q from finite clean samples and replacing exact posteriors with learned denoisers.
The main bound in Theorem 1 uses exact one-site posteriors. Learned posteriors appear only as an additive penalty. There is no experiment showing that an IGC schedule remains stable under the approximation error of a real network.
The √d and log d gaps for the reference law come from specific paired-binary constructions, plotted out to d around 10^8. Empirical gaps reported by Austin, DiGress, and later graph papers get an explanatory frame here, not a replication of how many function evaluations a reference swap saves on language or graphs.
The min{TC, DTC} bound can be loose. Following an independent-block argument of Dmitriev et al., the paper notes an Ω(d) gap between aggregate IGC and min{TC, DTC}. A provable equispaced bound can therefore sit far above the true difficulty.
A concurrent CTMC paper by Dmitriev, Huang, and Wei was posted on 24 August 2026. The author flags a planned comparison in a revision; this version has none.
Adaptive multi-block schedules need clean samples from the target. Generative training usually has no oracle draws from PZ, and the paper does not wire the estimator into a training loop.
The bivariate kernel is the true source of exact KL. Most usable guarantees still run through the univariate density q. How much the diagonal approximation loosens when q is peaked and the grid is coarse, beyond the exponential sandwich, is left unquantified.