arXiv:2605.06367 · v2

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

首次提交 2026-05-07 22:44

版本更新 2026-11-12 01:35

雷达观测 2026-11-13 14:16

摘要

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 and sampling imbalance play secondary roles. Specifically, sampling imbalance acts as a modulator that can reverse this ordering and, under strong imbalance, forces minority classes to acquire distinct, delayed inference 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.

作者备注：Accepted to NeurIPS 2026