Title: SteinDiff: Resolving the Contractivity Trap via Reference-free Stein Regularization
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
Keywords: Contractivity Trap; Reference-free Regularization; Stein's Identity; Closed-Form Optimization

Abstract: A fundamental tension arises when accelerating diffusion-based generative models via their deterministic probability flow ordinary differential equation (PF-ODE)  paths, which we formally identify as the *contractivity trap*: efficient inference requires large step sizes, but stable convergence demands strong contractivity that limits model expressiveness. This results in error accumulation in inference as contractivity weakens. In this work, we propose a principled inference approach, called *SteinDiff*, that relaxes the contractivity constraints through reference-free Stein adaptive regularization. Specifically, drawing on Krasnosel'skiĭ-Mann theory, we reformulate the ODE update to interpolate between predictions and current states.  Importantly, we contribute closed-form optimal regularization parameters via Stein's identity,  which is theoretically grounded in the continuous SDEs theory of  diffusion models. Our analytical approach eliminates the need for ground truth data to adapt to the local geometry of the data distribution while preserving the expressiveness of the vanilla model. Theoretically, our approach not only relaxes the strict contractivity requirements for robust convergence but also reveals a principle behind the stability of state-of-the-art (SOTA) pre-conditioned parameterizations. Practically, we offer a reference-free solution that reduces the risk of mode collapse in large-step inference. Extensive experiments validate our theoretical framework and demonstrate significant gains in generative inference.

TL;DR: We resolve the fundamental tension between stability and expressiveness in accelerating diffusion models by introducing a principled method, SteinDiff, that leverages closed-form adaptive regularization to achieve accurate and robust generation.