Geometry of Score Based Generative Models

TL;DR

This paper presents a geometric perspective on score-based generative models, proving their processes are Wasserstein gradient flows, and introduces a projection step to accelerate sampling for high-quality image generation.

Contribution

It is the first to establish the connection between score-based models and Wasserstein gradient flows, offering new insights and a geometric method to reduce sampling steps.

Findings

Proved the forward and backward processes are Wasserstein gradient flows.

Introduced a projection step to accelerate sampling.

Achieved high-quality image generation with fewer steps.

Abstract

In this work, we look at Score-based generative models (also called diffusion generative models) from a geometric perspective. From a new view point, we prove that both the forward and backward process of adding noise and generating from noise are Wasserstein gradient flow in the space of probability measures. We are the first to prove this connection. Our understanding of Score-based (and Diffusion) generative models have matured and become more complete by drawing ideas from different fields like Bayesian inference, control theory, stochastic differential equation and Schrodinger bridge. However, many open questions and challenges remain. One problem, for example, is how to decrease the sampling time? We demonstrate that looking from geometric perspective enables us to answer many of these questions and provide new interpretations to some known results. Furthermore, geometric perspective enables us to devise an intuitive geometric solution to the problem of faster sampling. By augmenting traditional score-based generative models with a projection step, we show that we can generate high quality images with significantly fewer sampling-steps.