A Variational Perspective On Solving Inverse Problems With Diffusion Models

Abstract

扩散模型已成为视觉领域基础模型的重要支柱之一。其关键应用之一是通过单一的扩散先验,不用为每个任务重新训练,就能普遍解决不同的下游逆任务。大多数逆任务可以表述为给定测量(例如,掩码图像)推断数据(例如,完整图像)的后验分布。然而,由于扩散过程的非线性和迭代性质,这在扩散模型中是具有挑战性的,因为后验分布是不可处理的。为应对这一挑战,我们提出了一种**变分方法,**旨在设计上逼近真实的后验分布。我们展示了我们的方法自然地引出了通过去噪扩散过程(RED-diff)进行正则化,其中不同时间步的去噪器同时对图像施加不同的结构约束。为了衡量来自不同时间步去噪器的贡献,我们提出了一种基于信噪比(SNR)的加权机制。我们的方法为用扩散模型解决逆问题提供了一个新的变分视角,使我们能够将采样表述为随机优化,其中可以简单地应用轻量迭代的现成求解器。

Instroduction

扩散模型,如Stable Diffusion(Rombach等,2021),正在成为当今视觉基础模型的重要组成部分。这些扩散模型的重要用途之一是将它们作为先验分布,用于在图像修复和渲染等各种下游逆问题中的采样。然而,这需要采样器满足以下要求:(i)通用且适应各种任务,无需为每个任务重新训练;(ii)高效且易于调节。最近有一些尝试开发用于逆问题的通用采样器;例如,DDRM(Kawar等,2022a)最初是为了扩展DDPM(Ho等,2020)以处理线性逆问题。它依赖于SVD将线性观测融入去噪过程中。然而,DDRM需要许多测量才能工作。随后,ΠGDM(Song等,2023)被引入以增强DDRM。ΠGDM的核心是通过反演从线性观测中获取指导,来增强去噪扩散评分。同样,DPS(Chung等,2022a)将评分修改框架扩展到一般(非线性)逆问题。然而,DPS和ΠGDM中的评分修改方法在很大程度上依赖于近似。实际上,反向扩散过程的非线性和递归性质使得后验分布变得相当难以处理且多峰。然而,DPS和ΠGDM依赖于评分的简单单峰近似,这在扩散过程的许多步骤中是相当松散的近似。

为了回避后验评分近似的挑战,我们提出了一种基于变分推断的完全不同的方法(Blei等,2017;Ahmed等,2012;Hoffman等,2013)。采用去噪扩散模型作为数据先验,并将测量模型表示为似然,我们使用变分推断来推断给定观测的后验数据分布。我们的方法本质上是使用KL散度将数据分布的模式与高斯分布匹配。这导致了一种简单的(加权)评分匹配标准,通过去噪扩散过程正则化观测的测量匹配损失。有趣的是,评分匹配正则化具有简单梯度的可解释形式。

这与Romano等人(2016年)提出的正则化去噪(RED)框架相似,在这个框架中,扩散过程中的不同阶段的去噪器施加了从高层语义到细节的不同结构约束。这是一个重要的联系,它将采样视为随机优化。因此,可以简单地使用现成的优化器库来进行采样,这使得推断过程既高效又可解释,并且容易调整。我们将这种方法命名为RED-diff。然而,值得注意的是,我们的框架与RED在几个方面有所不同:首先,我们从一个原则性的变分视角出发来推导我们的目标,这一视角已经被充分研究和理解;其次,我们的正则化利用了来自所有扩散步骤的反馈,涵盖了不同的噪声水平,而RED只使用了一个去噪模型。

为了使评分匹配正则化成功,不同时间步的去噪器需要进行适当的加权。为此,我们提出了一种基于每个时间步去噪信噪比(SNR)的加权机制,该机制在反向扩散过程中对早期步骤加权,并对后期时间步进行减权。为了验证这一提议,我们进行了各种线性和非线性逆问题的实验和消融研究。主要发现包括:

  1. RED-diff在图像逆问题中实现了比最先进的采样器更优越的图像保真度和感知质量;
  2. RED-diff具有轻量级迭代,不涉及DPS和ΠGDM中的评分Jacobian,因此更加内存高效且对GPU友好;
  3. 消融研究表明,优化器参数如学习率和步骤数量是调整保真度和感知质量之间权衡的合适调节器。

Related work

Diffusion models for inverse problems:

最近有几个研究工作将扩散模型以“即插即用”的方式应用于各个领域的逆问题,例如自然图像(Jalal等,2021;Kawar等,2022a;Song等,2023;Chung等,2022a;b;Graikos等,2022;Chung等,2023a;b)、医学图像(Jalal等,2021)和音频处理(Kong等,2020)。我们主要关注图像领域,这些工作在处理测量数据的方式上有所不同。

  • 作为早期工作的一部分,Kadkhodaie和Simoncelli(2021)以及Jalal等(2021)采用了Langevin动态方法来处理线性逆问题,并通过投影(Kadkhodaie和Simoncelli,2021)或最小二乘保真度的梯度(Jalal等,2021)来整合观测引导。其他一些工作采用DDPM(Ho等,2020)扩散模型,在扩散去噪和投影步骤之间交替进行(Choi等,2021;Chung等,2022c)。然而,这些迭代可能会累积误差,使得轨迹偏离先验流形,
  • MCG方法(Chung等,2022b)提出了一个额外的修正项,受流形约束的启发,以保持迭代接近流形。DDRM(Kawar等,2022a)扩展了DDPM,用矩阵SVD解决线性逆问题,但在测量有限的情况下效果不佳。
  • 为了克服这一不足,最近的方法旨在通过重建引导(Ho等,2022)来提供指导,这在DPS(Chung等,2022a)中进一步扩展到非线性逆问题。ΠGDM(Song等,2023)引入了伪逆引导,通过反转测量模型来改进引导近似。然而,它的范围仅限于线性和某些半线性任务(如JPEG(Kawar等,2022b))。然而,ΠGDM和DPS都严重依赖于对不可处理后验评分的近似,这在许多扩散过程的步骤中对非小噪声水平来说非常粗糙。
  • Graikos等(2022)最近还提出了一种不同的方法,该方法通过扩散误差损失来正则化逆问题的重建项。这类似于传统的即插即用先验(P3)方法(Venkatakrishnan等,2013),其起源于ADMM优化(Boyd等,2011)。然而,我们的方法在精神上更接近RED框架,这提供了更多的优化器和调节灵活性;(Romano等,2016;Cohen等,2021)。

Background

Denoising Diffusion Models

扩散模型(Sohl-Dickstein等,2015;Ho等,2020;Song等,2021b)包括两个过程:前向过程和反向过程。前向过程逐渐向输入图像添加噪声,而反向过程通过迭代去噪来学习生成图像。形式上,前向过程可以通过方差保持随机微分方程(VP-SDE)(Song等,2021b)来表达:
d x = − 1 2 β ( t ) x   d t + β ( t )   d W dx = -\frac{1}{2} \beta(t)x \, dt + \beta(t) \, dW dx=21β(t)xdt+β(t)dW

其中 t ∈ [ 0 , T ] t \in [0, T] t[0,T] β ( t ) = β min + ( β max − β min ) t T \beta(t) = \beta_{\text{min}} + (\beta_{\text{max}} - \beta_{\text{min}}) \frac{t}{T} β(t)=βmin+(βmaxβmin)Tt 重新缩放时间变量, d W dW dW 是标准Wiener过程。前向过程的设计使得在过程结束时, x T x_T xT 的分布收敛到标准高斯分布(即, x T ∼ N ( 0 , I ) x_T \sim \mathcal{N}(0, I) xTN(0,I))。

反向过程通过以下公式定义:
d x = − 1 2 β ( t ) x   d t − β ( t ) ∇ x log ⁡ p ( x t )   d t + β ( t )   d W ˉ dx = -\frac{1}{2} \beta(t)x \, dt - \beta(t) \nabla_x \log p(x_t) \, dt + \beta(t) \, d\bar{W} dx=21β(t)xdtβ(t)xlogp(xt)dt+β(t)dWˉ

其中 ∇ x log ⁡ p ( x t ) \nabla_x \log p(x_t) xlogp(xt) 是时间 t t t 处的扩散数据的评分函数, d W ˉ d\bar{W} dWˉ 是反向标准Wiener过程。

求解反向生成过程需要估计评分函数。在实际应用中,这通过从前向扩散过程中采样,并使用去噪评分匹配目标(Vincent, 2011)来训练评分函数来完成。具体来说,扩散样本的生成由以下公式给出:
x t = α t x 0 + σ t ϵ , ϵ ∼ N ( 0 , I ) , t ∈ [ 0 , T ] (1) x_t = \alpha_t x_0 + \sigma_t \epsilon, \quad \epsilon \sim \mathcal{N}(0, I), \quad t \in [0, T] \tag{1} xt=αtx0+σtϵ,ϵN(0,I),t[0,T](1)

其中 x 0 ∼ p data x_0 \sim p_{\text{data}} x0pdata 从数据分布中抽取, σ t = 1 − e − ∫ 0 t β ( s )   d s \sigma_t = \sqrt{1 - e^{- \int_{0}^{t} \beta(s) \, ds}} σt=1e0tβ(s)ds α t = 1 − σ t 2 \alpha_t = \sqrt{1 - \sigma_t^2} αt=1σt2 。我们用参数化的评分函数(即扩散模型)来表示为 ϵ θ ( x t ; t ) ≈ − σ t ∇ x log ⁡ p ( x t ) \epsilon_\theta(x_t; t) \approx -\sigma_t \nabla_x \log p(x_t) ϵθ(xt;t)σtxlogp(xt),其中参数为 θ \theta θ,我们可以通过如下的欧几里得损失混合来训练 ϵ θ ( x t ; t ) \epsilon_\theta(x_t; t) ϵθ(xt;t)
min ⁡ θ E x 0 ∼ p data ( x 0 ) , ϵ ∼ N ( 0 , I ) , t ∼ U [ 0 , T ] [ ∥ ϵ − ϵ θ ( x t ; t ) ∥ 2 2 ] \min_\theta \mathbb{E}_{x_0 \sim p_{\text{data}}(x_0), \epsilon \sim \mathcal{N}(0, I), t \sim U[0, T]} \left[ \| \epsilon - \epsilon_\theta(x_t; t) \|^2_2 \right] θminEx0pdata(x0),ϵN(0,I),tU[0,T][ϵϵθ(xt;t)22]

也可以使用其他针对 t t t 的损失加权函数。给定训练好的评分函数,可以使用DDPM(Ho等,2020),DDIM(Song等,2020)或其他求解器(Lu等,2022;Zhang & Chen,2022;Dockhorn等,2022)来生成样本。

Score Approximation For Inverse Problems

逆问题可以被表述为从一个(非线性且带噪声的)观测值中找出 x 0 x_0 x0
y = f ( x 0 ) + v , v ∼ N ( 0 , σ v 2 I ) (2) y = f(x_0) + v, \quad v \sim \mathcal{N}(0, \sigma_v^2 I) \tag{2} y=f(x0)+v,vN(0,σv2I)(2)

其中,前向(也称为测量)模型 f f f 是已知的。在许多应用中,例如图像修补,这是一个严重病态的问题,需要强大的先验来找到一个合理的解决方案。我们的目标是利用(预训练的)扩散模型提供的先验,以即插即用的方式有效地从条件后验中采样。我们将扩散模型施加的先验分布表示为 p ( x 0 ) p(x_0) p(x0)。测量模型可以表示为: p ( y ∣ x 0 ) = N ( f ( x 0 ) , σ v 2 ) p(y | x_0) = \mathcal{N}(f(x_0), \sigma_v^2) p(yx0)=N(f(x0),σv2)

解决逆问题的目标是从后验分布 p ( x 0 ∣ y ) p(x_0 | y) p(x0y) 中采样。

正如我们在前面部分讨论的那样,扩散模型依赖于估计的评分函数来生成样本。在存在观测值 y y y 的情况下,只要可以获得所有扩散步骤的条件评分的近似,就可以用它们生成合理的 x 0 ∼ p ( x 0 ∣ y ) x_0 \sim p(x_0 | y) x0p(x0y)。这就是 ΠGDM(Song等,2023)和 DPS(Chung等,2022a)的核心思想。

具体来说,基于贝叶斯规则,条件评分 ∇ x log ⁡ p ( x t ∣ y ) \nabla_x \log p(x_t | y) xlogp(xty) 可以简单地得到为:
∇ x log ⁡ p ( x t ∣ y ) = ∇ x log ⁡ p ( y ∣ x t ) + ∇ x log ⁡ p ( x t ) (3) \nabla_x \log p(x_t | y) = \nabla_x \log p(y | x_t) + \nabla_x \log p(x_t)\tag{3} xlogp(xty)=xlogp(yxt)+xlogp(xt)(3)
总评分是模型似然和先验评分的叠加。虽然 ∇ x log ⁡ p ( x t ) \nabla_x \log p(x_t) xlogp(xt) 可以从预训练的扩散模型中轻松获得,但似然评分却非常困难,且在没有特定任务训练的情况下无法估计。这可以从以下事实看出:
p ( y ∣ x t ) = ∫ p ( y ∣ x 0 ) p ( x 0 ∣ x t )   d x 0 p(y | x_t) = \int p(y | x_0) p(x_0 | x_t) \, dx_0 p(yxt)=p(yx0)p(x0xt)dx0

尽管 p ( y ∣ x 0 ) p(y | x_0) p(yx0) 采用简单的高斯形式,但去噪分布 p ( x 0 ∣ x t ) p(x_0 | x_t) p(x0xt) 可以非常复杂且多模态(Xiao等,2022)。因此, p ( y ∣ x t ) p(y | x_t) p(yxt) 也可能非常复杂。为了规避这个问题,先前的工作(Song等,2023;Chung等,2022a;Kadkhodaie & Simoncelli,2021;Ho等,2022)依赖于在 MMSE(最小均方误差)估计点附近对 p ( x 0 ∣ x t ) p(x_0|x_t) p(x0xt) 的高斯近似:
E [ x 0 ∣ x t ] = 1 α t ( x t − σ t ϵ θ ( x t , t ) ) (4) \mathbb{E}[x_0 | x_t] = \frac{1}{\alpha_t} \left(x_t - \sigma_t \epsilon_\theta(x_t, t) \right) \tag{4} E[x0xt]=αt1(xtσtϵθ(xt,t))(4)

根据公式(1): x t = α t x 0 + σ t ϵ , ϵ ∼ N ( 0 , I ) , t ∈ [ 0 , T ] x_t = \alpha_t x_0 + \sigma_t \epsilon, \quad \epsilon \sim \mathcal{N}(0, I), \quad t \in [0, T] xt=αtx0+σtϵ,ϵN(0,I),t[0,T]

Variational Diffusion Sampling

在本节中,我们介绍了我们在解决逆问题方面的变分视角。为了应对之前方法在采样条件后验分布 p ( x 0 ∣ y ) p(x_0 | y) p(x0y) 时存在的不足,我们提出了一种基于KL最小化的变分方法:
min ⁡ q KL ( q ( x 0 ∣ y ) ∥ p ( x 0 ∣ y ) ) (5) \min_q \text{KL}(q(x_0 | y) \| p(x_0 | y)) \tag{5} qminKL(q(x0y)p(x0y))(5)

其中 q q q 是一个变分分布,定义为 q : = N ( μ , σ 2 I ) q := \mathcal{N}(\mu, \sigma^2 I) q:=N(μ,σ2I)。分布 q q q 试图在数据分布中找到最主要的模式,以匹配观测值。可以很容易地证明,方程 (5) 中的KL目标可以展开为:
KL ( q ( x 0 ∣ y ) ∥ p ( x 0 ∣ y ) ) = − E q ( x 0 ∣ y ) [ log ⁡ p ( y ∣ x 0 ) ] + KL ( q ( x 0 ∣ y ) ∥ p ( x 0 ) ) + log ⁡ p ( y ) (6) \text{KL}(q(x_0 | y) \| p(x_0 | y)) = - \mathbb{E}_{q(x_0 | y)} \left[ \log p(y | x_0) \right] + \text{KL}(q(x_0 | y) \| p(x_0)) + \log p(y) \tag{6} KL(q(x0y)p(x0y))=Eq(x0y)[logp(yx0)]+KL(q(x0y)p(x0))+logp(y)(6)

其中:

  • 项 (i) 是变分界限,通常用于训练变分自编码器(Kingma & Welling, 2013;Rezende et al., 2014)。
  • 项 (ii) 是观测似然,对 q q q 来说是常数。

因此,为了最小化方程 (5) 中的KL散度,相对于 q q q,只需最小化方程 (6) 中的变分界限(项 (i))。

这引出了下一个命题:

命题 1 假设评分函数被准确地学习,即 ϵ θ ( x t ; t ) = − σ t ∇ x t log ⁡ p ( x t ) \epsilon_\theta(x_t; t) = -\sigma_t \nabla_{x_t} \log p(x_t) ϵθ(xt;t)=σtxtlogp(xt)。那么,相对于 q q q 的KL最小化等同于最小化变分界限(方程 (6) 中的项 (i)),该界限本身遵循评分匹配损失:
min ⁡ μ , σ E q ( x 0 ∣ y ) [ ∥ y − f ( x 0 ) ∥ 2 2 2 σ v 2 ] + ∫ 0 T w ~ ( t ) E q ( x t ∣ y ) ( ∥ ∇ x t log ⁡ q ( x t ∣ y ) − ∇ x t log ⁡ p ( x t ) ∥ 2 2 ) d t (7) \min_{\mu, \sigma} \mathbb{E}_{q(x_0 | y)} \left[ \frac{\| y - f(x_0) \|^2_2}{2 \sigma_v^2} \right] + \int_0^T\tilde w(t)\mathbb E_{q(x_t|y)} \left(\left\| \nabla_{x_t} \log q(x_t | y) - \nabla_{x_t} \log p(x_t) \right\|^2_2 \right) dt \tag{7} μ,σminEq(x0y)[2σv2yf(x0)22]+0Tw~(t)Eq(xty)(xtlogq(xty)xtlogp(xt)22)dt(7)

其中 q ( x t ∣ y ) = N ( α t μ , α t 2 σ 2 + σ t 2 I ) q(x_t | y) = \mathcal{N}(\alpha_t \mu, \alpha_t^2 \sigma^2 + \sigma_t^2 I) q(xty)=N(αtμ,αt2σ2+σt2I) 通过从 q ( x 0 ∣ y ) q(x_0 | y) q(x0y) 中采样 x 0 x_0 x0 并应用方程 (1) 中的前向过程来生成样本 x t x_t xt,而 ω ~ ( t ) = β ( t ) 2 \tilde{\omega}(t) = \frac{\beta(t)}{2} ω~(t)=2β(t) 是损失加权项。


命题 1 的证明

如第 4 节所讨论的,通过贝叶斯规则,可以将方程 (5) 中的 KL 目标重新写为: KL ( q ( x 0 ∣ y ) ∥ p ( x 0 ∣ y ) ) = − E q ( x 0 ∣ y ) [ log ⁡ p ( y ∣ x 0 ) ] + KL ( q ( x 0 ∣ y ) ∥ p ( x 0 ) ) \text{KL}(q(x_0|y) \| p(x_0|y)) = - \mathbb{E}_{q(x_0|y)} \left[ \log p(y|x_0) \right] + \text{KL}(q(x_0|y) \| p(x_0)) KL(q(x0y)p(x0y))=Eq(x0y)[logp(yx0)]+KL(q(x0y)p(x0))

对于最小化目的,项 (ii) 是一个常量,我们可以忽略它。然而,项 (i) 还有两个部分。第一部分就是重建损失。根据主论文中的方程 (2) 中的测量模型,由于我们假设噪声是独立同分布的高斯噪声,第一部分可以简单地推导为:
E q ( x 0 ∣ y ) [ log ⁡ p ( y ∣ x 0 ) ] = − 1 2 σ v 2 E q ( x 0 ∣ y ) [ ∥ y − f ( x 0 ) ∥ 2 ] \mathbb{E}_{q(x_0|y)} \left[ \log p(y|x_0) \right] = - \frac{1}{2 \sigma_v^2} \mathbb{E}_{q(x_0|y)} \left[ \| y - f(x_0) \|^2 \right] Eq(x0y)[logp(yx0)]=2σv21Eq(x0y)[yf(x0)2]

使用 Song et al. (2021a) 中的定理 2,假设评分函数完全学习到,即 ϵ θ ( x t ; t ) = − σ t ∇ x t log ⁡ p ( x t ) \epsilon_\theta(x_t; t) = -\sigma_t \nabla_{x_t} \log p(x_t) ϵθ(xt;t)=σtxtlogp(xt),在对 log ⁡ q ( x t ∣ y ) \log q(x_t | y) logq(xty) p ( x t ) p(x_t) p(xt) 在无穷大处增长的某些温和假设下,我们有:

KL ( q ( x 0 ∣ y ) ∥ p ( x 0 ) ) = ∫ 0 T β ( t ) 2 E q ( x t ∣ y ) [ ∥ ∇ x t log ⁡ q ( x t ∣ y ) − ∇ x t log ⁡ p ( x t ) ∥ 2 ]   d t (12) \text{KL}(q(x_0|y) \| p(x_0)) = \int_0^T \frac{\beta(t)}{2} \mathbb{E}_{q(x_t|y)} \left[ \| \nabla_{x_t} \log q(x_t | y) - \nabla_{x_t} \log p(x_t) \|^2 \right] \, dt \tag{12} KL(q(x0y)p(x0))=0T2β(t)Eq(xty)[xtlogq(xty)xtlogp(xt)2]dt(12)

这是在正值 β ( t ) \beta(t) β(t) 的扩散去噪轨迹 { x t } \{x_t\} {xt} 上的。实际上,这意味着在连续去噪扩散轨迹上的加权评分匹配等于 KL 散度。然而,在实际应用中,我们通常感兴趣的是方程 (12) 中右侧的重加权,这会导致其他散度度量(见 Song et al., 2021a)。


上面提到的第一个项是测量匹配损失(即重建损失),它是通过定义 p ( y ∣ x 0 ) p(y|x_0) p(yx0) 获得的,而第二个项是通过将KL散度的项展开成评分匹配目标获得的,如(Vahdat et al., 2021; Song et al., 2021a)中所示, ω ~ ( t ) = β ( t ) 2 \tilde{\omega}(t) = \frac{\beta(t)}{2} ω~(t)=2β(t) 是基于最大似然的加权。

第二个项可以视为由扩散先验施加的评分匹配正则化项。该积分在扩散轨迹上进行评估,即 x t ∼ q ( x t ∣ y ) x_t \sim q(x_t | y) xtq(xty) 对于 t ∈ [ 0 , T ] t \in [0, T] t[0,T],这是应用于 q ( x 0 ∣ y ) q(x_0 | y) q(x0y) 的前向扩散过程。由于 q ( x 0 ∣ y ) q(x_0 | y) q(x0y) 采用简单的高斯形式,我们可以证明 q ( x t ∣ y ) q(x_t | y) q(xty) 也是高斯形式,表示为 q ( x t ∣ y ) = N ( α t μ , α t 2 σ 2 + σ t 2 ) q(x_t | y) = \mathcal{N}(\alpha_t \mu, \alpha_t^2 \sigma^2 + \sigma_t^2) q(xty)=N(αtμ,αt2σ2+σt2)(见 Vahdat et al., 2021)。因此,评分函数 ∇ x t log ⁡ q ( x t ∣ y ) \nabla_{x_t} \log q(x_t | y) xtlogq(xty) 可以通过解析方法计算。

假设变分分布的方差是一个接近于零的小常数(即 σ ≈ 0 \sigma \approx 0 σ0),那么方程 (7) 中的优化问题可以进一步简化为:
min μ { ∥ y − f ( μ ) ∥ recon 2 + ∫ 0 T 2 ω ( t ) σ v 2 σ t 2 ∥ ϵ θ ( x t ; t ) − ϵ ∥ 2 2   d t } (8) \text{min}_{\mu} \left\{ \| y - f(\mu) \|^2_{\text{recon}} + \int_0^T 2\omega(t) \frac{\sigma_v^2}{\sigma_t^2} \| \epsilon_\theta(x_t; t) - \epsilon \|^2_{2} \, dt \right\} \tag{8} minμ{yf(μ)recon2+0T2ω(t)σt2σv2ϵθ(xt;t)ϵ22dt}(8)

其中 x t = α t μ + σ t ϵ x_t = \alpha_t \mu + \sigma_t \epsilon xt=αtμ+σtϵ。简而言之,解决上述优化问题将会找到一个图像 μ \mu μ,使得在测量模型 f f f 下重建观测 y y y 的效果最好,同时在正则化项施加的先验下具有高的似然性。

Remark [无噪声观测] 如果观测噪声 σ v = 0 \sigma_v = 0 σv=0,则从方程 (6) 可知,重建项简化为一个硬约束,可以表示为指示函数 I y = f ( μ ) \mathbb{I}_{y = f(\mu)} Iy=f(μ),当 y = f ( μ ) y = f(\mu) y=f(μ) 时为零,否则为无穷大。然而,在实践中,我们仍然可以使用方程 (7) 并设置一个较小的 σ v \sigma_v σv 作为近似。

Sampling As Stochastic Optimization

正则化的评分匹配目标(方程 (8))使我们能够将采样问题公式化为优化问题,从而解决逆问题。实质上,不同扩散步骤上的集成损失支持将随机优化作为一种适合的采样策略。

然而,在实际应用中,权重项 ω ~ ( t ) \tilde{\omega}(t) ω~(t) 的选择在优化问题的成功中起着关键作用。多个关于训练扩散模型的先前工作(Ho et al., 2020; Vahdat et al., 2021; Karras et al., 2022; Choi et al., 2022)发现,对 t t t 上目标函数的重加权在不同扩散步骤中平衡内容与细节方面发挥了关键作用,我们在本研究中也观察到了这一点(更多信息见第 4.3 节)。此外,方程 (8) 中标记为“reg”的第二项需要通过预训练的评分函数进行反向传播,这可能会导致优化过程缓慢和不稳定。接下来,我们考虑一种通用的加权机制 ω ~ ( t ) = β ( t ) ω ( t ) 2 \tilde{\omega}(t) = \frac{\beta(t) \omega(t)}{2} ω~(t)=2β(t)ω(t),其中 ω ( t ) \omega(t) ω(t) 是一个正值函数。我们将展示,如果选择的权重使得 ω ( 0 ) = 0 \omega(0) = 0 ω(0)=0,则可以高效地计算正则化项的梯度,而不需要通过预训练的评分函数进行反向传播。

命题 2. 如果 ω ( 0 ) = 0 \omega(0) = 0 ω(0)=0 σ = 0 \sigma = 0 σ=0,则评分匹配正则化的梯度为:
∇ μ reg ( μ ) = E t ∼ U [ 0 , T ] , ϵ ∼ N ( 0 , I ) [ λ t ( ϵ θ ( x t ; t ) − ϵ ) ] \nabla_\mu \text{reg}(\mu) = \mathbb{E}_{t \sim U[0, T], \epsilon \sim \mathcal{N}(0, I)} \left[ \lambda_t \left( \epsilon_\theta(x_t; t) - \epsilon \right) \right] μreg(μ)=EtU[0,T],ϵN(0,I)[λt(ϵθ(xt;t)ϵ)]
其中 λ t = 2 T σ v 2 α t σ t d ω ( t ) d t \lambda_t = \frac{2T \sigma_v^2 \alpha_t}{\sigma_t} \frac{d\omega(t)}{dt} λt=σt2Tσv2αtdtdω(t)

一阶随机优化器。 基于命题 2 中评分匹配正则化梯度的简单表达式,我们可以将时间视为一个均匀随机变量。因此,通过对时间和噪声进行随机采样,我们可以轻松获得梯度的无偏估计。因此,可以应用一阶随机优化方法来搜索 μ \mu μ。我们在算法 1 中列出了迭代步骤。注意,我们基于瞬时梯度定义每个时间步的损失,这可以视为线性损失的梯度。我们引入了 (sg) 符号,表示在优化过程中评分不进行反向传播。第 D.4 节中的消融实验表明,像 DDPM 和 DDIM 等标准反向扩散采样器中从 t = T t = T t=T t = 0 t = 0 t=0 的时间递减在实践中表现优于随机时间采样。

Remark [非零扩散]. 注意,命题 2 为简化起见推导了零扩散情况(即 σ = 0 \sigma = 0 σ=0)的梯度。对非零扩散的扩展请参见附录 A.3。

Regularization By Denoising

请注意,我们的变分采样器与 正则化通过去噪(RED) 框架(Romano et al., 2016)具有相似之处。实质上,RED 是一种灵活的方法,用于利用给定的去噪引擎来处理一般的逆问题。RED 正则化有效地根据某种图像自适应的拉普拉斯先验来促进图像的平滑。为了更好地理解与 RED 的关系,让我们看一下我们提出的变分采样器在每个时间步的损失。从命题 2 中的梯度表达式来看,形成时间步 t t t 的损失是有用的:

∥ y − f ( μ ) ∥ 2 + λ t ( sg ( ϵ θ ( x t ; t ) − ϵ ) ) T μ (9) \| y - f(\mu) \|^2 + \lambda_t \left( \text{sg}(\epsilon_\theta(x_t; t) - \epsilon) \right)^T \mu \tag{9} yf(μ)2+λt(sg(ϵθ(xt;t)ϵ))Tμ(9)

这个正则化项与 RED 相似。较小的正则化意味着要么扩散过程达到一个固定点,即 ϵ θ ( x t ; t ) = ϵ \epsilon_\theta(x_t; t) = \epsilon ϵθ(xt;t)=ϵ,要么残差仅包含噪声,而不再对图像贡献任何信息。然而,需要注意的是,这里不需要对称雅可比矩阵或原始 RED 中所需的假设,因为公式 9 的梯度自然是 ϵ θ ( x t ; t ) − ϵ \epsilon_\theta(x_t; t) - \epsilon ϵθ(xt;t)ϵ(注意停止梯度操作 sg)。尽管如此,RED 与我们的方法有根本的不同,包括扩散先验的生成性质,以及我们使用整个扩散轨迹进行正则化的事实。尽管如此,我们相信这是一个重要的联系,可以利用 RED 的工具来改进扩散模型在逆问题中的采样。同时值得提到的是,早期的研究(Reehorst & Schniter, 2018)也基于单个(确定性)去噪器探索了 RED 和评分匹配之间的联系。

Weighting Mechanism

时间步加权在训练扩散模型中起着关键作用。不同的时间步负责生成不同的结构,从最后的时间步生成大尺度内容到早期时间步生成精细的细节(Choi et al., 2022)。因此,为了有效的正则化,关键是要适当地调整算法 1 中的去噪器权重 { λ t } \{ \lambda_t \} {λt}。我们观察到,公式 9 中的正则化项对噪声调度非常敏感。例如,在方差保持的情况下,当 t t t 接近零时,正则化项会剧烈增加。

为了减轻正则化对权重的敏感性,更希望在信号域中定义正则化,这与拟合项兼容,如下所示:

∥ y − f ( μ ) ∥ 2 + λ ∥ sg ( μ − μ ^ t ) ∥ J \| y - f(\mu) \|^2 + \lambda \|\text{sg}(\mu - \hat{\mu}_t)\|_J yf(μ)2+λsg(μμ^t)J

其中, λ \lambda λ 是一个超参数,用于平衡先验和似然之间的权衡,而 μ ^ t \hat{\mu}_t μ^t 是干净数据的最小均方误差(MMSE)预测器。这里,我们希望常数 λ \lambda λ 来控制偏差(对观察的拟合)和方差(对先验的拟合)之间的权衡。为了得出公式 10 中的可解释损失,需要重新缩放噪声残差项 ϵ θ ( x t ; t ) − ϵ \epsilon_\theta(x_t; t) - \epsilon ϵθ(xt;t)ϵ

回顾一下,时间步 t t t 的去噪器观察到 x t = α t x 0 + σ t ϵ x_t = \alpha_t x_0 + \sigma_t \epsilon xt=αtx0+σtϵ。MMSE 估计器还提供去噪为:

μ ^ t = E μ ∣ x t [ μ ∣ x t ] = 1 α t ( x t − σ t ϵ θ ( x t ; t ) ) (10) \hat{\mu}_t = \mathbb{E}_{\mu | x_t} \left[ \mu \mid x_t \right] = \frac{1}{\alpha_t} \left( x_t - \sigma_t \epsilon_\theta(x_t; t) \right) \tag{10} μ^t=Eμxt[μxt]=αt1(xtσtϵθ(xt;t))(10)

因此,可以得到:

μ − μ ^ t = σ t α t ( ϵ θ ( x t ; t ) − ϵ ) (11) \mu - \hat{\mu}_t = \frac{\sigma_t}{\alpha_t} \left( \epsilon_\theta(x_t; t) - \epsilon \right) \tag{11} μμ^t=αtσt(ϵθ(xt;t)ϵ)(11)

这里我们定义了信噪比 SNR t : = α t σ t \text{SNR}_t := \frac{\alpha_t}{\sigma_t} SNRt:=σtαt。因此,通过选择 λ t = λ SNR t \lambda_t = \frac{\lambda}{\text{SNR}_t} λt=SNRtλ,我们可以将公式 9 中的噪声预测公式简单地转换为公式 10 中的干净数据公式。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值