Search papers, labs, and topics across Lattice.
The paper introduces Dynamic Gradient Weighting (DyWeight), a learned multi-step ODE solver for diffusion model sampling that adaptively weights historical gradients. DyWeight relaxes numerical constraints and learns time-varying parameters to align the solver's trajectory with the denoising dynamics, enabling larger integration steps. Experiments across various datasets and diffusion models demonstrate that DyWeight achieves state-of-the-art visual fidelity and stability with fewer function evaluations compared to existing efficient diffusion solvers.
Ditch the handcrafted coefficients: DyWeight learns how to dynamically weight gradients in diffusion model sampling, slashing compute while boosting image quality.
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