Error as Signal: Stiffness-Aware Diffusion Sampling via Embedded Runge-Kutta Guidance
Inho Kong*, Sojin Lee*, Youngjoon Hong†, Hyunwoo J. Kim†
1
MLV Lab
2
Background
Gaussian Noise
Generated Data
Diffusion ODE formulation
Diffusion sampling can be formulated as solving an ordinary differential equation (ODE).
Sampling errors stem from two main sources: model prediction and integration.
MLV Lab
3
Model-induced Error
Solver-induced Error�
Overview
Inaccurate Score Estimation
Models are trained to estimate the score function via score matching loss.
Inaccurate ODE Integration�We rely on numerical ODE solvers since the diffusion ODE lacks an exact solution.
Extrapolate outputs from two models
(e.g., CFG, AG)
Extrapolate outputs from two solvers
(Our Method)
Target Problem
Referred to as
Mitigation Strategy
Conventional guidance methods primarily address model-induced errors.
ERK-Guid directly targets solver-induced local truncation errors to improve generation quality.
MLV Lab
4
Key Intuition (Drift Field)
Stiffness in diffusion ODEs arises when drift directions change rapidly.
A 2D drift field illustrates stiff and non-stiff regions.
Drift Field
non-stiff region
stiff region
MLV Lab
5
Key Intuition (Local Truncation Error)
Local Truncation Error (LTE)
Drift Field
In these stiff regions, a significantly larger LTE occurs during sampling.
MLV Lab
6
Key Intuition (Eigenvector Alignment)
Local Truncation Error (LTE)
Dominant Eigenvector Direction
We observe that in stiff regions, the LTE closely aligns with the dominant eigenvector of the Jacobian.
This alignment provides a reliable directional proxy to correct solver-induced errors.
MLV Lab
7
Method Formulation (Cost-free Estimators)
Calculating exact stiffness and eigenvectors via Jacobian-vector products is inefficient for practical sampling.
We leverage Embedded Runge-Kutta (ERK) pairs to obtain cost-free estimators.
Embedded Runge-Kutta pair
Stiffness Estimator
Dominant Eigenvector Estimator
MLV Lab
8
Method Formulation (ERK-Guid)
Adaptive Guidance Scale�determined by stiffness and hyperparameters
Step-size
Guidance Vector�drift differences between distinct solver states
MLV Lab
9
Experiments
Accuracy of the estimators
Our estimated stiffness has a strong correlation with the exact stiffness.
MLV Lab
10
Experiments
Accuracy of the estimators
The cosine similarity between ERK-Guid and the eigenvector is notably high in high-stiffness regions.
These results demonstrate that our estimator is sufficiently accurate.
MLV Lab
11
Experiments
ERK-Guid consistently improves all sample quality metrics without any extra inference cost.
Notably, the performance gains are most dramatic at low sampling steps where solver errors dominate.
Effectiveness of ERK-Guid
MLV Lab
12
Experiments
ERK-Guid provides an orthogonal guidance signal with existing methods such as CFG and Autoguidance. Integrating our method improves sample quality metrics across different sampling steps.
Guidance compatibility
MLV Lab
13
Experiments
ERK-Guid can be integrated into numerical solvers such as Heun, DPM-Solver, and DEIS.
Applying our method as a plug-and-play module improves generation quality across baseline solvers.
Plug-and-play adaptation
MLV Lab
14
Experiments
Comparison with adaptive step-size control
Conventional adaptive solvers resize the step size when stiffness exceeds a specific threshold.
While this approach improves performance, ERK-Guid achieves better sample quality without extra NFEs.
MLV Lab
15
Experiments
ERK-Guid calculates the guidance vector without additional network evaluations.
As a result, it introduces minimal inference overhead and maintains the same memory consumption
Wall-clock time (seconds per image) and Memory consumption
MLV Lab
16
Experiments
Qualitative result demonstrates how ERK-Guid helps reduce structural distortions.
It generates an image with much better preserved details than the baseline.
Qualitative result
DPM Solver
DPM Solver + ERK-Guid
MLV Lab
17
Summary
MLV Lab