Technical Summary: The Interplay of Data Structure and Imbalance in the Learning Dynamics of Diffusion Models

Problem Statement

Real-world datasets are inherently heterogeneous, characterized by varying per-class structural differences (such as intrinsic variance and centroid geometry) and sampling imbalances. While existing theoretical frameworks for score-based diffusion models characterize the transition from generalization to memorization, they largely assume homogeneous data distributions. Consequently, it remains poorly understood how class-specific structural properties and sampling frequencies shape the training dynamics of diffusion models, potentially leading to unequal learning quality or the exacerbation of disparities where some classes are memorized while others remain insufficiently learned.

Methodology

The authors develop a high-dimensional analytical framework to study class-dependent learning in score-based diffusion models. The core methodology involves:

1. Model Architecture: The study utilizes a Random Feature (RF) model to approximate the score function. The model is trained via denoising score matching on a dataset sampled from a Gaussian Mixture Model (GMM) with $C$ components, where each class $c$ is defined by a centroid $m_c$, variance $\sigma_c^2$, and mixing proportion $b_c$.
2. Theoretical Framework:
   • Gaussian Equivalence Principle (GEP): The authors leverage GEP techniques to decompose the data into a low-dimensional structured component and high-dimensional Gaussian fluctuations. This allows for the derivation of an asymptotically equivalent formulation for the empirical feature-feature covariance matrix $U$.
   • Random Matrix Theory (RMT): In the high-dimensional proportional limit ($N, P, M \to \infty$ with fixed ratios), the learning dynamics are analyzed through the eigenspectrum of the feature covariance matrix $U$. The authors derive the spectral density and identify critical eigenvalues that dictate convergence timescales.
   • Dynamical Analysis: The training dynamics are solved exactly as a linear Ordinary Differential Equation (ODE) for the readout weights. The authors analyze two distinct time domains:
     • Backward Diffusion Time ($t$): To determine when specific data structures (class centroids) become inferable (speciation times).
     • Training Time ($\tau$): To determine the onset of generalization and the transition to memorization based on the spectral edges of $U$.
3. Empirical Validation: The theoretical predictions are validated using:
   • Numerical simulations of the RF model on synthetic GMM data.
   • Training of standard U-Net Diffusion Probabilistic Models (DDPMs) on the Fashion MNIST dataset, specifically pairing the "Sneaker" class (lowest variance) with other classes under varying sampling proportions.

Key Contributions and Results

1. Class-Specific Speciation Times (Result 1)

The paper establishes that the "speciation time"—the point in backward diffusion where a class becomes inferable—depends on the interplay between class variance, centroid norm, and sampling weight.

• Weak Imbalance: When classes are proportionally represented, all classes speciate simultaneously.
• Strong Imbalance: Under severe undersampling ($b_c = O(N^{-a_c})$), the system exhibits a hierarchy of speciation times. Minority classes are delayed, becoming inferable only at distinct, later diffusion times (closer to $t=0$). This creates a sequential inference process where overrepresented classes are resolved before minority classes.

2. Hierarchy of Learning Dynamics (Result 2)

The authors identify a strict hierarchy governing the order in which classes are generalized and memorized during training:

• Primary Determinant: Class Variance. High-variance classes present smoother and larger score targets, leading to significantly faster generalization and memorization compared to low-variance classes.
• Secondary Determinant: Centroid Geometry. Classes with smaller-norm centroids are prioritized over those with large-norm centroids.
• Modulator: Sampling Imbalance. While usually a subleading effect, strong imbalance can act as a modulator capable of reversing the learning order established by variance and geometry.

3. The Generalization-Memorization Trade-off (Result 3)

The study reveals a fundamental trade-off driven by sample scarcity:

• Scarcity delays generalization: Undersampling a high-variance class forces the model to generalize the low-variance class first.
• Scarcity accelerates memorization: Conversely, oversampling a high-variance class deprives the low-variance class of data, causing the low-variance class to be memorized first (i.e., overfitting occurs earlier for the minority class).
• Implication: Adjusting sampling weights to close the generalization gap between classes inevitably widens the memorization gap, and vice versa.

4. Consistency with Real Architectures (Result 4)

Theoretical predictions derived from the RF model on synthetic Gaussian data show strong qualitative agreement with U-Net DDPMs trained on Fashion MNIST.

• The memorization gap trends observed in real image data (e.g., the "Bag" class vs. "Sneaker") align with the theoretical predictions based on extracted class moments (variance, centroid norm, and cosine similarity).
• This suggests that a Gaussian structural representation captures the leading-order effects of the memorization hierarchy in real-world image data, despite the complexity of non-Gaussian distributions.

Significance and Claims

The paper claims to provide the first solvable analytical framework that explicitly derives how the combination of class imbalance and cluster variance dictates distinct, class-specific critical times for both learning and memorization in diffusion models.

• Mechanistic Understanding: The work bridges the gap between abstract statistical mechanics and practical generative modeling, explaining why diffusion models might memorize one population while failing to learn another.
• Practical Implication: The authors suggest that a single early-stopping time cannot regularize all classes equally under heterogeneous data. Since classes with higher variance or smaller centroid norms are learned earlier, reweighting or rebalancing toward lower-signal classes may be necessary to mitigate per-class disparities.
• Limitations: The authors modestly note that their framework focuses on binary Gaussian mixtures with isotropic covariances and captures only first- and second-order structures. They acknowledge that more complex heterogeneities (e.g., multimodal classes, outliers) and non-isotropic settings remain open questions for future work.

The paper concludes that while variance is the dominant force in learning speed, centroid norms and class imbalance play crucial roles, becoming decisive factors in dictating the learning sequence when class variances are comparable.