Local Curvature Smoothing with Stein's Identity for Efficient Score Matching

TL;DR

This paper introduces LCSS, a novel score matching method using Stein's identity that reduces computational cost and improves sample quality in diffusion models, especially at high resolutions.

Contribution

LCSS is a new score matching approach that bypasses Jacobian trace computation, enhancing efficiency and stability in training score-based diffusion models.

Findings

LCSS outperforms existing methods in sample quality metrics.

LCSS matches denoising score matching performance.

Enables realistic high-resolution image generation at 1024x1024.

Abstract

The training of score-based diffusion models (SDMs) is based on score matching. The challenge of score matching is that it includes a computationally expensive Jacobian trace. While several methods have been proposed to avoid this computation, each has drawbacks, such as instability during training and approximating the learning as learning a denoising vector field rather than a true score. We propose a novel score matching variant, local curvature smoothing with Stein's identity (LCSS). The LCSS bypasses the Jacobian trace by applying Stein's identity, enabling regularization effectiveness and efficient computation. We show that LCSS surpasses existing methods in sample generation performance and matches the performance of denoising score matching, widely adopted by most SDMs, in evaluations such as FID, Inception score, and bits per dimension. Furthermore, we show that LCSS enables realistic image generation even at a high resolution of $1024 \times 1024$.