Mitigating the Contractivity Trap in Diffusion ODEs via Stein Stabilization

A fundamental tension exists in the large-step inference of diffusion models via their deterministic probability flow ordinary differential equation (PF-ODE) trajectories, which we identify as the contractivity trap: efficient inference favors large step sizes, while aggressive steps and highly expressive denoisers can undermine contraction-based stability certificates for error suppression. To address this, we propose SteinDiff, a step-wise inference-time stabilization framework that employs Stein-derived corrections without requiring reference samples. Specifically, SteinDiff introduces a geometry-aware residual correction mechanism that regularizes large-step solver updates without retraining. To this end, we derive a closed-form Stein correction coefficient for step-wise solver adjustment, enabling reference-free adaptation to local data geometry. We further establish a score-controlled perturbation bound under distributional shifts and provide a complementary Stein perspective on EDM-style parameterizations. Extensive experiments demonstrate that SteinDiff mitigates severe artifacts and improves generative quality across large-step inference settings.

Diffusion models are powerful AI systems that generate images by gradually turning random noise into a clear picture. To make this process faster, modern systems often take much larger denoising steps. However, these large steps can make the generation process less stable, causing distorted images, visual artifacts, or loss of detail. We introduce SteinDiff, a method that helps stabilize fast diffusion sampling without retraining the original model. Instead of forcing the model to use smaller steps or changing its architecture, SteinDiff adds a principled correction during inference. This correction is computed from information already available along the ongoing generation trajectory, so it does not require reference images or additional training data. Our experiments show that SteinDiff improves image quality in fast, large-step generation across different samplers and datasets. This makes diffusion models more practical by reducing sampling cost while preserving visual quality.