Uniform discrete diffusion models efficiently learn small support data

Uniform Discrete Diffusion Models are Minimax Optimal for Estimating Distributions with Small Effective Support Size

Machine Learning

Summary

Understanding how to model data like text or DNA sequences is tricky because such data only use a tiny part of all possible combinations. The authors studied a type of machine learning called uniform discrete diffusion models and showed that these models can learn these kinds of data accurately without getting bogged down by the massive complexity of all possible sequences. They also proved their results are as good as possible based on the sample size and data complexity. This means their approach can effectively handle complex data that is focused on a small number of important cases.

What this means in practice

  • For machine learning engineers: Design diffusion-based generative models that efficiently learn data concentrated on small subsets of large discrete spaces.
  • For computational biologists: Model biological sequence distributions with guarantees tied to the actual diversity rather than the huge ambient sequence space.

A theory result. No direct application yet.

Authors

Dongsun Yoon, Saptarshi Chakraborty

Abstract

Discrete diffusion models have emerged as a practically successful framework for generative modeling on discrete product spaces, yet their statistical generalization properties remain poorly understood. Discrete real-world data such as text or biological sequences often concentrate on a small fraction of the astronomically large ambient space because of semantic or physical constraints, but existing bounds fail to capture this distributional structure and instead scale with the size of the ambient space, giving rise to almost vacuous error bounds. We address this gap for uniform discrete diffusion, one of the two dominant discrete diffusion paradigms alongside masking diffusion, by deriving statistical guarantees governed by the effective support size $s_n(P_0)$, a sample-size-dependent measure of distributional complexity. Given $n$ independent and identically distributed (i.i.d.) samples from an unknown data distribution $P_0$ on $[K]^d$, we show that, with appropriate choices of network size and hyperparameters, the expected total variation (TV) loss scales as $O(\sqrt{s_n(P_0)/n})$, while the expected Kullback--Leibler (KL) divergence is bounded by $O(\frac{1}{n}s_n(P_0)\log(eK^d/s_n(P_0))\log n)$. Furthermore, we show that the TV rate is minimax optimal and that the KL rate is minimax optimal up to a factor of $\log n$. Together, these upper and lower bounds show that uniform discrete diffusion successfully avoids the curse of dimensionality for distributions with small effective support size: the TV error rate depends on the ambient state-space size only through $s_n(P_0)$, while the corresponding KL rate incurs only an additional logarithmic dependence on the ambient state-space size.