SteinDiff: Resolving the Contractivity Trap via Reference-free Stein Regularization

Shigui Li, Delu Zeng

Keywords: Contractivity Trap; Reference-free Regularization; Stein's Identity; Closed-Form Optimization

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.

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 regularization. Specifically, drawing on Krasnosel'skiĭ-Mann theory, we reformulate the discretized ODE update operator to interpolate between predictions and current states. Importantly, we contribute efficient closed-form regularization estimators via Stein's identity, which is grounded in the continuous SDE theory of diffusion models. Our step-wise 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.

Primary Area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)

Submission Number: 1392