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

arXiv: https://arxiv.org/abs/2605.06367 · HTML: https://mlsanity.com/paper/2605.06367/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 20 of its 61 references that are in the map.
- First posted on arXiv 7 May 2026.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 20 of 20)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Decoupled Weight Decay Regularization](https://mlsanity.com/paper/1711.05101/) (2017)
- [Diffusion Models already have a Semantic Latent Space](https://mlsanity.com/paper/2210.10960/) (2022)
- [On Memorization in Diffusion Models](https://mlsanity.com/paper/2310.02664/) (2023)
- [Breaking the Curse of Dimensionality: Diffusion Models Efficiently Learn Low-Dimensional Distributions](https://mlsanity.com/paper/2409.02426/) (2024)
- [The statistical thermodynamics of generative diffusion models: Phase transitions, symmetry breaking and critical instability](https://mlsanity.com/paper/2310.17467/) (2023)
- [Analysis of learning a flow-based generative model from limited sample complexity](https://mlsanity.com/paper/2310.03575/) (2023)
- [Losing dimensions: Geometric memorization in generative diffusion](https://mlsanity.com/paper/2410.08727/) (2024)
- [An Analytical Theory of Spectral Bias in the Learning Dynamics of Diffusion Models](https://mlsanity.com/paper/2503.03206/) (2025)
- [Dynamical Decoupling of Generalization and Overfitting in Large Two-Layer Networks](https://mlsanity.com/paper/2502.21269/) (2025)
- [Bigger Isn't Always Memorizing: Early Stopping Overparameterized Diffusion Models](https://mlsanity.com/paper/2505.16959/) (2025)
- [Denoising Score Matching with Random Features: Insights on Diffusion Models from Precise Learning Curves](https://mlsanity.com/paper/2502.00336/) (2025)
- [Stage-wise Dynamics of Classifier-Free Guidance in Diffusion Models](https://mlsanity.com/paper/2509.22007/) (2025)
- [Non-Uniform Diffusion Models](https://mlsanity.com/paper/2207.09786/) (2022)
- [A solvable model of learning generative diffusion: theory and insights](https://mlsanity.com/paper/2501.03937/) (2025)
- [Exploring Bias in over 100 Text-to-Image Generative Models](https://mlsanity.com/paper/2503.08012/) (2025)
- [Generalization Dynamics of Linear Diffusion Models](https://mlsanity.com/paper/2505.24769/) (2025)
- [A solvable high-dimensional model where nonlinear autoencoders learn structure invisible to PCA while test loss misaligns with generalization](https://mlsanity.com/paper/2602.10680/) (2026)

## Built on by (most cited first, 2 of 2)

- [Weighting Schedules Govern What and When Score-Based Generative Models Learn from Multimodal Data](https://mlsanity.com/paper/2609.35322/) (2026)
- [When Does a Diffusion Model Decide What to Draw ?](https://mlsanity.com/paper/2610.05645/) (2026)

## Abstract (arXiv)

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.
