MCMC Correction of Score-Based Diffusion Models for Model Composition

TL;DR

This paper introduces a new way to improve sampling in score-based diffusion models by combining them with MCMC techniques, enabling better performance in model composition.

Contribution

A novel MH-like acceptance rule for score-based diffusion models using line integration of the score function is introduced.

Findings

The proposed MH-like samplers achieve similar improvements in sampling quality as energy-based models.

Existing score-based diffusion models can be reused with MCMC techniques without explicit energy parameterization.

Abstract

Diffusion models can be parameterized in terms of either score or energy function. The energy parameterization is attractive as it enables sampling procedures such as Markov Chain Monte Carlo (MCMC) that incorporates a Metropolis–Hastings (MH) correction step based on energy differences between proposed samples. Such corrections can significantly improve sampling quality, particularly in the context of model composition, where pre-trained models are combined to generate samples from novel distributions. Score-based diffusion models, on the other hand, are more widely adopted and come with a rich ecosystem of pre-trained models. However, they do not, in general, define an underlying energy function, making MH-based sampling inapplicable. In this work, we address this limitation by retaining score parameterization and introducing a novel MH-like acceptance rule based on line integration of the score function. This allows the reuse of existing diffusion models while still combining the reverse process with various MCMC techniques, viewed as an instance of annealed MCMC. Through experiments on synthetic and real-world data, we show that our MH-like samplers yield relative improvements of similar magnitude to those observed with energy-based models, without requiring explicit energy parameterization.

Figures