Limits on understanding confidence in token-based diffusion models
Limits of Confidence in Diffusion
Artificial Intelligence
Summary
This paper looks at how discrete diffusion models produce sequences by updating multiple parts at once, each chosen from probabilities at each position. The authors found that if the parts being updated depend on each other, the model’s step won’t match the training data distribution well. They show mathematically and with a test example that token dependencies matter and that just knowing single-token probabilities isn’t enough to capture these dependencies. This means that the model’s confidence in certain tokens can be overestimated when dependencies exist.
What this means in practice
- •For machine learning engineers: Design diffusion-based generative models with awareness of token dependency limits for improved training and sampling accuracy.
- •For natural language processing teams: Evaluate and adjust token confidence estimations in sequence models to account for dependencies between words or tokens.
Tested on simulated data.
Authors
Russ Webb, Amitis Shidani, Alice Bizeul, Dan Busbridge
Abstract
Discrete diffusion, including remasking and uniform-state samplers, generate a sequence by writing multiple token positions per step, drawing each from a per-position distribution and choosing which positions to write from those same distributions. For domains of general interest (pixels, phonemes, or words) there are inherent dependencies between tokens. We show that a step matches the training distribution only when the positions it writes are conditionally independent given the tokens already fixed, that no product of per-position distributions can match a dependent group, and that per-position distributions do not determine whether a group is dependent: two joint distributions can have identical per-position marginals while differing in which combinations of values occur. On ScanAndAdd, a synthetic task whose joint distribution is available in closed form, we verify that every group of two or more undetermined positions a confidence ranking writes is dependent, and measure the generated distribution to be $29\times$ the sampling-noise floor total variation while per-sample metrics are $1.0$.