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

Abstract: Real-world datasets are inherently heterogeneous, yet how per-class structural differences and sampling imbalance shape the training dynamics of diffusion models-and potentially exacerbate disparities-remains poorly understood. While models typically transition from an initial phase of generalization to memorizing the training set, existing theory assumes homogeneous data, leaving open how class imbalance and heterogeneity reshape these dynamics. In this work, we develop a high-dimensional analytical framework to study class-dependent learning in score-based diffusion models. Analyzing a random-features model trained on Gaussian mixtures, we derive the feature-covariance spectrum to characterize per-class generalization and memorization times. We reveal the explicit hierarchy governing these dynamics: class variance is the primary determinant of learning order-consistently favoring higher-variance classes-while centroid geometry plays a secondary role. Sampling imbalance acts as a modulator that can reverse this ordering and, under strong imbalance, forces minority classes to acquire distinct, delayed speciation times during backward diffusion. Together, these results suggest that diffusion models can memorize some classes while others remain insufficiently learned. We validate our theoretical predictions empirically using U-Net models trained on Fashion MNIST.

• The paper provides an explicit high-dimensional framework that derives class-resolved learning timescales using Gaussian-mixture data and a Random Features model.
• It reveals that class variance, centroid geometry, and sampling imbalance drive distinct learning phases, governing speciation, generalization, and memorization in diffusion models.
• Empirical validations using DDPMs on Fashion MNIST confirm that reweighted class sampling modulates memorization gaps while exposing a trade-off between generalization and memorization.

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

Introduction

This paper (2605.06367) proposes a rigorous high-dimensional theoretical framework to dissect how structural heterogeneities and sample imbalance in heterogeneous datasets impact the class-dependent generalization and memorization dynamics of denoising diffusion models. By focusing on score-based models trained on Gaussian-mixture data using a Random Features (RF) architecture, the authors derive explicit, class-resolved timescales for learning and memory formation, and validate these predictions with empirical training of U-Net-based DDPMs on Fashion MNIST.

Theoretical Framework and Problem Setting

Classical analyses of overparameterized diffusion models consider homogeneous data distributions and consequently predict a global separation between generalization and memorization phases, governed by the feature-covariance spectrum. This work instead explicitly models the data distribution as a heterogeneous Gaussian Mixture Model (GMM) with class-wise variances, centroids, and mixing proportions. The diffusion score is approximated by an RF model, enabling exact gradient flow dynamics for the readout, determined analytically in terms of the feature-feature covariance matrix $\mathbf{U}$ and its spectral decomposition (see Section 2).

By leveraging techniques from random matrix theory and Gaussian equivalence, the analysis yields closed-form, high-dimensional-limit expressions for the eigenspectrum of $\mathbf{U}$ that capture the impact of class-conditional variance, centroid geometry, and imbalance.

Spectral Transitions: Speciation, Generalization, and Memorization

The spectrum of $\mathbf{U}$ is shown to drive distinct qualitative phases during backward diffusion and finite-time score learning. The emergence of isolated eigenvalues signals “speciation,” at which the model acquires the ability to infer distinct class modes. Under weak imbalance ($b_c = \mathcal{O}(1)$), all classes speciate simultaneously, but under strong imbalance ($b_c = \mathcal{O}(N^{-a_c})$ with $0 < a_c < 1$), speciation times become class-specific and hierarchically separated, with minority classes speciation increasingly delayed (see Result 1, Section 3.1).

Figure 1: Strong imbalance induces separation in speciation times; minority classes become inferable only later in the backward diffusion process.

The generalization and memorization phases in training time are linked to bulk extremal eigenvalues: the rightmost edge marks onset of generalization, central and leftmost bulks correspond to the onset and completion of memorization, respectively. Strong class variance heterogeneity ($\sigma_c^2 \neq \sigma_{c'}^2$) contracts the generalization window and expands the memorization window. Sample imbalance modulates this effect, with skewed sampling either mitigating or accentuating the contraction depending on which class is undersampled.

Figure 2: Theoretical and empirical spectral densities and loss curves illustrate correspondence between bulk edges and macroscopic learning phases.

Hierarchy and Modulation of Class-Specific Learning Dynamics

A central result is the explicit hierarchical ordering governing per-class generalization/memorization:

• Class variance: High-variance classes are learned and memorized first.
• Centroid geometry: For comparable variances, classes with smaller-norm centroids are prioritized.
• Sampling imbalance: Usually subleading, but can override intrinsic ordering in the strong imbalance regime.

This is confirmed by explicit tracking of class-conditional test errors and their optimal generalization/memorization times. The gaps between classes in these times can be arbitrarily amplified or reversed via imbalance, which delays generalization for minority classes yet can accelerate their memorization if majority classes have dominant variance.

Figure 3: Disaggregated class-conditional errors display the effect of variance, centroid, and imbalance perturbations on training trajectories.

Figure 4: Gap analysis shows how memorization and generalization order can be controlled by sampling weight.

This structure induces an unavoidable generalization-memorization trade-off: closing the generalization gap by reweighting in favor of a minority class comes at the cost of widening the memorization disparity, and vice versa.

Empirical Validation: Diffusion Models on Fashion MNIST

The theoretical findings are validated using standard U-Net DDPMs trained on Fashion MNIST under controlled class-pairing and imbalance and compared to the RF theory predictions using Gaussian moment statistics. Per-class memorization times—evaluated through nearest-neighbor criteria—demonstrate that high-variance classes (e.g., Bag) are memorized earlier than low-variance ones (e.g., Sneaker), and the observed memorization gaps are quantitatively consistent with the predictions from synthetic GMMs with matched empirical statistics.

Figure 5: DDPM memorization fractions confirm that high-variance classes are systematically memorized first and that the gap scales with empirical variance.

Varying sampling proportions, the empirical memorization gap between pairs is controlled in alignment with the theory: increasing the sampling weight of the high-variance class systematically reduces the gap.

Figure 6: DDPM–RF model comparisons reveal near-identical descriptor-dependent memorization gap patterns across class pairs.

Implications and Future Directions

Theoretically, this framework establishes generalization and memorization as fundamentally class-specific phenomena in structured, imbalanced data, driven primarily by variance hierarchy and secondarily by centroid and sampling geometry. A single early-stopping time in overparameterized diffusion models cannot regularize all classes equitably. Practically, this suggests that mitigation of per-class disparities—especially for minority or low-signal classes—requires explicit class reweighting or data augmentation.

The analysis is currently restricted to isotropic binary mixtures and first/second-order geometry, but generalizes naturally to multiclass settings and can, in principle, accommodate non-isotropic or higher-order structure. Key open questions concern the impact of multi-modal/intra-class heterogeneity and the effect of outliers in the data landscape on generalization/memorization ordering.

Conclusion

By providing an exact, high-dimensional treatment of class-aware learning dynamics in score-based diffusion models, this work rigorously characterizes the interplay of data structure and imbalance at all stages of generative modeling. It demonstrates both theoretically and empirically that class variance is primary in setting learning order, with centroid geometry and sample weighting as critical modulators, and shows that these insights carry over robustly to real image data and neural architectures. Such foundational understanding is directly actionable for developing bias-aware training protocols and architectures in generative modeling.