Soft Diffusion: Score Matching for General Corruptions

TL;DR

This paper introduces Soft Score Matching, a new objective for learning score functions in general diffusion models, achieving state-of-the-art results on CelebA with improved efficiency.

Contribution

It proposes a unified framework for general linear corruption processes, including a novel loss, corruption level selection, and sampling method, advancing diffusion model capabilities.

Findings

State-of-the-art FID score 1.85 on CelebA-64

Effective for Gaussian blur and masking processes

Reduced computational costs compared to traditional diffusion models

Abstract

We define a broader family of corruption processes that generalizes previously known diffusion models. To reverse these general diffusions, we propose a new objective called Soft Score Matching that provably learns the score function for any linear corruption process and yields state of the art results for CelebA. Soft Score Matching incorporates the degradation process in the network. Our new loss trains the model to predict a clean image, \textit{that after corruption}, matches the diffused observation. We show that our objective learns the gradient of the likelihood under suitable regularity conditions for a family of corruption processes. We further develop a principled way to select the corruption levels for general diffusion processes and a novel sampling method that we call Momentum Sampler. We show experimentally that our framework works for general linear corruption processes, such as Gaussian blur and masking. We achieve state-of-the-art FID score $1.85$ on CelebA-64, outperforming all previous linear diffusion models. We also show significant computational benefits compared to vanilla denoising diffusion.