LDM: 在隐空间用diffusion model合成高质量的图片！
[论文地址] High-Resolution Image Synthesis with Latent Diffusion Models
[github] https://github.com/compvis/latent-diffusion
文章目录
LDM: 在隐空间用diffusion model合成高质量的图片！ 系列阅读 1 背景 2 方法 2.1 整体架构 2.2 更多细节 2.2.1 感知压缩的权衡 2.2.2 LDM的训练策略与预测 2.2.3 给生成过程引入控制信号 参考文献系列阅读
diffusion model（一）DDPM技术小结 (denoising diffusion probabilistic) diffusion model（二）—— DDIM技术小结 diffusion model（三）—— classifier guided diffusion model diffusion model（四）文生图diffusion model（classifier-free guided） diffusion model（五）stable diffusion底层原理（latent diffusion model， LDM1 背景
近期扩散模型（diffusion model，DM）在图像生成取得了最先进的结果。但是传统的扩散模型是在像素空间(pixel space)进行优化的，高分辨率图像的训练往往需要更大的显存，更多的梯度回传，导致DM模型训练成本高，推理时延长。当下很多工作从优化采样策略12、分层方法3来解决这两个缺点，但效果有限。latent diffusion model提供了一个新的思路：它将扩散过程从传统的像素空间转到了隐空间（DM模型输入参数量大大降低），极大提升了DM的训练效率和推理效率，使得DM模型得以在单张消费级显卡应用，降低了AI图片生成的上手成本。目前火出圈的Dalle24, Stable Diffusion5都用到了LDM技术。下面我们来看LDM是如何做的。
2 方法
2.1 整体架构
LDM主要参考了VQGAN6的思路，其整体架构如下图所示。与传统Diffusion model在像素空间(pixel-based diffusion model)重建不同的是，LDM是在隐空间进行重建。为了得到图片的隐空间表征，LDM需要预选训练一个VAE模型。记原始图像为 x , x ∈ R H × W × 3 x,x\in\mathbb{R}^{H \times W \times 3} x,x∈RH×W×3通过VAE的编码器 E \mathcal{E} E获得图片的隐空间表示 z , z ∈ R H f × W f × c z,z\in \mathbb{R}^{\frac{H}{f} \times \frac{W}{f} \times c} z,z∈RfH×fW×c， f f f为下采样率。在重建阶段，首先通过DM预测隐空间的表征 z ~ \tilde{z} z~，再将其送入到VAE的解码器中 D \mathcal{D} D重建像素空间的表征 x ~ \tilde{x} x~。本质上LDM是一种二阶段的图片生成方法。
通过在隐空间重建的操作，相较原本的像素空间重建，一个时间步的计算量近似降低了近 f 2 f^2 f2倍。有同学会问通道数不是从 3 → c 3 \rightarrow c 3→c了吗，为什么整体计算量没有扩到 c 3 \frac{c}{3} 3c倍，其实这只会增加第一层卷积的计算量，相较整体模型而言较低。
原本DM模型的优化目标
L
D
M
=
E
x
,
ϵ
∼
N
(
0
,
1
)
,
t
[
∥
ϵ
−
ϵ
θ
(
x
t
,
t
)
∥
2
2
]
L_{DM} = \mathbb{E}_{x, \epsilon \sim \mathcal{N}(0, 1), t} [\parallel \epsilon - \epsilon_\theta(x_t, t) \parallel_2^2]
LDM=Ex,ϵ∼N(0,1),t[∥ϵ−ϵθ(xt,t)∥22]
LDM的优化目标
L
L
D
M
=
E
E
(
x
)
,
ϵ
∼
N
(
0
,
1
)
,
t
[
∥
ϵ
−
ϵ
θ
(
z
t
,
t
)
∥
2
2
]
L_{LDM} = \mathbb{E}_{\mathcal{E}(x), \epsilon \sim \mathcal{N}(0, 1), t} [\parallel \epsilon - \epsilon_\theta(z_t, t) \parallel_2^2]
LLDM=EE(x),ϵ∼N(0,1),t[∥ϵ−ϵθ(zt,t)∥22]
t
t
t在
{
1
,
.
.
.
,
T
}
\{1, ..., T\}
{1,...,T}以服从均匀分布进行采样
通过上面的描述不难看出LDM就是套了一层VAE的DM！利用VAE搭建起隐空间和像素空间的桥梁。
2.2 更多细节
2.2.1 感知压缩的权衡
上文提到LDM先通过VAE的encode将图片从像素空间 x , x ∈ R H × W × 3 x,x\in\mathbb{R}^{H \times W \times 3} x,x∈RH×W×3变换到隐空间 z , z ∈ R H f × W f × c z,z\in \mathbb{R}^{\frac{H}{f} \times \frac{W}{f} \times c} z,z∈RfH×fW×c。当下采样率 f f f越大，DM的过程越快，但f越大可能导致丢失过多的高频信号，影响重建结果。因此需要对不同的下采样率进行权衡。论文中分别对 f ∈ { 1 , 2 , 4 , 8 , 16 , 32 } f \in \{1,2,4,8,16,32\} f∈{1,2,4,8,16,32}进行了试验。基于试验结果，作者建议采样率控制在 4 − 16 4-16 4−16能够取得效率和质量的均衡。
从下图可以看出，当采样率为4-16时，LDM的训练收益比最大。
2.2.2 LDM的训练策略与预测
LDM采用了二阶段的训练策略(two-stage)。需要先训练VAE模型再训练DM模型。DM训练时会利用已训练好的VAE的encode将数据从像素空间切换到隐空间。
预测时，在隐空间随机采样一个噪声，用DM进行去噪。最后将DM的预测结果传给VAE的Decode进行解码。
2.2.3 给生成过程引入控制信号
如果说不带控制信号的LDM是对 p ( z ) p(z) p(z)建模的话，那么加上控制信号的建模可以表述成 p ( z ∣ y ) p(z|y) p(z∣y)。控制信号 y y y可以是文本、图片布局、轮廓图等。当有有一些更为细粒度的生成存在多个控制信号，即对 p ( z ∣ y 1 , ⋯ , y N ) p(z|y_1, \cdots,y_N) p(z∣y1,⋯,yN)。为了使得生成过程考虑控制信息，作者在原有的Unet backbone上引入了一个交叉注意力机制，来融入控制信号。首先通过一个 domain specific encoder τ θ \text{domain specific encoder} \space \tau_\theta domain specific encoder τθ将 y y y投影到为一个中间表征 τ θ ( y ) ∈ R M × d τ \tau_\theta(y) \in \mathbb{R}^{M \times d_\tau} τθ(y)∈RM×dτ,随后和DM某层的输出进行融合。
假定
φ
i
(
z
t
)
\varphi_i(z_t)
φi(zt)是隐变量
z
z
z在Unet在层
i
i
i时间步为
t
t
t时的输出，下面需要将
φ
i
(
z
t
)
\varphi_i(z_t)
φi(zt)与
τ
θ
(
y
)
\tau_\theta(y)
τθ(y)用交叉注意力机制融合
A
t
t
e
n
t
i
o
n
(
Q
,
K
,
V
)
=
s
o
f
t
m
a
x
(
Q
K
T
d
)
⋅
V
Q
=
W
Q
(
i
)
⋅
φ
i
(
z
t
)
;
K
=
W
K
(
i
)
⋅
τ
θ
(
y
)
;
V
=
W
V
(
i
)
⋅
τ
θ
(
y
)
\begin{aligned} \mathrm{Attention}(Q, K, V) &= \mathrm{softmax}(\frac{QK^T}{\sqrt{d}})\cdot V \notag \\ Q = W_Q^{(i)} \cdot \varphi_i(z_t); \quad K &= W_K^{(i)} \cdot \tau_\theta(y); \quad V = W_V^{(i)} \cdot \tau_\theta(y) \end{aligned}
Attention(Q,K,V)Q=WQ(i)⋅φi(zt);K=softmax(d
QKT)⋅V=WK(i)⋅τθ(y);V=WV(i)⋅τθ(y)
此时模型的优化目标为
L
L
D
M
c
o
n
d
=
E
E
(
x
)
,
ϵ
∼
N
(
0
,
1
)
,
t
[
∥
ϵ
−
ϵ
θ
(
z
t
,
t
,
τ
θ
(
y
)
)
∥
2
2
]
L_{LDM_{cond}} = \mathbb{E}_{\mathcal{E}(x), \epsilon \sim \mathcal{N}(0, 1), t} [\parallel \epsilon - \epsilon_\theta(z_t, t, \tau_\theta(y)) \parallel_2^2]
LLDMcond=EE(x),ϵ∼N(0,1),t[∥ϵ−ϵθ(zt,t,τθ(y))∥22]
参考文献
On fast sampling of diffusion probabilistic models. ↩︎
Noise estimation for generative diffusion models. ↩︎
Cascaded diffusion models for high fidelity image generation ↩︎
Hierarchical Text-Conditional Image Generation with CLIP Latents,GitHub ↩︎
stable diffusion GitHub ↩︎
Taming transformers for high-resolution image synthesis ↩︎