DyWeight: Dynamic Gradient Weighting for Few-Step Diffusion Sampling

TL;DR

DyWeight is a novel learning-based multi-step solver for diffusion models that adaptively combines historical gradients, enabling faster sampling with fewer evaluations while maintaining high visual quality.

Contribution

Introduces DyWeight, a lightweight, implicit multi-step solver that learns adaptive parameters to improve diffusion sampling efficiency and accuracy.

Findings

Achieves superior visual fidelity with fewer function evaluations.

Demonstrates state-of-the-art efficiency across multiple datasets.

Maintains stability and quality with large integration steps.

Abstract

Diffusion Models (DMs) have achieved state-of-the-art generative performance across multiple modalities, yet their sampling process remains prohibitively slow due to the need for hundreds of function evaluations. Recent progress in multi-step ODE solvers has greatly improved efficiency by reusing historical gradients, but existing methods rely on handcrafted coefficients that fail to adapt to the non-stationary dynamics of diffusion sampling. To address this limitation, we propose Dynamic Gradient Weighting (DyWeight), a lightweight, learning-based multi-step solver that introduces a streamlined implicit coupling paradigm. By relaxing classical numerical constraints, DyWeight learns unconstrained time-varying parameters that adaptively aggregate historical gradients while intrinsically scaling the effective step size. This implicit time calibration accurately aligns the solver's numerical trajectory with the model's internal denoising dynamics under large integration steps, avoiding complex decoupled parameterizations and optimizations. Extensive experiments on CIFAR-10, FFHQ, AFHQv2, ImageNet64, LSUN-Bedroom, Stable Diffusion and FLUX.1-dev demonstrate that DyWeight achieves superior visual fidelity and stability with significantly fewer function evaluations, establishing a new state-of-the-art among efficient diffusion solvers. Code is available at https://github.com/Westlake-AGI-Lab/DyWeight