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

Abstract: How diffusion models circumvent the curse of dimensionality to learn complex distributions over high dimensional spaces from a finite training set, instead of memorizing it, remains a fundamental mystery. To address this, we introduce analytically tractable Bayesian information restricted diffusion (BIRD) models, in which each pixel observes restricted information about noisy data. A BIRD model time-reverses diffusion by inferring which past training sample produced its current restricted observation using the Bayesian posterior. This model class generalizes existing analytical diffusion models that use spatially local information restriction. We show that spatially local BIRD models closely approximate trained diffusion models \textit{early in training}, across different architectures such as UNets and DiTs. Under minimal assumptions on the data distribution, we identify an information-theoretic phase boundary between memorization and generalization in the joint space of amount of training data, time in the reverse generative process, and amount of information restriction: a BIRD model memorizes when the mutual information between its restricted noisy observations and the training data exceeds the log number of training points, and it generalizes otherwise. Experiments across a range of datasets confirm our theoretically predicted location for the transition. We find that generation proceeds near the edge of memorization: both spatially local BIRD models and early-training diffusion models track the memorization-generalization phase boundary by increasingly restricting information over time. Overall, our results reveal a fundamental role for information restriction in generative AI to circumvent the curse of dimensionality.

• The paper develops the BIRD framework that quantifies the transition between memorization and generalization using a precise mutual information criterion.
• It employs spatially restricted Bayesian posterior means for denoising, with validations on datasets like CIFAR-10 and CelebA showing high agreement between theoretical predictions and neural models.
• The study demonstrates how optimal information bottlenecks can eliminate the curse of dimensionality, guiding the design of scalable generative diffusion models.

An Exact Information Theory of Generalization Phase Transitions in Bayesian Diffusion Models

Overview

The paper "An exact information theory of generalization phase transitions in Bayesian diffusion models" (2607.08041) develops an exact, analytic information-theoretic framework for understanding when and how generative diffusion models transition from memorizing a finite training set to generalizing the underlying data distribution. The authors introduce the Bayesian Information Restricted Diffusion (BIRD) class of models and derive a set of results that accurately predict the transition between memorization and generalization, quantifying the impact of dataset size, the amount of information accessible by a model (e.g., patch size), and reverse diffusion time/noise. Of special note is the rigorous proof that generalized diffusion modeling can avoid the exponential curse of dimensionality under realistic data statistics—explaining the empirical effectiveness of modern diffusion models on high-dimensional image data. The work is both theoretically and empirically validated across standard image datasets and neural architectures such as UNets and DiTs.

BIRD Models and the Information Bottleneck Approach

The BIRD framework formalizes a family of diffusion models whereby each pixel observes a restricted channel $\mathcal{C}_{x,t}$ of the noisy image $\phi_t$ during the reverse diffusion process. Instead of using the entire image to make a denoising decision, each pixel is constrained to operate only on a local or otherwise restricted view. Crucially, the BIRD denoiser at each pixel is constructed as the Bayesian posterior mean over the training set, conditioned on this restricted observation. This construction is analytically tractable and enables derivation of the mutual information between observations and data.

One of the essential outcomes of this formalism is that information restriction (as opposed to unrestricted Bayes-optimal score models) makes the Bayesian guessing game harder, thereby preventing the collapse to memorization unless the observation channel is sufficiently informative to resolve individual data points in the training set.

Mutual Information Phase Transition: Criteria and Empirical Validation

A central theoretical result is the identification of an exact criterion for the memorization-generalization transition in BIRD models:

$\ln |\mathcal{D}| = I(\mathcal{C}_{x,t}; \varphi)$

where $|\mathcal{D}|$ is the size of the training set, and $I(\mathcal{C}_{x,t}; \varphi)$ is the mutual information between a pixel’s restricted observation and the underlying (clean) data under the forward diffusion channel.

• Generalizing phase: $I(\mathcal{C}_{x,t}; \varphi) < \ln |\mathcal{D}|$
• Memorizing phase: $I(\mathcal{C}_{x,t}; \varphi) > \ln |\mathcal{D}|$

The mutual information thereby demarcates a phase boundary in the (amount of training data, time/noise, information-restriction) space.

This is empirically validated using both analytical forms (on Gaussian mixtures and other tractable data distributions) and actual experiments on real datasets such as CIFAR-10 and CelebA. The authors show close agreement between mutual information theory curves and observed transitions, as evidenced by entropy deficit metrics saturating to zero precisely as $I(\mathcal{C}_{x,t}; \varphi)$ meets $\ln |\mathcal{D}|$.

Figure 1: A plot of the mutual information $I(\phi_{\Omega};\varphi)$ quantifying the phase boundary between memorization and generalization as a function of noise and patch size.

BIRD Models as Predictors for Neural Diffusion Training Dynamics

A notable empirical finding is that spatially local BIRD models do not merely provide a theoretical construct but also quantitatively predict the detailed behavior of neural diffusion models (both UNets and DiTs) early in training across diverse datasets. That is, when trained on disjoint data subsets and given the same noise input, both neural models and the corresponding BIRD analytic model generate nearly identical images, with median pixelwise $\phi_t$0 as high as 0.85–0.93 in early epochs.

Figure 2: Median $\phi_t$1 and MSE trends showing high agreement between BIRD models and neural diffusion models across datasets as a function of training epochs.

This demonstrates that, prior to overfitting/late memorization, neural diffusion models are well-approximated by the optimal BIRD model for the network’s effective inductive bias—namely, its spatial locality and information restriction constraints.

Figure 3: Uncurated samples comparing DiT, UNet, and local BIRD model outputs on CelebA64 early in training, evidencing strong visual similarity.

Theoretical Analysis of Critical Scale and Curse of Dimensionality

The paper rigorously analyzes the relationship between patch scale, noise, dataset size, and generalization. When model observations are spatially restricted to image patches of length scale $\phi_t$2, there exists a critical scale $\phi_t$3 at each noise level below which generalization occurs and above which memorization dominates. The derivation relies on the monotonicity of mutual information with respect to both patch size and noise.

For natural image data characterized by a power-law spectrum $\phi_t$4, the required dataset size for generalization scales only as $\phi_t$5 for images of size $\phi_t$6, where $\phi_t$7 is small for scale-invariant, natural scenes. Thus, for $\phi_t$8 (full scale invariance), the exponential curse of dimensionality is eliminated—consistent with the empirical scalability of diffusion models.

Figure 4: BIRD and neural model outputs on CIFAR-10 early in training showing generative consistency with local information restriction.

The shape of the memorization-generalization boundary in $\phi_t$9 space is shown to align with the optimal denoising (Wiener filter) patch scale, and empirical critical scales closely track theoretical predictions from Gaussian constraints and information-theoretic bounds.

Figure 5: Output comparisons on FashionMNIST, further validating the predictive accuracy of calibrated local BIRD models in early training.

Figure 6: Analogous output samples on MNIST for increasing epochs across BIRD, UNet, and DiT instantiations.

Multi-Scale Generation and Entropic Criticality

Generation proceeds at the edge of memorization, with reverse diffusion trajectories following the phase boundary as information is increasingly restricted over time. This manifests empirically as an evolution from patchwork “mosaic” output at high noise (large $\ln |\mathcal{D}| = I(\mathcal{C}_{x,t}; \varphi)$0) toward globally consistent, creative samples once restricted to the critical scale for generalization.

Figure 7: Evolution of model generations from local “patch mosaic” regimes toward coherent global samples as training proceeds.

Implications and Future Directions

This work provides rigorous, general conditions under which finite-sample diffusion models can generalize rather than memorize—showing that model inductive biases (especially spatial information restriction) play a fundamental role in compressing the required data for high-dimensional generalization. It clarifies why naive empirical (unrestricted) score models fail and justifies why convolutional/transformer architectures, with appropriate information bottlenecks, scale so effectively without requiring exponential data.

The theoretical framework and information-theoretic phase diagram established here create a foundation for future studies on:

• The interplay between architecture, training dynamics, and generalization in non-Bayesian neural models.
• Extension to more complex/structured data and tasks beyond images.
• Development of new model classes exploiting optimal information restriction to achieve efficient, scalable generative modeling.

Conclusion

The paper presents a comprehensive, exact information-theoretic account of memorization and generalization in generative diffusion models, encapsulated in the BIRD framework. It identifies mutual information as the critical quantity governing the phase transition, with spatial information restriction providing the mechanism to circumvent the curse of dimensionality on natural data. Experimental validations substantiate the predictive and descriptive power of BIRD models, establishing them as both a theoretical and practical tool to understand and design scalable generative AI systems.