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

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

Problem Statement
A fundamental mystery in generative AI is how diffusion models circumvent the "curse of dimensionality" to learn complex distributions over high-dimensional spaces from finite training sets without merely memorizing them. Naively, transitioning from memorization to generalization for arbitrary data distributions should require an intractable (exponential in dimension) amount of data. While empirical score models (which reverse diffusion using the exact empirical score of a finite dataset) are known to memorize, trained neural diffusion models often generalize robustly. Existing theoretical approaches have struggled to provide a precise, predictive theory for when and where the transition from memorization to generalization occurs for trained, nonlinear diffusion models across diverse architectures and datasets.

Methodology
The authors introduce a new class of analytically tractable models called Bayesian Information Restricted Diffusion (BIRD) models.

• Core Concept: In BIRD models, each pixel $x$ acts as an agent that observes only a restricted, possibly stochastic, portion of the noisy image $\phi_t$ (denoted $C_{x,t}$) rather than the full image. The pixel then uses optimal Bayesian inference to guess which past training sample $\phi \in \mathcal{D}$ produced this restricted observation.
• Mechanism: The model computes a denoised pixel value as the posterior mean over the training set, weighted by the posterior probability $P_{\text{train}}(\phi | C_{x,t})$. By restricting the information available to each pixel (e.g., observing only a local spatial patch), the "Bayesian guessing game" becomes harder, theoretically preventing the model from collapsing to a single training point (memorization).
• Theoretical Framework: The authors analyze the memorization-generalization transition using information theory. They derive a phase boundary in the joint space of:
  1. The amount of training data ($|\mathcal{D}|$).
  2. The time/noise level in the reverse process ($t$ or $\sigma_t$).
  3. The amount of information restriction (e.g., patch size $L$).
• Key Criterion: The transition occurs when the mutual information $I(\phi; C_{x,t})$ between the restricted observation and the true data distribution equals the logarithm of the dataset size: $\ln |\mathcal{D}| = I(\phi; C_{x,t})$.
  • Memorization Phase: $I(\phi; C_{x,t}) > \ln |\mathcal{D}|$ (The observation is too informative; the model guesses the training point perfectly).
  • Generalization Phase: $I(\phi; C_{x,t}) < \ln |\mathcal{D}|$ (The observation is insufficient to uniquely identify a training point; the model averages over possibilities).

Key Contributions

1. Introduction of BIRD Models: A general, analytically solvable framework that encompasses and generalizes existing local score models (LS), equivariant score models (ES), and equivariant-local score models (ELS). It requires no explicit architectural assumptions or learning dynamics, relying solely on Bayesian inference over a finite dataset with information restriction.
2. Prediction of Early Training Dynamics: The authors demonstrate that spatially local BIRD models (where $C_{x,t}$ is a local image patch) closely approximate the behavior of trained neural diffusion models (including UNets and Diffusion Transformers) during the early stages of training. They achieve high predictive accuracy (median $r^2 \sim 0.85-0.93$) on individual image outputs across various architectures and datasets (Celeba64, CIFAR10, MNIST, FashionMNIST).
3. Exact Phase Boundary Theory: Under minimal assumptions on data distributions, the paper derives a precise information-theoretic criterion for the memorization-generalization transition. It proves that generalization occurs if and only if the mutual information between the restricted observation and the data is less than the log-number of training points.
4. Critical Patch Scale and Spectral Alignment: For spatially local models, the authors identify a time-dependent critical patch scale $L_c(t)$. They show that optimal generation proceeds near the edge of memorization, where the patch scale $L$ tracks the critical scale $L_c$. Furthermore, they demonstrate that for natural images with power-law spectral densities ($P(k) \sim k^{-2-\epsilon}$), the critical scale aligns with the "spectral scale" (the scale of optimal linear denoising).
5. Circumventing the Curse of Dimensionality: The theory reveals that for scale-invariant images ($\epsilon=0$), BIRD models do not suffer from the curse of dimensionality; the required dataset size does not need to scale with image size. For near scale-invariant images ($\epsilon \sim 0.1-0.3$), the data requirement scales as $\exp(L_I^\epsilon)$, which is significantly sub-exponential compared to the naive $\exp(L_I^2)$ requirement.

Results

• Empirical Validation: Experiments confirm the theoretically predicted location of the phase transition. The mutual information curves derived from theory closely track the entropy deficit observed in numerical experiments on both toy datasets and real-world image datasets.
• Consistency: Both BIRD models and early-trained neural networks exhibit "consistent generalization," where models trained on disjoint subsets of data produce nearly identical outputs when given the same noise input.
• Scaling Laws: The critical patch scale $L_c(\sigma_t)$ is shown to monotonically decrease as noise decreases (time progresses in the reverse process), aligning with the observed coarse-to-fine generation process. The critical scale is found to be highly correlated with the spectral scale $L_{\text{spec}}$ derived from the data's power spectral density.
• Failure Modes: The paper illustrates that if the patch scale is too small ($L \ll L_c$), the model generalizes but performs poorly (high denoising error). If the patch scale is too large ($L > L_c$), the model memorizes the training data and fails to generalize to test data.

Significance
The paper claims to provide a fundamental insight into how generative AI circumvents the curse of dimensionality. By framing the problem through the lens of information restriction, the authors show that successful generalization in diffusion models is not merely a result of architectural inductive biases (like smoothness or low-dimensionality) but is fundamentally driven by the limitation of information available to the generative process at any given step. The results suggest that generation proceeds "at the edge of memorization," where information restriction is carefully tuned to prevent the model from collapsing to specific training examples while retaining enough information to reconstruct the underlying data distribution. This work offers a new theoretical laboratory for studying the memorization-generalization transition across a wide variety of data distributions and model classes.