An exact information theory of generalization phase transitions in Bayesian diffusion models

Diffusion models have demonstrated a remarkable ability to generate high-quality data without simply memorizing their training sets. However, the theoretical mechanisms that allow these models to circumvent the curse of dimensionality—effectively learning complex distributions from finite data—have remained elusive. This paper introduces Bayesian Information Restricted Diffusion (BIRD) models to provide an analytically tractable framework for studying this phenomenon.

BIRD models treat each pixel as an agent that observes a restricted, noisy version of the data. Instead of attempting to learn the full score function of the entire image, the model performs optimal Bayesian inference to guess which training sample produced the current restricted observation. By limiting the information available to each pixel (e.g., through local spatial patches), the model is forced to generalize rather than memorize. The authors demonstrate that these BIRD models serve as an excellent proxy for the behavior of modern neural architectures, such as UNets and Diffusion Transformers (DiTs), particularly during the early stages of training.

Using information theory, the authors identify a precise phase boundary between memorization and generalization. The transition is governed by the mutual information between the restricted observation and the training data: if this information exceeds the log of the number of training samples, the model memorizes the data. If it falls below this threshold, the model generalizes.

Crucially, the study finds that successful generative models operate near the edge of this phase boundary. As the reverse generative process progresses (moving from high noise to low noise), the models dynamically adjust their information restriction to track this boundary. This suggests that the inductive biases inherent in diffusion models—such as spatial locality—are essential for maintaining the balance between learning general features and avoiding the trap of memorization.