Riemannian Diffusion Schrödinger Bridge

In Euclidean spaces, Score-based Generative Modeling (SGM) (Song & Ermon, 2019; Song et al., 2021b) consists of two main components. The first is a forward noising process $(\mathbf{X}_{t})_{t\geq 0}$ defined via the stochastic differential equation (SDE) (1) and the initial distribution $\mathbf{X}_{0}\sim p_{\textup{data}}$, targeting an easy-to-sample prior $p_{\textup{prior}}$. The second component is a backward denoising process $(\mathbf{Y}_{t})_{t\geq 0}=(\mathbf{X}_{T-t})_{t\in\left[0,T\right]}$ defined by the time-reversal of the noising SDE (1) (Cattiaux et al., 2021; Haussmann & Pardoux, 1986) from $p_{\textup{prior}}=p_{T}$ to $p_{\textup{data}}=p_{0}$ . Here $f$ is the drift, $g$ is the time-dependent volatility, while $(\mathbf{B}_{t})_{t\geq 0}$ is a d-dimensional Brownian motion and $\mathbb{P}_{t}$ is the distribution of $\mathbf{X}_{t}$ with corresponding density $p_{t}$. The denoising process defines a generative model by sampling $\mathbf{Y}_{0}\sim p_{\textup{prior}}$

The intractable score term, $\nabla\log p_{T-t}(\mathbf{Y}_{t})$, may be approximated by using the following score matching identity $\textstyle{\nabla_{x_{t}}\log p_{t}(x_{t})=\int_{\mathbb{R}^{d}}\nabla_{x_{t}}\log p_{t|0}(x_{t}|x_{0})~{}p_{0|t}(x_{0}|x_{t})\mathrm{d}x_{0}}$, with tractable transition density $p_{t|0}$. This is then used to train a neural network $\bm{s}_{\theta}:\mathbb{R}\times\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ with regression objective $\bm{s}_{\theta^{*}}=\arg\min_{\theta}\mathbb{E}_{\mathbf{X}_{t},\mathbf{X}_{0}}\|\nabla_{x}\log p_{t|0}(\mathbf{X}_{t}|\mathbf{X}_{0})-\bm{s}_{\theta}(t,\mathbf{X}_{t})\|^{2}$. Once trained, one can generate samples which are approximately distributed according to $p_{\textup{data}}$ via (1) (e.g. considering the Euler–Maruyama discretization of this SDE) by substituting $\bm{s}_{\theta^{*}}(t,x_{t})\approx\nabla_{x_{t}}\log p_{t}(x_{t})$.

De Bortoli et al. (2022) extended SGM to compact Riemannian manifolds, denoted $\mathcal{M}$, henceforth abbreviated RSGM for Riemannian Score-based Generative Modeling. Given a $\mathcal{M}$-valued diffusion process (2, left), where $\mathbf{B}_{t}^{\mathcal{M}}$ denotes Brownian motion on $\mathcal{M}$; the time reversal process may be written as (2, right) (De Bortoli et al., 2022, Theorem 1)

For a more thorough background on Riemannian geometry and time-reversal on manifolds see De Bortoli et al. (2022, App. B and G). Simulating a diffusion on a manifold is crucial to this approach. In contrast to Euclidean SGMs, closed-form sampling schemes are generally not available for diffusions on manifolds, hence approximation schemes such as the Geodesic Random Walk from Algo. 1 are required. Note that $\mathbf{B}_{t}^{\mathcal{M}}$ may be simulated using Algo 1 for $g=1,f=0$. This consists of applying an Euler–Maruyama step in the tangent space, whereby Gaussian noise has also been projected to the tangent space using the projection operator $\mathrm{P}$, then projecting back to the manifold using the exponential mapping.

Score matching may be used to approximate the time-reversal (2) in the Riemannian setting (De Bortoli et al., 2022). Unlike for Euclidean SGM, $\nabla\log p_{t|0}(x_{t}|x_{0})$ is not available in closed-form. Instead, one may use the score identity ¹¹1Here all gradients are considered w.r.t. the Riemannian metric of $\mathcal{M}$. $\textstyle{\nabla_{x_{t}}\log p_{t}(x_{t})=\int_{\mathcal{M}}\nabla_{x_{t}}\log p_{t|s}(x_{t}|x_{s})\mathbb{P}_{s|t}(x_{t},\mathrm{d}x_{s})}$ for $s\approx t$, where again $\mathbb{P}$ is the distribution of $\mathbf{X}$ and $p_{t}$ corresponds to the density of $\mathbb{P}_{t}$ with respect to the uniform distribution on $\mathcal{M}$, then $\nabla\log p_{t}(x_{t})=\operatorname*{arg\,min}_{\mathbf{s}_{t}}\int_{\mathcal{M}^{2}}\|\nabla_{x}\log p_{t|s}(x_{t}|x_{s})-\mathbf{s}_{t}(x_{t})\|^{2}\mathrm{d}\mathbb{P}_{s,t}(x_{s},x_{t})$.

A Schrödinger bridge extension of SGM has been introduced to reduce the number of diffusion steps for SGM by learning both forward and backward diffusions (De Bortoli et al., 2021). We briefly recall the notion of dynamical Schrödinger bridge (Léonard, 2012; Chen et al., 2016; Vargas et al., 2021; De Bortoli et al., 2021; Chen et al., 2022). We consider a reference path probability measure $\mathbb{P}\in\mathcal{P}(\mathrm{C}(\left[0,T\right]),\mathcal{M})$ where $\mathcal{P}(\mathrm{C}(\left[0,T\right]),\mathcal{M})$ is the space measures on continuous paths $(\mathbf{X}_{t})_{t\in[0,1]}$ in $\mathcal{M}$ . In practice, we set $\mathbb{P}$ to be the distribution of the Brownian motion $(\mathbf{B}_{t}^{\mathcal{M}})_{t\in\left[0,T\right]}$ initialized from $p_{\textup{data}}$, i.e. $\mathbf{B}_{0}^{\mathcal{M}}$ has distribution $p_{\textup{data}}$. We consider the dynamical Schrödinger bridge problem

The idealised solution $\mathbb{Q}^{\star}$ is called the Schrödinger Bridge (SB). Given a backward process $(\mathbf{Y}_{t}^{\star})_{t\in\left[0,T\right]}$ associated to $\mathbb{Q}^{\star}$, one can obtain a generative model as follows. First sample from $\mathbf{Y}^{\star}_{T}\sim p_{\textup{prior}}$ and then follow the (backward) dynamics of $(\mathbf{Y}^{\star}_{t})_{t\in\left[0,T\right]}$. By definition, we obtain that $\mathbf{Y}^{\star}_{0}\sim p_{\textup{data}}$, the data distribution.

In practice, the solution of the SB problem is approximated using the Iterative Proportional Fitting (IPF) algorithm, which coincides with the Sinhkorn algorithm in discrete space (Sinkhorn, 1967; Peyré & Cuturi, 2019). IPF defines a sequence of path probability measures $(\mathbb{Q}^{n})_{n\in\mathbb{N}}\in(\mathcal{P}(\mathrm{C}(\left[0,T\right],\mathcal{M})))^{\mathbb{N}}$, such that $\mathbb{Q}^{0}=\mathbb{P}$ and for any $n\in\mathbb{N}$

Under mild assumptions on $\mathbb{P}$, $p_{\textup{data}}$ and $p_{\textup{prior}}$, we have that $(\mathbb{Q}^{n})_{n\in\mathbb{N}}$ converges towards $\mathbb{Q}^{\star}$ (Nutz & Wiesel, 2022).

RDSB enables interpolation between datasets by choosing $p_{\textup{prior}}$ to be another dataset, and not necessary $p_{T}$, the terminal distribution of the inital noising process. Figure 3 illustrates RDSB applied to samples on the sphere taken with probability proportional to the real component of spherical harmonic function with parameters $(l=2,m=4)$ and $(l=2,m=6)$.