DyWeight: Dynamic Gradient Weighting for Few-Step Diffusion Sampling

2026-03-12

AI Summary

This work addresses the inefficiency of diffusion model sampling, which stems from its inherently slow iterative process, and the limited adaptability of existing multi-step solvers that rely on handcrafted coefficients and struggle to capture the non-stationary dynamics of diffusion. To overcome this, we propose DyWeight—a lightweight, learning-based multi-step ODE solver that adaptively aggregates historical gradients and dynamically adjusts the effective step size through an implicit coupling paradigm, aligning the numerical trajectory under large steps with the true denoising dynamics. By introducing learnable, unconstrained time-varying parameters, DyWeight eliminates conventional numerical constraints and complex decoupled optimization, achieving both structural simplicity and precise adaptation to diffusion dynamics. Experiments demonstrate that DyWeight significantly reduces the number of function evaluations while improving generation quality and stability across benchmarks including CIFAR-10, ImageNet64, and Stable Diffusion, establishing a new state of the art in efficient diffusion sampling.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Deep Generative Models & Autoencoders

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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

Problem

• Diffusion Models
• Sampling Efficiency
• Multi-step ODE Solvers
• Non-stationary Dynamics
• Gradient Weighting

Innovation

• Dynamic Gradient Weighting
• Diffusion Sampling
• Multi-step ODE Solver
• Implicit Coupling
• Time-varying Parameters