SCORE-BASED CONTINUOUS-TIME DISCRETE DIFFU-SION MODELS

Abstract

Score-based modeling through stochastic differential equations (SDEs) has provided a new perspective on diffusion models, and demonstrated superior performance on continuous data. However, the gradient of the log-likelihood function, i.e., the score function, is not properly defined for discrete spaces. This makes it non-trivial to adapt the score-based modeling to categorical data. In this paper, we extend diffusion models to discrete variables by introducing a stochastic jump process where the reverse process denoises via a continuous-time Markov chain. This formulation admits an analytical simulation during backward sampling. To learn the reverse process, we extend score matching to general categorical data, and show that an unbiased estimator can be obtained via simple matching of the conditional marginal distributions. We demonstrate the effectiveness of the proposed method on a set of synthetic and real-world music and image benchmarks.

1. INTRODUCTION

Diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020) have emerged as an important technique for data distribution modeling, where a data-corrupting forward process is coupled with a denoising reverse process to simulate a diffusion relationship between the data distribution and an uninformed source. Such models admit stable learning procedures and have demonstrated superior performance on continuous data modeling in challenging scenarios (Dhariwal & Nichol, 2021) , leading to rapidly increasing popularity. Song et al. (2020) established a stochastic differential equation view of diffusion models by forming the limit of finer corruption and denoising steps in the forward and backward processes, rendering a continuum of distributions. This perspective has provided a unified framework under a new score-based learning objective, and inspired a variety of simulation methods for efficient sampling and inference (De Bortoli et al., 2021; Zhang & Chen, 2022) . Given the advantages of diffusion models in terms of flexibility, learning tractability, and sampling, there have been several attempts to extend the approach to discrete data. Recent attempts have investigated alternative corruption operations for discrete data in the forward process, yielding promising results (Hoogeboom et al., 2021b; a; Austin et al., 2021) . However, these extensions still execute a finite sequence of corruption and restoration steps, and remain restricted to a fixed reverse sampling strategy that can be sub-optimal. To overcome this limitation, we investigate whether a continuoustime discrete diffusion formulation might admit more effective estimation and generation. Such an extension is highly non-trivial however. The continuous-time diffusion framework is based on a stochastic differential equation (SDE) with respect to the score function, which itself is the gradient of the log-likelihood with respect to a continuous variable. Although this can be used to characterize a continuum of infinitesimally evolving distributions over a continuous space, such a formulation no longer exists for discrete variables, since the gradient of the log-likelihood does not exist with respect to a discrete variable. Recently, Campbell et al. (2022) made significant progress in * Work done during an internship at Google.

