Diffusion Posterior Sampling (DPS)

前置

• posterior mean

1. 设定

考虑一般的线性高斯前向扩散过程。记 $a_t,b_t>0$ 为时间 $t$ 的确定性系数，

$${\mathbf{x}}_t=a_t{\mathbf{x}}_0+b_t\mathbf{ϵ},\mathbf{ϵ}\sim \mathcal{N}(0,\mathbf{I}).\text{\hspace{1em}(1)}$$

无条件 score 模型给出 ${\nabla }_{{\mathbf{x}}_t}\log p_t({\mathbf{x}}_t)$。由 posterior mean 的 (5)，后验均值是

$${\hat{\mathbf{x}}}_0:=\mathbb{E}[{\mathbf{x}}_0\mid {\mathbf{x}}_t]=\frac{b_t^2}{a_t}{\nabla }_{{\mathbf{x}}_t}\log p_t({\mathbf{x}}_t)+\frac{1}{a_t}{\mathbf{x}}_t.\text{\hspace{1em}(2)}$$

测量模型与文末定理 1 相同：

$$\mathbf{y}=\mathcal{A}({\mathbf{x}}_0)+\mathbf{n},\mathbf{n}\sim \mathcal{N}(0,{\sigma }^2\mathbf{I}),\text{\hspace{1em}(3)}$$

其中 $\mathcal{A}$ 已知。目标是从后验 $p({\mathbf{x}}_0\mid \mathbf{y})$ 采样。

扩散核 $p({\mathbf{x}}_t\mid {\mathbf{x}}_0)$ 与测量似然 $p(\mathbf{y}\mid {\mathbf{x}}_0)$ 都已知。$p({\mathbf{x}}_0\mid {\mathbf{x}}_t)$ 的完整分布未知，因此 $p(\mathbf{y}\mid {\mathbf{x}}_t)$ 没有闭式。

2. 后验 score

Bayes 公式把条件 score 拆成先验 score 与似然 score：

$${\nabla }_{{\mathbf{x}}_t}\log p_t({\mathbf{x}}_t\mid \mathbf{y})={\nabla }_{{\mathbf{x}}_t}\log p_t({\mathbf{x}}_t)+{\nabla }_{{\mathbf{x}}_t}\log p(\mathbf{y}\mid {\mathbf{x}}_t).\text{\hspace{1em}(4)}$$

先验 score 由无条件扩散模型提供。反向过程在每个时刻需要似然 score ${\nabla }_{{\mathbf{x}}_t}\log p(\mathbf{y}\mid {\mathbf{x}}_t)$。测量模型 (3) 直接给出的是干净数据上的 $p(\mathbf{y}\mid {\mathbf{x}}_0)$。

3. 积分与插件近似

在 $\mathbf{y}\perp {\mathbf{x}}_t\mid {\mathbf{x}}_0$ 下，

$$p(\mathbf{y}\mid {\mathbf{x}}_t)=\int p(\mathbf{y}\mid {\mathbf{x}}_0)p({\mathbf{x}}_0\mid {\mathbf{x}}_t)d{\mathbf{x}}_0=\mathbb{E}[p(\mathbf{y}\mid {\mathbf{x}}_0)\mid {\mathbf{x}}_t].\text{\hspace{1em}(5)}$$

积分需要完整的 $p({\mathbf{x}}_0\mid {\mathbf{x}}_t)$。扩散模型通常只给 score，并由 (2) 给出条件均值 ${\hat{\mathbf{x}}}_0$。DPS 用这个均值做插件近似 (Chung et al., 2022)：

$$p(\mathbf{y}\mid {\mathbf{x}}_t)\approx p(\mathbf{y}\mid {\hat{\mathbf{x}}}_0),{\hat{\mathbf{x}}}_0=\mathbb{E}[{\mathbf{x}}_0\mid {\mathbf{x}}_t].\text{\hspace{1em}(6)}$$

记 $f({\mathbf{x}}_0)=p(\mathbf{y}\mid {\mathbf{x}}_0)$，则 (5) 是 $\mathbb{E}[f({\mathbf{x}}_0)\mid {\mathbf{x}}_t]$，(6) 是 $f(\mathbb{E}[{\mathbf{x}}_0\mid {\mathbf{x}}_t])$。二者一般不相等。文末定理 1 在 Lipschitz 条件下给出这个间隙的上界。

4. 似然 score 与采样

对高斯测量 (3)，

$$\log p(\mathbf{y}\mid {\mathbf{x}}_0)=-\frac{\Vert \mathbf{y}-\mathcal{A}({\mathbf{x}}_0){\Vert }^2}{2{\sigma }^2}+C.$$

把 ${\mathbf{x}}_0$ 换成 ${\hat{\mathbf{x}}}_0({\mathbf{x}}_t)$ 后对 ${\mathbf{x}}_t$ 求导。梯度经过 ${\hat{\mathbf{x}}}_0$，而 ${\hat{\mathbf{x}}}_0$ 通过 (2) 依赖 ${\mathbf{x}}_t$ 与 score：

$${\nabla }_{{\mathbf{x}}_t}\log p(\mathbf{y}\mid {\hat{\mathbf{x}}}_0)=-\frac{1}{2{\sigma }^2}{\nabla }_{{\mathbf{x}}_t}\Vert \mathbf{y}-\mathcal{A}({\hat{\mathbf{x}}}_0){\Vert }^2.\text{\hspace{1em}(7)}$$

条件反向过程把 (4) 中的似然 score 换成 (7)，并把其中的系数收进步长 $\zeta$。先按无条件反向过程从 ${\mathbf{x}}_t$ 得到候选 ${\mathbf{x}}_{t-1}^{\prime }$，再减去残差范数对当前 ${\mathbf{x}}_t$ 的梯度 (Chung et al., 2022)：

$${\mathbf{x}}_{t-1}={\mathbf{x}}_{t-1}^{\prime }-\zeta {\nabla }_{{\mathbf{x}}_t}\Vert \mathbf{y}-\mathcal{A}({\hat{\mathbf{x}}}_0){\Vert }_2.\text{\hspace{1em}(8)}$$

5. 特例

公式 (2) 在两种常见前向过程中收回熟悉的后验均值，代入方式与 posterior mean 的 2.3 相同。

DDPM / VP-SDE：取 $a_t=\sqrt{{\bar{\alpha }}_t}$、$b_t=\sqrt{1-{\bar{\alpha }}_t}$，其中 ${\bar{\alpha }}_t={\prod }_{s=1}^t{\alpha }_s$。代入 (2) 得

$${\hat{\mathbf{x}}}_0=\frac{1}{\sqrt{{\bar{\alpha }}_t}}({\mathbf{x}}_t+(1-{\bar{\alpha }}_t){\nabla }_{{\mathbf{x}}_t}\log p_t({\mathbf{x}}_t)).\text{\hspace{1em}(9)}$$

VE-SDE：取 $a_t=1$、$b_t={\sigma }_t$，则

$${\hat{\mathbf{x}}}_0={\mathbf{x}}_t+{\sigma }_t^2{\nabla }_{{\mathbf{x}}_t}\log p_t({\mathbf{x}}_t).\text{\hspace{1em}(10)}$$

命题 2（Jensen 间隙上界）

设 $\mathbf{X}\in {\mathbb{R}}^n$ 是随机向量，其 $\alpha$ 阶绝对中心矩有限（$\alpha >0$）。记

$$\mathbf{\mu }=\mathbb{E}[\mathbf{X}],m_{\alpha }={\left(\mathbb{E}\Vert \mathbf{X}-\mathbf{\mu }{\Vert }_2^{\alpha }\right)}^{1/\alpha }.$$

若存在常数 $K>0$，使对任意 $\mathbf{x}\in {\mathbb{R}}^n$ 都有

$$|f(\mathbf{x})-f(\mathbf{\mu })|\le K\Vert \mathbf{x}-\mathbf{\mu }{\Vert }_2^{\alpha },$$

则

$$\left|\mathbb{E}[f(\mathbf{X})]-f(\mathbb{E}[\mathbf{X}])\right|\le K{m}_{\alpha }^{\alpha }.$$

证明. 由期望的线性性与三角不等式，

$$\left|\mathbb{E}[f(\mathbf{X})]-f(\mathbf{\mu })\right|=\left|\mathbb{E}[f(\mathbf{X})-f(\mathbf{\mu })]\right|\le \mathbb{E}|f(\mathbf{X})-f(\mathbf{\mu })|\le K\mathbb{E}\Vert \mathbf{X}-\mathbf{\mu }{\Vert }_2^{\alpha }=K{m}_{\alpha }^{\alpha }.\square$$

引理 2（一维高斯密度是 Lipschitz 的）

设 $\varphi :\mathbb{R}\to \mathbb{R}$ 是均值为 $\mu$、方差为 ${\sigma }^2>0$ 的一维高斯密度，则存在有限常数 $L>0$，使对任意 $x,y\in \mathbb{R}$ 有

$$|\varphi (x)-\varphi (y)|\le L|x-y|.$$

证明. 由 ${\varphi }^{\prime }(x)=-\frac{x-\mu }{{\sigma }^2}\varphi (x)$ 知 ${\varphi }^{\prime }$ 连续，且 ${\lim }_{x\to \pm \infty }{\varphi }^{\prime }(x)=0$（指数衰减快于线性增长），故 ${\varphi }^{\prime }$ 在 $\mathbb{R}$ 上有界。取 $L={\sup }_{x\in \mathbb{R}}|{\varphi }^{\prime }(x)|<\infty$，由中值定理，

$$|\varphi (x)-\varphi (y)|=|{\varphi }^{\prime }(\xi )||x-y|\le L|x-y|.\square$$

引理 3（各向同性多元高斯密度是 Lipschitz 的）

设 $\varphi :{\mathbb{R}}^d\to \mathbb{R}$ 是均值为 $\mathbf{\mu }$、协方差为 ${\sigma }^2\mathbf{I}$ 的高斯密度（${\sigma }^2>0$），则存在有限常数 $L>0$，使对任意 $\mathbf{x},\mathbf{y}\in {\mathbb{R}}^d$ 有

$$|\varphi (\mathbf{x})-\varphi (\mathbf{y})|\le L\Vert \mathbf{x}-\mathbf{y}{\Vert }_2.$$

证明. 梯度为

$$\nabla \varphi (\mathbf{x})=-\frac{\mathbf{x}-\mathbf{\mu }}{{\sigma }^2}\varphi (\mathbf{x}),$$

故

$$\Vert \nabla \varphi (\mathbf{x}){\Vert }_2=\frac{\Vert \mathbf{x}-\mathbf{\mu }{\Vert }_2}{{\sigma }^2(2\pi {\sigma }^2{)}^{d/2}}\exp \left(-\frac{\Vert \mathbf{x}-\mathbf{\mu }{\Vert }_2^2}{2{\sigma }^2}\right)$$

连续，且在 $\Vert \mathbf{x}-\mathbf{\mu }{\Vert }_2\to \infty$ 时趋于 $0$，从而在 ${\mathbb{R}}^d$ 上有界。取 $L={\sup }_{\mathbf{z}\in {\mathbb{R}}^d}\Vert \nabla \varphi (\mathbf{z}){\Vert }_2<\infty$，沿线段 $\mathbf{\gamma }(s)=\mathbf{y}+s(\mathbf{x}-\mathbf{y})$（$s\in [0,1]$）积分并用 Cauchy–Schwarz 不等式，

$$|\varphi (\mathbf{x})-\varphi (\mathbf{y})|\le {\int }_0^1\Vert \nabla \varphi (\mathbf{\gamma }(s)){\Vert }_2\Vert \mathbf{x}-\mathbf{y}{\Vert }_{2}ds\le L\Vert \mathbf{x}-\mathbf{y}{\Vert }_2.\square$$

定理 1（DPS 的近似误差上界）

考虑测量模型

$$\mathbf{y}=\mathcal{A}({\mathbf{x}}_0)+\mathbf{n},\mathbf{n}\sim \mathcal{N}(0,{\sigma }^2{\mathbf{I}}_d),$$

其中 $\mathcal{A}:{\mathbb{R}}^n\to {\mathbb{R}}^d$ 连续可微且 Jacobian 有界：

$$C_{\mathcal{A}}:=\underset{\mathbf{x}\in {\mathbb{R}}^n}{\sup }\Vert J_{\mathcal{A}}(\mathbf{x}){\Vert }_{\mathrm{op}}<\infty .$$

令 ${\hat{\mathbf{x}}}_0=\mathbb{E}[{\mathbf{x}}_0\mid {\mathbf{x}}_t]$，

$$m_1:=\mathbb{E}[\Vert {\mathbf{x}}_0-{\hat{\mathbf{x}}}_0{\Vert }_2\mid {\mathbf{x}}_t],$$

则存在有限常数 $L>0$，使

$$\left|p(\mathbf{y}\mid {\mathbf{x}}_t)-p(\mathbf{y}\mid {\hat{\mathbf{x}}}_0)\right|\le L{C}_{\mathcal{A}}{m}_1.$$

证明.

第一步：把边缘似然写成期望. 记

$$h(\mathbf{z})=\frac{1}{(2\pi {\sigma }^2{)}^{d/2}}\exp \left(-\frac{\Vert \mathbf{y}-\mathbf{z}{\Vert }_2^2}{2{\sigma }^2}\right),f(\mathbf{x})=h(\mathcal{A}(\mathbf{x})),$$

则 $p(\mathbf{y}\mid {\mathbf{x}}_0)=f({\mathbf{x}}_0)$。由全概率公式，

$$p(\mathbf{y}\mid {\mathbf{x}}_t)=\int f({\mathbf{x}}_0)p({\mathbf{x}}_0\mid {\mathbf{x}}_t)d{\mathbf{x}}_0=\mathbb{E}[f({\mathbf{x}}_0)\mid {\mathbf{x}}_t],$$

而 $p(\mathbf{y}\mid {\hat{\mathbf{x}}}_0)=f({\hat{\mathbf{x}}}_0)$，故误差为

$$\mathcal{J}:=\left|\mathbb{E}[f({\mathbf{x}}_0)\mid {\mathbf{x}}_t]-f({\hat{\mathbf{x}}}_0)\right|.$$

第二步：$h$ 与 $\mathcal{A}$ 都是 Lipschitz 的. 对固定的 $\mathbf{y}$，$h$ 是协方差为 ${\sigma }^2{\mathbf{I}}_d$ 的高斯密度，由引理 3 存在 $L<\infty$ 使

$$|h({\mathbf{z}}_1)-h({\mathbf{z}}_2)|\le L\Vert {\mathbf{z}}_1-{\mathbf{z}}_2{\Vert }_2.$$

另一方面，由 Jacobian 有界以及

$$\mathcal{A}(\mathbf{x})-\mathcal{A}({\mathbf{x}}^{\prime })={\int }_0^1{J}_{\mathcal{A}}({\mathbf{x}}^{\prime }+s(\mathbf{x}-{\mathbf{x}}^{\prime }))(\mathbf{x}-{\mathbf{x}}^{\prime })ds$$

可得 $\Vert \mathcal{A}(\mathbf{x})-\mathcal{A}({\mathbf{x}}^{\prime }){\Vert }_2\le C_{\mathcal{A}}\Vert \mathbf{x}-{\mathbf{x}}^{\prime }{\Vert }_2$，即 $\mathcal{A}$ 是 $C_{\mathcal{A}}$-Lipschitz 的。二者复合给出

$$|f({\mathbf{x}}_0)-f({\hat{\mathbf{x}}}_0)|\le L{C}_{\mathcal{A}}\Vert {\mathbf{x}}_0-{\hat{\mathbf{x}}}_0{\Vert }_2.$$

第三步：估计误差. 对左端取条件期望，用三角不等式即得

$$\mathcal{J}\le \mathbb{E}[|f({\mathbf{x}}_0)-f({\hat{\mathbf{x}}}_0)|\mid {\mathbf{x}}_t]\le L{C}_{\mathcal{A}}\mathbb{E}[\Vert {\mathbf{x}}_0-{\hat{\mathbf{x}}}_0{\Vert }_2\mid {\mathbf{x}}_t]=L{C}_{\mathcal{A}}{m}_1.$$

这正是命题 2 在 $\alpha =1$、$K=L{C}_{\mathcal{A}}$、$f=h\circ \mathcal{A}$ 下的特例。$\square$

备注

• 全程只要求 Lipschitz 常数存在且有限，不必算出其显式表达式；引理 2、3 只需导数或梯度连续且有界。
• 仅有 $h$ 的 Lipschitz 性推不出 $h\circ \mathcal{A}$ 的 Lipschitz 性，因此需要对 $\mathcal{A}$ 加全局 Lipschitz 条件，证明中取的是充分条件 ${\sup }_{\mathbf{x}}\Vert J_{\mathcal{A}}(\mathbf{x}){\Vert }_{\mathrm{op}}<\infty$。若 $\mathcal{A}$ 线性，则自动成立且 $C_{\mathcal{A}}=\Vert \mathcal{A}{\Vert }_{\mathrm{op}}$。
• 为使 $\mathcal{J}$ 严格成为命题 2 意义下的 Jensen 间隙，需要 ${\hat{\mathbf{x}}}_0=\mathbb{E}[{\mathbf{x}}_0\mid {\mathbf{x}}_t]$；若它只是普通估计，上界依然成立，但应理解为对该估计的误差控制。

综上，在所列条件下

$$\mathcal{J}\le L{C}_{\mathcal{A}}{m}_1,L<\infty .$$

Chung, H., Kim, J., McCann, M. T., Klasky, M. L., & Ye, J. C. (2022). Diffusion Posterior Sampling for General Noisy Inverse Problems. arXiv Preprint arXiv:2209.14687.