Beyond Log-Concavity and Score Regularity: Improved Convergence Bounds for Score-Based Generative Models in W2-distance

TL;DR

This paper introduces a new theoretical framework for analyzing the convergence of score-based generative models in Wasserstein-2 distance, relaxing previous assumptions like log-concavity and score regularity, and leveraging properties of the Ornstein-Uhlenbeck process.

Contribution

It provides a PDE-based analysis showing how weak log-concavity evolves into log-concavity over time, broadening the applicability of convergence bounds for SGMs.

Findings

Weak log-concavity becomes log-concavity over time due to OU process.

The drift of the reversed process alternates between contractive and non-contractive regimes.

The framework applies to Gaussian mixture models, demonstrating versatility.

Abstract

Score-based Generative Models (SGMs) aim to sample from a target distribution by learning score functions using samples perturbed by Gaussian noise. Existing convergence bounds for SGMs in the W2-distance rely on stringent assumptions about the data distribution. In this work, we present a novel framework for analyzing W2-convergence in SGMs, significantly relaxing traditional assumptions such as log-concavity and score regularity. Leveraging the regularization properties of the Ornstein--Uhlenbeck (OU) process, we show that weak log-concavity of the data distribution evolves into log-concavity over time. This transition is rigorously quantified through a PDE-based analysis of the Hamilton--Jacobi--Bellman equation governing the log-density of the forward process. Moreover, we establish that the drift of the time-reversed OU process alternates between contractive and non-contractive regimes, reflecting the dynamics of concavity. Our approach circumvents the need for stringent regularity conditions on the score function and its estimators, relying instead on milder, more practical assumptions. We demonstrate the wide applicability of this framework through explicit computations on Gaussian mixture models, illustrating its versatility and potential for broader classes of data distributions.