Beyond Scores: Proximal Diffusion Models

TL;DR

This paper introduces Proximal Diffusion Models (ProxDM), a novel approach that replaces score-based reverse SDE discretization with proximal maps, leading to faster convergence and theoretical guarantees in generative modeling.

Contribution

The paper proposes a new backward discretization method for diffusion models using proximal maps, with theoretical analysis and empirical evidence of improved convergence speed.

Findings

ProxDM achieves faster convergence with fewer sampling steps.

Theoretical guarantee of $ ilde{O}(d/\sqrt{\varepsilon})$ steps for $\varepsilon$-accuracy.

Empirical results show significant speedup over traditional score-based methods.

Abstract

Diffusion models have quickly become some of the most popular and powerful generative models for high-dimensional data. The key insight that enabled their development was the realization that access to the score -- the gradient of the log-density at different noise levels -- allows for sampling from data distributions by solving a reverse-time stochastic differential equation (SDE) via forward discretization, and that popular denoisers allow for unbiased estimators of this score. In this paper, we demonstrate that an alternative, backward discretization of these SDEs, using proximal maps in place of the score, leads to theoretical and practical benefits. We leverage recent results in proximal matching to learn proximal operators of the log-density and, with them, develop Proximal Diffusion Models (ProxDM). Theoretically, we prove that $\widetilde{O}(d/\sqrt{\varepsilon})$ steps suffice for the resulting discretization to generate an $\varepsilon$-accurate distribution w.r.t. the KL divergence. Empirically, we show that two variants of ProxDM achieve significantly faster convergence within just a few sampling steps compared to conventional score-matching methods.