Denoising Diffusion Bridge Models
去噪扩散桥模型

Diffusion process. We are interested in modeling the distribution $q_{\rm{data}}({\mathbf{x}})$, for ${\mathbf{x}}\in\mathbb{R}^{d}$. We do this by constructing a diffusion process, which is represented by a set of time-indexed variables $\{{\mathbf{x}}_{t}\}_{t=0}^{T}$ such that ${\mathbf{x}}_{0}\sim p_{0}({\mathbf{x}}):=q_{\rm{data}}({\mathbf{x}})$ and ${\mathbf{x}}_{T}\sim p_{T}({\mathbf{x}}):=p_{\rm{prior}}({\mathbf{x}})$. Here $q_{\rm{data}}({\mathbf{x}})$ is the initial “data" distribution and $p_{\rm{prior}}({\mathbf{x}})$ is the final “prior" distribution. The process can be modeled as the solution to the following SDE
扩散过程。我们感兴趣的是对 ${\mathbf{x}}\in\mathbb{R}^{d}$ 的分布 $q_{\rm{data}}({\mathbf{x}})$ 进行建模。我们通过构建一个扩散过程来实现这一点，该过程由一组时间索引变量 $\{{\mathbf{x}}_{t}\}_{t=0}^{T}$ 表示，例如 ${\mathbf{x}}_{0}\sim p_{0}({\mathbf{x}}):=q_{\rm{data}}({\mathbf{x}})$ 和 ${\mathbf{x}}_{T}\sim p_{T}({\mathbf{x}}):=p_{\rm{prior}}({\mathbf{x}})$ 。这里 $q_{\rm{data}}({\mathbf{x}})$ 是初始“数据”分布， $p_{\rm{prior}}({\mathbf{x}})$ 是最终“先验”分布。该过程可以建模为以下 SDE 的解决方案

$\displaystyle\quad d{\mathbf{x}}_{t}={\mathbf{f}}({\mathbf{x}}_{t},t)dt+g(t)d{\mathbf{w}}_{t}$

(1)

where ${\mathbf{f}}:\mathbb{R}^{d}\times[0,T]\rightarrow\mathbb{R}^{d}$ is vector-valued drift function, $g:[0,T]\rightarrow\mathbb{R}$ is a scalar-valued diffusion coefficient, and ${\mathbf{w}}_{t}$ is a Wiener process. Following this diffusion process forward in time constrains the final variable ${\mathbf{x}}_{T}$ to follow distribution $p_{\rm{prior}}({\mathbf{x}})$. The reverse of this process is given by
其中 ${\mathbf{f}}:\mathbb{R}^{d}\times[0,T]\rightarrow\mathbb{R}^{d}$ 是矢量值漂移函数， $g:[0,T]\rightarrow\mathbb{R}$ 是标量值扩散系数， ${\mathbf{w}}_{t}$ 是维纳过程。在时间上遵循此扩散过程会限制最终变量 ${\mathbf{x}}_{T}$ 遵循分布 $p_{\rm{prior}}({\mathbf{x}})$ 。该过程的逆过程由下式给出

$\displaystyle\quad d{\mathbf{x}}_{t}={\mathbf{f}}({\mathbf{x}}_{t},t)-g(t)^{2}\nabla_{{\mathbf{x}}_{t}}\log p({\mathbf{x}}_{t}))dt+g(t)d{\mathbf{w}}_{t}$

(2)

where $p({\mathbf{x}}_{t})\vcentcolon=p({\mathbf{x}}_{t},t)$ is the marginal distribution of ${\mathbf{x}}_{t}$ at time $t$. Furthermore, one can derive an equivalent deterministic process called the probability flow ODE (Song et al., 2020b), which has the same marginal distributions:
其中 $p({\mathbf{x}}_{t})\vcentcolon=p({\mathbf{x}}_{t},t)$ 是 ${\mathbf{x}}_{t}$ 在时间 $t$ 的边际分布。此外，我们还可以推导出一种等效的确定性过程，称为概率流 ODE（Song 等人，2020b），它具有相同的边际分布：

$\displaystyle d{\mathbf{x}}_{t}=\Big{[}{\mathbf{f}}({\mathbf{x}}_{t},t)-\frac{1}{2}g(t)^{2}\nabla_{{\mathbf{x}}_{t}}\log p({\mathbf{x}}_{t})\Big{]}dt$

(3)

In particular, one can draw ${\mathbf{x}}_{T}\sim q_{\rm{data}}({\mathbf{y}})$ and sample $q_{\rm{data}}$ by solving either the above reverse SDE or ODE backward in time.
特别是，可以通过在时间上向后求解上述逆 SDE 或 ODE 来绘制 ${\mathbf{x}}_{T}\sim q_{\rm{data}}({\mathbf{y}})$ 和样本 $q_{\rm{data}}$ 。

Denoising score-matching. The score, $\nabla_{{\mathbf{x}}_{t}}\log p({\mathbf{x}}_{t})$, can be learned by the score-matching loss
去噪分数匹配。分数 $\nabla_{{\mathbf{x}}_{t}}\log p({\mathbf{x}}_{t})$ 可以通过分数匹配损失来学习

$\displaystyle\mathcal{L}(\theta)=\mathbb{E}_{{\mathbf{x}}_{t}\sim p({\mathbf{x}}_{t}\mid{\mathbf{x}}_{0}),{\mathbf{x}}_{0}\sim q_{\rm{data}}({\mathbf{x}}),t\sim{\mathcal{U}}(0,T)}\Big{[}\norm{{\mathbf{s}}_{\theta}({\mathbf{x}}_{t},t)-\nabla_{{\mathbf{x}}_{t}}\log p({\mathbf{x}}_{t}\mid{\mathbf{x}}_{0})}^{2}\Big{]}$

(4)

such that the minimizer ${\mathbf{s}}^{*}_{\theta}({\mathbf{x}}_{t},t)$ of the above loss approximates the true score. Crucially, the above loss is tractable because the transition kernel $p({\mathbf{x}}_{t}\mid{\mathbf{x}}_{0})$, which depends on specific choices of drift and diffusion functions, is designed to be Gaussian ${\mathbf{x}}_{t}=\alpha_{t}{\mathbf{x}}_{0}+\sigma_{t}{\bm{\epsilon}}$, where $\alpha_{t}$ and $\sigma_{t}$ are functions of time and ${\bm{\epsilon}}\sim{\mathcal{N}}({\bm{0}},{\bm{I}})$. It is also common to view the diffusion process in terms of the ${\mathbf{x}}_{t}$’s signal-to-noise ratio (SNR), defined as $\alpha_{t}^{2}/\sigma_{t}^{2}$.
使得上述损失的最小化 ${\mathbf{s}}^{*}_{\theta}({\mathbf{x}}_{t},t)$ 接近真实分数。至关重要的是，上述损失是易于处理的，因为取决于漂移和扩散函数的特定选择的转换内核 $p({\mathbf{x}}_{t}\mid{\mathbf{x}}_{0})$ 被设计为高斯 ${\mathbf{x}}_{t}=\alpha_{t}{\mathbf{x}}_{0}+\sigma_{t}{\bm{\epsilon}}$ ，其中 $\alpha_{t}$ 和 $\sigma_{t}$ 是时间和 ${\bm{\epsilon}}\sim{\mathcal{N}}({\bm{0}},{\bm{I}})$ 的函数。通常还根据 ${\mathbf{x}}_{t}$ 的信噪比 (SNR) 来查看扩散过程，定义为 $\alpha_{t}^{2}/\sigma_{t}^{2}$ 。

Diffusion models are limited because they can only transport complex data distributions to a standard Gaussian distribution and cannot be naturally adapted to translating between two arbitrary distributions, e.g. in the case of image-to-image translation. Luckily, classical results have shown that one can condition a diffusion process on a fixed known endpoint via the famous Doob’s $h$-transform:
扩散模型是有限的，因为它们只能将复杂的数据分布传输到标准高斯分布，并且不能自然地适应两个任意分布之间的转换，例如在图像到图像翻译的情况下。幸运的是，经典结果表明，我们可以通过著名的 Doob 的 $h$ 变换来调节固定已知端点上的扩散过程：

Stochastic bridges via $h$-transform. Specifically, a diffusion process defined in Eq. (1) can be driven to arrive at a particular point of interest $y\in\mathbb{R}^{d}$ almost surely via Doob’s $h$-transform (Doob and Doob, 1984; Rogers and Williams, 2000),
通过 $h$ 变换的随机桥。具体来说，方程中定义的扩散过程。 (1) 几乎肯定可以通过 Doob 的 $h$ 变换到达特定的兴趣点 $y\in\mathbb{R}^{d}$ （Doob 和 Doob，1984；Rogers 和 Williams，2000），

$\displaystyle d{\mathbf{x}}_{t}={\mathbf{f}}({\mathbf{x}}_{t},t)dt+g(t)^{2}{\mathbf{h}}({\mathbf{x}}_{t},t,y,T)+g(t)d{\mathbf{w}}_{t},\quad{\mathbf{x}}_{0}\sim q_{\rm{data}}({\mathbf{x}}),\quad{\mathbf{x}}_{T}=y$

(5)

where ${\mathbf{h}}(x,t,y,T)=\nabla_{{\mathbf{x}}_{t}}\log p({\mathbf{x}}_{T}\mid{\mathbf{x}}_{t})\big{\rvert}_{{\mathbf{x}}_{t}=x,{\mathbf{x}}_{T}=y}$ is the gradient of the log transition kernel of from $t$ to $T$ generated by the original SDE, evaluated at points ${\mathbf{x}}_{t}=x$ and ${\mathbf{x}}_{T}=y$, and each ${\mathbf{x}}_{t}$ now explicitly depends on $y$ at time $T$. Furthermore, $p({\mathbf{x}}_{T}=y\mid{\mathbf{x}}_{t})$ satisfies the Kolmogorov backward equation (specified in Appendix A). With specific drift and diffusion choices, e.g. ${\mathbf{f}}({\mathbf{x}}_{t},t)={\bm{0}}$, ${\mathbf{h}}$ is tractable due to the tractable (Gaussian) transition kernel of the underlying diffusion process.
其中 ${\mathbf{h}}(x,t,y,T)=\nabla_{{\mathbf{x}}_{t}}\log p({\mathbf{x}}_{T}\mid{\mathbf{x}}_{t})\big{\rvert}_{{\mathbf{x}}_{t}=x,{\mathbf{x}}_{T}=y}$ 是原始 SDE 生成的从 $t$ 到 $T$ 的对数转换内核的梯度，在点 ${\mathbf{x}}_{t}=x$ 处评估，并且 ${\mathbf{x}}_{T}=y$ ，并且每个 ${\mathbf{x}}_{t}$ 现在在 $T$ 时间明确依赖于 $y$ 。此外， $p({\mathbf{x}}_{T}=y\mid{\mathbf{x}}_{t})$ 满足Kolmogorov后向方程（在附录A中指定）。具有特定的漂移和扩散选择，例如 ${\mathbf{f}}({\mathbf{x}}_{t},t)={\bm{0}}$ 、 ${\mathbf{h}}$ 由于底层扩散过程的易处理（高斯）转换核而易于处理。

When the initial point ${\mathbf{x}}_{0}$ is fixed, the process is often called a diffusion bridge (Särkkä and Solin, 2019; Heng et al., 2021; Delyon and Hu, 2006; Schauer et al., 2017; Peluchetti, ; Liu et al., 2022b), and its ability to connect any given ${\mathbf{x}}_{0}$ to a given value of ${\mathbf{x}}_{T}$ is promising for image-to-image translation. Furthermore, the transition kernel may be tractable, which serves as further motivation.

A schematic of the bridge process is shown in Figure 1. Note that this process is defined up to $T-\epsilon$. To recover the initial distribution in the SDE case, we make an approximation that ${\mathbf{x}}_{T-\epsilon}\approx y$ for some small $\epsilon$ simulate SDE backward in time. For the ODE case, since we need to sample from $p({\mathbf{x}}_{T-\epsilon})$ which cannot be Dirac delta, we cannot approximate ${\mathbf{x}}_{T-\epsilon}$ with a single $y$. Instead, we can first approximate ${\mathbf{x}}_{T-\epsilon^{\prime}}\approx y$ where $\epsilon>\epsilon^{\prime}>0$, and then take an Euler-Maruyama step to ${\mathbf{x}}_{T-\epsilon}$, and Eq. (7) can be used afterward. A toy visualization of VE bridge and VP bridges are shown in Figure 2. The top and bottom shows the respective SDE and ODE paths for VE and VP bridges.

The sampling process in Theorem 1 requires approximation of the score ${\mathbf{s}}(x,t,y,T)=\nabla_{{\mathbf{x}}_{t}}\log q({\mathbf{x}}_{t}\mid{\mathbf{x}}_{T})\big{|}_{{\mathbf{x}}_{t}=x,{\mathbf{x}}_{T}=y}$ where $q({\mathbf{x}}_{t}\mid{\mathbf{x}}_{T})=\int_{{\mathbf{x}}_{0}}q({\mathbf{x}}_{t}\mid{\mathbf{x}}_{0},{\mathbf{x}}_{T})q_{\rm{data}}({\mathbf{x}}_{0}\mid{\mathbf{x}}_{T})d{\mathbf{x}}_{0}$. However, as the true score is not known in closed-form, we take inspiration from denoising score-matching (Song et al., 2020b) and use a neural network to approximate the true score by matching against a tractable quantity. This usually results in closed-form marginal sampling of ${\mathbf{x}}_{t}$ given data (e.g. ${\mathbf{x}}_{0}$ in the case of diffusion models and $({\mathbf{x}}_{0},{\mathbf{x}}_{T})$ in our case), and given ${\mathbf{x}}_{t}$, the model is trained to match against the closed-form denoising score objective. We are motivated to follow a similar approach because (1) tractable marginal sampling of ${\mathbf{x}}_{t}$ and (2) closed-form objectives enable a simple and scalable algorithm. We specify how to design the marginal sampling distribution and the tractable score objective below to approximate the ground-truth conditional score $\nabla_{{\mathbf{x}}_{t}}\log q({\mathbf{x}}_{t}\mid{\mathbf{x}}_{T})$.
定理 1 中的采样过程需要分数 ${\mathbf{s}}(x,t,y,T)=\nabla_{{\mathbf{x}}_{t}}\log q({\mathbf{x}}_{t}\mid{\mathbf{x}}_{T})\big{|}_{{\mathbf{x}}_{t}=x,{\mathbf{x}}_{T}=y}$ 的近似值，其中 $q({\mathbf{x}}_{t}\mid{\mathbf{x}}_{T})=\int_{{\mathbf{x}}_{0}}q({\mathbf{x}}_{t}\mid{\mathbf{x}}_{0},{\mathbf{x}}_{T})q_{\rm{data}}({\mathbf{x}}_{0}\mid{\mathbf{x}}_{T})d{\mathbf{x}}_{0}$ 。然而，由于真实分数在封闭形式中是未知的，因此我们从去噪分数匹配（Song et al., 2020b）中获得灵感，并使用神经网络通过与易于处理的量进行匹配来近似真实分数。这通常会导致 ${\mathbf{x}}_{t}$ 给定数据的封闭形式边际采样（例如，在扩散模型的情况下为 ${\mathbf{x}}_{0}$ ，在我们的情况下为 $({\mathbf{x}}_{0},{\mathbf{x}}_{T})$ ），并给定 < b5>，训练模型以匹配封闭式去噪分数目标。我们有动力遵循类似的方法，因为 (1) ${\mathbf{x}}_{t}$ 的易于处理的边际抽样和 (2) 封闭形式的目标支持简单且可扩展的算法。我们在下面指定如何设计边际抽样分布和易处理的分数目标来近似地面实况条件分数 $\nabla_{{\mathbf{x}}_{t}}\log q({\mathbf{x}}_{t}\mid{\mathbf{x}}_{T})$ 。

Sampling distribution. Fortunately, for the former condition, we can design our sampling distribution $q(\cdot)$ such that $q({\mathbf{x}}_{t}\mid{\mathbf{x}}_{0},{\mathbf{x}}_{T})\vcentcolon=p({\mathbf{x}}_{t}\mid{\mathbf{x}}_{0},{\mathbf{x}}_{T})$, where $p(\cdot)$ is the diffusion distribution pinned at both endpoints as in Eq. (5). For diffusion processes with Gaussian transition kernels, e.g. VE, VP (Song et al., 2020b), our sampling distribution is a Gaussian distribution of the form
抽样分配。幸运的是，对于前一种情况，我们可以设计采样分布 $q(\cdot)$ ，使得 $q({\mathbf{x}}_{t}\mid{\mathbf{x}}_{0},{\mathbf{x}}_{T})\vcentcolon=p({\mathbf{x}}_{t}\mid{\mathbf{x}}_{0},{\mathbf{x}}_{T})$ ，其中 $p(\cdot)$ 是固定在两个端点的扩散分布，如方程式 1 所示。 （5）。对于具有高斯转移核的扩散过程，例如VE、VP（Song et al., 2020b），我们的抽样分布是以下形式的高斯分布

	$\displaystyle q({\mathbf{x}}_{t}\mid{\mathbf{x}}_{0},{\mathbf{x}}_{T})$	$\displaystyle={\mathcal{N}}(\hat{\mu}_{t},\hat{\sigma}_{t}^{2}\bm{I}),\quad\text{where}$		(8)
	$\displaystyle\hat{\mu}_{t}$	$\displaystyle=\frac{\text{SNR}_{T}}{\text{SNR}_{t}}\frac{\alpha_{t}}{\alpha_{T}}{\mathbf{x}}_{T}+\alpha_{t}{\mathbf{x}}_{0}(1-\frac{\text{SNR}_{T}}{\text{SNR}_{t}})$
	$\displaystyle\hat{\sigma}_{t}^{2}$	$\displaystyle=\sigma_{t}^{2}(1-\frac{\text{SNR}_{T}}{\text{SNR}_{t}})$

where $\alpha_{t}$ and $\sigma_{t}$ are pre-defined signal and noise schedules and $\text{SNR}_{t}=\alpha_{t}^{2}/\sigma_{t}^{2}$ is the signal-to-noise ratio at time $t$. For VE schedule, we assume $\alpha_{t}=1$ and derivation details are provided in Appendix A.1. Notably, the mean of this distribution is a linear interpolation between the (scaled) endpoints, and the distribution approaches a Dirac distribution when nearing either end. For concreteness, we present the bridge processes generated by both VP and VE diffusion in Table 1 and recommend choosing ${\mathbf{f}}$ and $g$ specified therein.

Training objective. For the latter condition, diffusion bridges benefit from a similar setup as in diffusion models, since a pre-defined signal/noise schedule gives rise to a closed-form conditional score $\nabla_{{\mathbf{x}}_{t}}\log q({\mathbf{x}}_{t}\mid{\mathbf{x}}_{0},{\mathbf{x}}_{T})$. We show in the following theorem that with ${\mathbf{x}}_{t}\sim q({\mathbf{x}}_{t}\mid{\mathbf{x}}_{0},{\mathbf{x}}_{T})$, a neural network ${\mathbf{s}}_{\theta}({\mathbf{x}}_{t},{\mathbf{x}}_{T},t)$ that matches against this closed-form score approximates the true score.

In short, we establish a tractable diffusion bridge over two endpoints and, by matching the conditional score of the Gaussian bridge, we can learn the score of the new distribution $q({\mathbf{x}}_{t}\mid{\mathbf{x}}_{T})$ that satisfies the boundary distribution $q_{\rm{data}}({\mathbf{x}},{\mathbf{y}})$.