Understanding and Improving Diffusion Training via Cross-Timestep Task Affinity

Non-uniform training schemes have been widely used to accelerate diffusion model training; however, most existing approaches rely on static strategies and often exhibit unstable performance. Empirical evidence suggests that prior static methods are built on the assumption that timestep importance remains constant throughout training, an assumption that is inherently limited because it fails to capture the dynamically evolving interdependencies among denoising tasks during optimization. To address this limitation, we introduce task affinity to quantify the optimization influences among denoising tasks at different timesteps. By further relating task dependencies to the training optimization state, an affinity-aware adaptive timestep sampling strategy is derived. This strategy dynamically adjusts timestep preferences and yields robust and consistent acceleration across a wide range of model architectures, model scales, and datasets.