💻 computer science

DyWeight: Dynamic Gradient Weighting for Few-Step Diffusion Sampling

The paper proposes DyWeight, a lightweight, learning-based multi-step solver that employs dynamic gradient weighting and an implicit coupling paradigm to adaptively align with non-stationary diffusion dynamics, achieving superior visual fidelity and stability with significantly fewer function evaluations across various models and datasets.

Original authors: Tong Zhao, Mingkun Lei, Liangyu Yuan, Yanming Yang, Chenxi Song, Yang Wang, Beier Zhu, Chi Zhang

Published 2026-03-13

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Tong Zhao, Mingkun Lei, Liangyu Yuan, Yanming Yang, Chenxi Song, Yang Wang, Beier Zhu, Chi Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to bake a perfect cake, but you only have three minutes to do it.

Normally, baking a cake requires following a recipe step-by-step: mix, rest, bake, check, cool, etc. In the world of AI image generation (Diffusion Models), this "recipe" is called sampling. To get a high-quality image, the AI usually needs to take hundreds of tiny steps, checking its work constantly. This is like baking a cake perfectly, but it takes hours.

The Problem:
We want the cake (the image) in seconds, not hours. So, we try to take big jumps instead of tiny steps. But here's the catch: if you take giant leaps in a complex process, you often miss the details or end up with a burnt mess.

Current methods try to speed things up by using "handcrafted rules." Imagine a chef who says, "For the first jump, use 30% of the last mix; for the second, use 70%." These rules are fixed. They work okay for small steps, but when you try to take huge leaps (few steps), the rules break down because the "dough" (the image data) changes too fast and unpredictably.

The Solution: DyWeight
The paper introduces DyWeight (Dynamic Gradient Weighting). Think of DyWeight as a super-intelligent, adaptive chef who doesn't follow a static recipe book.

Here is how it works, using simple analogies:

1. The "Flexible Recipe" (Dynamic Weighting)

Old solvers are like a rigid robot: "I will always mix 50% of ingredient A and 50% of ingredient B."
DyWeight is like a human chef who tastes the batter as they go. It asks: "Wait, the mixture is changing faster than I thought. I need to rely more on the last taste I had, and less on the one before that."

Instead of fixed numbers, DyWeight learns the perfect mix ratio for every single step. It dynamically decides how much "history" (past guesses) to trust. Sometimes it trusts the recent past heavily; other times, it looks further back. This allows it to adapt to the chaotic, fast-changing nature of generating an image in just a few steps.

2. The "Time Warp" (Implicit Time Calibration)

This is the paper's cleverest trick.
Usually, when you take a big step, you have to calculate exactly when you are in the process. If you get the timing wrong, the AI gets confused (like trying to frost a cake that is still liquid).

Old methods try to calculate the time separately, which is like having a separate clock and a separate timer that sometimes disagree.
DyWeight does something magical: it hides the time calculation inside the mixing ratios.

The Analogy: Imagine you are walking up a steep hill. Instead of stopping to check your watch to see how far you've gone, you just look at how tired your legs feel. If your legs feel like you've walked 100 meters, you are at the 100-meter mark.
How it works: DyWeight adjusts the "size" of its step automatically. If the math says the step should be bigger, the weights naturally expand. If it should be smaller, they shrink. It doesn't need a separate "time clock" to tell it where it is; the way it mixes the ingredients is the clock. This keeps the AI perfectly synchronized with the image's internal rhythm, even when taking giant leaps.

3. The "Coach and Student" (Distillation)

How does this chef learn to be so good so fast?
They don't learn by trial and error on the final cake. Instead, they use a Teacher-Student system.

The Teacher: A slow, perfect chef who takes 35 steps to make a flawless cake.
The Student (DyWeight): A fast chef who only has 3 steps.

The Student watches the Teacher make the cake. But here's the secret: The Student isn't told to copy every single tiny movement the Teacher makes. The Student is only graded on the final result.
"Did you end up with a cake that looks like the Teacher's?"
This freedom allows the Student to find shortcuts. It might take a weird path in the middle that the Teacher never took, but as long as the final cake is perfect, the Student wins. This makes the learning process incredibly efficient.

The Result

In the paper, they tested this on famous AI models (like FLUX.1 and Stable Diffusion).

Old methods (taking 3-5 steps) produced blurry, weird, or distorted images.
DyWeight (taking the same 3-5 steps) produced images that were sharp, clear, and followed the instructions perfectly.

In Summary:
DyWeight is a new way to tell an AI to "draw a picture fast." Instead of using a rigid, pre-written rulebook, it gives the AI a flexible, learning-based strategy that adapts its own speed and timing on the fly. It's the difference between a robot following a broken map and a human navigator who knows exactly how to get to the destination in record time, no matter how bumpy the road is.

1. Problem Statement

Diffusion Models (DMs) have achieved state-of-the-art generative performance but suffer from prohibitively slow sampling speeds, often requiring hundreds of function evaluations (NFEs) to generate a single high-quality image. While multi-step Ordinary Differential Equation (ODE) solvers (e.g., DPM-Solver++, iPNDM) have improved efficiency by reusing historical gradients, they face two critical limitations in few-step sampling regimes:

Non-Stationary Dynamics: Existing solvers rely on handcrafted, time-invariant coefficients derived under the assumption that denoising gradients at neighboring timesteps share similar statistics. However, diffusion sampling is highly non-stationary; early noisy gradients differ significantly from later reliable ones. Fixed weighting schemes fail to adapt, leading to accumulated bias and degraded image quality, especially with large step sizes.
Rigid Optimization Paradigms: Recent learning-based solvers attempt to optimize coefficients but typically use a decoupled optimization paradigm. They treat solver coefficients and time schedules as independent variables, often requiring complex, multi-loop bi-level optimization strategies to ensure valid time sequences (e.g., strictly decreasing steps). This leads to cumbersome training and suboptimal convergence.

2. Methodology: DyWeight

The authors propose DyWeight, a lightweight, learning-based multi-step solver that introduces a streamlined implicit coupling paradigm to address the above limitations.

Core Mechanisms

Relaxed Weight Formulation (Implicit Time Shifting):
- Classical multi-step solvers constrain gradient weights to sum to one ( $\sum w_i = 1$ ). DyWeight relaxes this constraint, allowing weights to be unconstrained and learnable.
- Dual Function:
  1. The normalized part of the weights adaptively learns the optimal combination of historical gradients.
  2. The unnormalized sum acts as an implicit time shifting mechanism. It dynamically adjusts the effective step size ( $\tilde{h}_n$ ) and the macro-level integration boundary, aligning the solver's numerical trajectory with the model's internal denoising dynamics without needing explicit time parameterization.
Micro-Level Time Scaling:
- To correct residual misalignment between the numerical timeline and the network's internal noise expectations (exposure bias), DyWeight introduces a learnable per-step scalar $s_n$ . This scales the time $t$ fed into the neural network ( $\epsilon_\theta(x_t, s_n t)$ ), ensuring the model queries the correct noise level.
Training Framework (Endpoint Supervision):
- DyWeight employs a teacher-student distillation framework. A high-fidelity teacher (e.g., 35-step iPNDM) generates a trajectory, and the student (DyWeight) learns to match the final output.
- Crucially, it uses endpoint supervision (minimizing distance only at the final step $x_0$ ) rather than path supervision. This grants the solver flexibility to deviate from the teacher's intermediate path to discover more efficient trajectories that better align with the data manifold.

Algorithm Flow

Buffering: Historical gradients are buffered.
Weighting: Unconstrained learned weights aggregate these gradients.
Time Calibration: The unnormalized weight sum shifts the time step, and the scaling factor $s_n$ adjusts the query time for the network.
Update: The state is updated using the weighted gradient and the calibrated time step.

3. Key Contributions

Implicit Coupling Paradigm: DyWeight replaces complex, decoupled parameterizations with a relaxed weight formulation. By allowing weights to sum to anything other than 1, it implicitly couples gradient aggregation with timestep adjustments, transforming a difficult joint-optimization problem into a single-round learning process.
Time Calibration Strategy: The combination of implicit time shifting (via unnormalized weights) and micro-level time scaling allows the solver to maintain temporal alignment with the diffusion model's internal dynamics in few-step regimes, significantly reducing truncation errors.
Training Efficiency: The method is training-data-free (using on-the-fly teacher generation) and requires minimal optimization overhead. It converges rapidly (often within ~600 training images) due to strong inductive bias from initializing with classical coefficients (e.g., Adams-Bashforth).
State-of-the-Art Performance: Extensive experiments demonstrate that DyWeight achieves superior visual fidelity and stability with significantly fewer NFEs compared to both handcrafted and existing learning-based solvers.

4. Experimental Results

The authors evaluated DyWeight across diverse benchmarks, including pixel-space datasets (CIFAR-10, FFHQ, AFHQv2, ImageNet64) and latent-space models (LSUN-Bedroom, Stable Diffusion v1.5, FLUX.1-dev).

Quantitative Performance:
- CIFAR-10 (3 NFE): DyWeight achieves an FID of 8.16, outperforming the next best learnable method (EPD-Solver, 10.40) and handcrafted iPNDM (24.55).
- FLUX.1-dev (7 NFE): DyWeight achieves an FID of 20.69 and a CLIP score of 26.49, surpassing DPM-Solver++ and iPNDM.
- Generalization: On the challenging DrawBench prompt suite, DyWeight shows robust generalization, outperforming baselines in human preference metrics (HPSv2.1, ImageReward).
Ablation Studies:
- Removing either time shifting or time scaling significantly degrades performance, confirming the necessity of the dual time calibration mechanism.
- Using a history of $K=3$ or $4$ steps is optimal; using all available history ( $K=N$ ) introduces noise from stale gradients.
- Endpoint vs. Path Supervision: Endpoint supervision consistently outperforms path supervision in few-step settings, as it allows the solver to find "shortcuts" rather than strictly mimicking the teacher's trajectory.

5. Significance

DyWeight represents a significant advancement in efficient diffusion sampling. By challenging the classical numerical constraint that gradient weights must sum to one, it elegantly solves the problem of adapting to non-stationary diffusion dynamics without the computational burden of complex time schedule optimization.

Practical Impact: It enables high-fidelity image generation with drastically reduced inference time (fewer NFEs), making real-time and interactive applications of large diffusion models (like FLUX.1-dev) more feasible.
Theoretical Insight: The paper highlights that for few-step sampling, the destination (final data alignment) is more critical than the path (intermediate trajectory), validating the use of endpoint supervision over path supervision.
Accessibility: The method is lightweight, requires no modification to the underlying diffusion model backbone, and can be quickly distilled for new models, establishing a new state-of-the-art for efficient solvers.