万字长文解读深度学习——VQ-VAE和VQ-VAE-2
🌺历史文章列表🌺
推荐阅读:
图像生成发展起源:从VAE、VQ-VAE、扩散模型DDPM、DETR到ViT、Swin transformer
文生图模型演进:AE、VAE、VQ-VAE、VQ-GAN、DALL-E 等 8 模型
AE(自动编码器)、VAE(变分自动编码器)、VQ-VAE(向量量化变分自编码器) 的区别和联系?
Understanding VQ-VAE (DALL-E Explained Pt. 1)
AE和VAE
参考:深度学习——AE、VAE
VQ-VAE
论文:Neural Discrete Representation Learning
VQ-VAE(Vector Quantized Variational AutoEncoder,向量量化变分自编码器) 主要是将连续潜在空间的点映射到最近的一组离散的向量(即codebook中的向量)。模型引入了离散潜在空间的思想,克服了传统 VAE 中连续潜在空间表示的局限性,能够有效学习高质量的离散特征表示。
传统 VAE 的问题
-
连续潜在空间的限制:
- VAE 的潜在变量 z z z 是连续值,这会导致模型生成的表示较为分散、不够紧凑,无法高效捕获复杂数据的离散结构(如图像中的清晰边缘、重复纹理,或离散的语音特征)。
-
后验坍塌问题:
- 潜在变量的表示能力未被充分利用。指编码器生成的潜在表示 z z z 对解码器的输出贡献非常小,可能部分或完全被忽略。
- 当 KL 散度正则化过强时,编码器可能输出接近于先验分布(如 N ( 0 , 1 ) \mathcal{N}(0, 1) N(0,1)),导致潜在变量 z z z 的信息丢失。
VQ-VAE 与 VAE 的对比
特点 | VAE | VQ-VAE |
---|---|---|
潜在空间 | 连续空间 | 离散空间 |
潜在变量 z z z | 每一维是连续的实数值,包括所有的有理数(如整数、小数和分数)以及无理数 | 每一维是离散的整数 |
潜在分布建模 | 高斯分布 | 离散分布(通过 codebook 表示) |
VQ-VAE 的主要改进
-
解决后验坍塌问题:
- 通过 codebook 的离散化潜在变量,解码器被迫使用潜在变量的所有信息,避免后验坍塌。
-
增强生成质量:
- VAE 生成的图像质量相对有限,而 VQ-VAE 可以生成更高分辨率和更清晰的图像。
-
为后续模型奠定基础:
- VQ-VAE 的离散表示为 DALL·E 和 VQ-GAN 等模型的开发提供了关键技术支持。
VQ-VAE 的核心思想
关键点:
- VQ-VAE 的最大特点是引入了 向量量化(Vector Quantization),将连续的潜在变量 z z z 离散化。这有助于模型捕捉和表示更加丰富和复杂的数据分布。
- 将编码器输出的连续潜在向量 z e ( x ) z_e(x) ze(x) 映射到离散的代码簇(codebook)中最近的离散向量 e k e_k ek,编码范围更加可控。
Codebook 是嵌入空间中一组可学习的向量集合,用来将连续的编码器输出 z e ( x ) z_e(x) ze(x) 映射为离散表示 z q ( x ) z_q(x) zq(x)。这些嵌入向量也称为 离散嵌入向量,它们定义了潜在空间中离散值的分布。
高效的重建与生成:
- 编码器将输入映射为潜在空间的连续值。
- 向量量化将这些连续值投影到离散的代码簇中。
- 解码器根据离散化的潜在变量重建输入。
VQ-VAE 模型结构
VQ-VAE 与 VAE 的结构非常相似,只是中间部分不是学习概率分布,而是换成 VQ 来学习 Codebook。
VQ-VAE 的整体结构如下:
Figure 1:
- 左侧:描述 VQ-VAE 的图示。编码器 z ( x ) z(x) z(x) 的输出被映射到最近的嵌入点 e 2 e_2 e2。梯度 ∇ z L \nabla_z L ∇zL(红色箭头)将推动编码器改变输出,从而在下一次前向传播中调整配置。
- 右侧:嵌入空间的可视化。编码器输出 z e ( x ) z_e(x) ze(x) 被映射到最近的离散点 e 2 e_2 e2。
VQ-VAE 的核心思想是输入的图片通过编码器、向量量化模块和解码器的协同作用,将连续潜在变量映射为离散的潜在表示,并最终重建输入数据,得到新图像。以下是模型的主要流程:
1. 编码器 (Encoder)
- 输入:左侧图片(例如狗的图像)被输入到 卷积神经网络 (CNN) 编码器中。
- 输出:编码器将输入数据
x
x
x 转化为连续的潜在变量
z
e
(
x
)
z_e(x)
ze(x)【图中绿色】:
z e ( x ) = Encoder ( x ) z_e(x) = \text{Encoder}(x) ze(x)=Encoder(x)- z e ( x ) z_e(x) ze(x) 是一个连续高维向量,表示输入数据的特征。
- 从图中可以看到,输出 z e ( x ) z_e(x) ze(x) 是一个 D D D 维特征张量。
2. 向量量化层 (Vector Quantization Layer)
- 嵌入空间:
- 嵌入空间由一个 代码簇 (Codebook) 构成【紫色】,包含 K K K 个离散的嵌入向量 { e 1 , e 2 , … , e k } \{e_1, e_2, \dots, e_k\} {e1,e2,…,ek},每个向量是可学习的。
- 嵌入空间用于将编码器的连续表示 z e ( x ) z_e(x) ze(x) 离散化。
Codebook 是嵌入空间中一组可学习的向量集合,用来将连续的编码器输出 z e ( x ) z_e(x) ze(x) 映射为离散表示 z q ( x ) z_q(x) zq(x)。这些嵌入向量也称为 离散嵌入向量,它们定义了潜在空间中离散值的分布。
Codebook 的初始化:使用均匀分布初始化嵌入向量,范围为 [ − 1 num_embeddings , 1 num_embeddings ] [-\frac{1}{\text{num\_embeddings}}, \frac{1}{\text{num\_embeddings}}] [−num_embeddings1,num_embeddings1],保证稳定性和覆盖性。Codebook 大小对模型性能有较大影响,过小可能导致欠拟合,过大则增加计算开销。VQ-VAE 的论文中,num_embeddings 的典型值在 512 到 4096 之间,具体选择取决于任务和数据集。
-
量化过程:
- 对编码器输出
z
e
(
x
)
z_e(x)
ze(x) 中的每一个特征向量,找到嵌入空间中与之最近的嵌入向量
e
k
e_k
ek【紫色】。
z q ( x ) = arg min e k ∥ z e ( x ) − e k ∥ 2 z_q(x) = \arg\min_{e_k} \|z_e(x) - e_k\|_2 zq(x)=argekmin∥ze(x)−ek∥2 - 得到量化后的潜在表示 z q ( x ) z_q(x) zq(x)【紫色】,每个元素对应一个离散的代码【青色】。
- 对编码器输出
z
e
(
x
)
z_e(x)
ze(x) 中的每一个特征向量,找到嵌入空间中与之最近的嵌入向量
e
k
e_k
ek【紫色】。
-
梯度反向传播:
- 红色箭头 ∇ z L \nabla_z L ∇zL 表示通过量化后的误差对编码器进行反向传播,以调整编码器的输出。
- 代码簇本身也会更新,以更好地表示数据。
最近邻搜索通过计算嵌入空间 { e 1 , e 2 , … , e k } \{e_1, e_2, \dots, e_k\} {e1,e2,…,ek} 中每个嵌入向量与 z e ( x ) z_e(x) ze(x) 的欧几里得距离,选择最近的嵌入向量 e k e_k ek:
z q ( x ) = arg min e k ∥ z e ( x ) − e k ∥ 2 z_q(x) = \arg\min_{e_k} \|z_e(x) - e_k\|_2 zq(x)=argekmin∥ze(x)−ek∥2 最近邻搜索是一个不可导操作,因为 arg min \arg\min argmin 或 arg max \arg\max argmax 是一个离散操作,它选择一个代码簇索引 k k k,这个操作本身不具有梯度,因此编码器的参数 θ \theta θ, 包括 CNN的所有权重(weights)和偏置(biases),无法直接通过反向传播更新。为了绕过最近邻搜索的不可导问题,VQ-VAE 使用了 Stop Gradient 操作。具体步骤如下:
1. 停止梯度传播(Stop Gradient)
- 当计算嵌入向量 z q ( x ) z_q(x) zq(x) 时,停止对最近邻搜索的梯度传播。
- 操作:
- z q ( x ) z_q(x) zq(x) 直接从 codebook 中选取最近的嵌入 e k e_k ek,但不对 e k e_k ek 的选择过程传播梯度。
- 换句话说, z q ( x ) z_q(x) zq(x) 的选择被认为是固定的。
2. 梯度的复制
- 解码器的损失对 z q ( x ) z_q(x) zq(x) 的梯度 ∇ z L \nabla_z L ∇zL 会通过 红色线 被直接复制到编码器的输出 z e ( x ) z_e(x) ze(x) 上。
- 这样,虽然最近邻搜索不可导,但解码器提供的损失信号仍然能够通过 z q ( x ) z_q(x) zq(x) 的路径更新编码器参数。
3. 对嵌入向量的优化
- 嵌入空间中的代码簇 { e 1 , e 2 , … , e k } \{e_1, e_2, \dots, e_k\} {e1,e2,…,ek} 是可学习的。
- 使用以下损失更新嵌入向量: L vq = ∥ sg [ z e ( x ) ] − e k ∥ 2 L_{\text{vq}} = \| \text{sg}[z_e(x)] - e_k \|^2 Lvq=∥sg[ze(x)]−ek∥2
- 这里 sg [ z e ( x ) ] \text{sg}[z_e(x)] sg[ze(x)] 停止梯度传播,嵌入向量 e k e_k ek 会被直接更新。
红色线的意义 在图中,红色线的主要作用是通过停止梯度传播的技巧,让解码器的梯度绕过不可导的最近邻搜索,直接作用在编码器的输出 z e ( x ) z_e(x) ze(x) 上,从而实现编码器的参数更新:
梯度传播路径:
- 从解码器计算的重建损失 ∇ L rec \nabla L_{\text{rec}} ∇Lrec 沿着 z q ( x ) z_q(x) zq(x) 的路径,跳过最近邻搜索的不可导部分,直接作用在 z e ( x ) z_e(x) ze(x) 上。
- 通过这种方式,编码器仍然能够接收有效的梯度信号,从而调整其参数。
编码器更新:
- z e ( x ) z_e(x) ze(x) 被优化为更好地靠近其最近的嵌入向量 e k e_k ek,减少量化误差。
总结
- ∇ z L \nabla_z L ∇zL 是解码器的损失对 z q ( x ) z_q(x) zq(x) 的梯度,但由于使用 Stop Gradient 操作,它被传递到编码器的输出 z e ( x ) z_e(x) ze(x) 上。
- 它是解码器反向传播过程中提供的梯度信号,用于间接优化编码器的输出 z e ( x ) z_e(x) ze(x) 和参数 θ \theta θ。
- ∇ z L \nabla_z L ∇zL 的最终目标是减少重建误差 L rec L_{\text{rec}} Lrec 和量化误差 ∥ z e ( x ) − e k ∥ 2 \|z_e(x) - e_k\|^2 ∥ze(x)−ek∥2,提升模型的重建能力和表示质量。
3. 解码器 (Decoder)
-
输入:量化后的离散潜在表示 z q ( x ) z_q(x) zq(x) 被传递给解码器。
-
解码器使用 CNN 将离散的潜在变量 z q ( x ) z_q(x) zq(x) 转化为重建图像 x ^ \hat{x} x^:
x ^ = Decoder ( z q ( x ) ) \hat{x} = \text{Decoder}(z_q(x)) x^=Decoder(zq(x)) -
输出:最终生成的重建图像 x ^ \hat{x} x^ 与输入图像 x x x 尽可能相似。
4. 嵌入空间的动态调整
- 嵌入空间的作用:
- 从右侧的嵌入空间可视化可以看到,编码器输出 z e ( x ) z_e(x) ze(x) 被映射到离嵌入点最近的 e 2 e_2 e2。
- 梯度 ∇ z L \nabla_z L ∇zL 将推动编码器的输出朝向最近的嵌入点 e 2 e_2 e2 移动,以减少量化误差。
损失函数
1. 总损失函数
VQ-VAE 的总损失函数为:
L
=
L
rec
+
L
vq
+
β
L
commit
L = L_{\text{rec}} + L_{\text{vq}} + \beta L_{\text{commit}}
L=Lrec+Lvq+βLcommit
其中:
- L rec L_{\text{rec}} Lrec:重建损失,优化解码器,衡量输入 x x x 和重建 x ^ \hat{x} x^ 的差异,提升重建质量。
- L vq L_{\text{vq}} Lvq:量化损失,用于优化嵌入空间中的代码簇(codebook),减少量化误差。
- L commit L_{\text{commit}} Lcommit:承诺损失,约束编码器的输出靠近选择的嵌入向量。
- β \beta β:权衡参数,用于调节 L commit L_{\text{commit}} Lcommit 的影响。
2. 各部分损失的定义
(1) 重建损失( L rec L_{\text{rec}} Lrec)
- 衡量模型的输入数据 x x x 与解码器生成的重建 x ^ \hat{x} x^ 之间的差异。
- 通常采用均方误差(MSE)作为损失:
L rec = ∥ x − x ^ ∥ 2 L_{\text{rec}} = \|x - \hat{x}\|^2 Lrec=∥x−x^∥2 - 作用:
- 直接优化解码器,使其生成与输入 x x x 尽可能一致的输出。
(2) 量化损失( L vq L_{\text{vq}} Lvq)
-
用于优化嵌入空间中的代码簇(codebook)。
-
量化损失是为了将编码器的连续输出 z e ( x ) z_e(x) ze(x) 静态地绑定到最近邻的嵌入向量 e k e_k ek:
L vq = ∥ sg [ z e ( x ) ] − e k ∥ 2 L_{\text{vq}} = \| \text{sg}[z_e(x)] - e_k \|^2 Lvq=∥sg[ze(x)]−ek∥2 -
其中:
- sg [ ⋅ ] \text{sg}[\cdot] sg[⋅] 表示停止梯度传播,确保该损失只优化嵌入向量 e k e_k ek,不影响编码器。
- z e ( x ) z_e(x) ze(x):编码器的连续输出。
- e k e_k ek:嵌入空间中与 z e ( x ) z_e(x) ze(x) 最近的向量。
-
作用:
- 确保嵌入空间中的向量 e k e_k ek 能够准确表示编码器的输出 z e ( x ) z_e(x) ze(x)。
(3) 承诺损失( L commit L_{\text{commit}} Lcommit)
-
约束编码器的输出 z e ( x ) z_e(x) ze(x) 靠近所选的嵌入向量 e k e_k ek,避免输出偏离。
L commit = ∥ z e ( x ) − sg [ e k ] ∥ 2 L_{\text{commit}} = \| z_e(x) - \text{sg}[e_k] \|^2 Lcommit=∥ze(x)−sg[ek]∥2 -
其中:
- sg [ e k ] \text{sg}[e_k] sg[ek] 表示停止梯度传播,确保嵌入向量 e k e_k ek 不被优化,只作用在编码器上。
-
β \beta β:用于平衡承诺损失的权重。
-
作用:
- 确保编码器的输出 z e ( x ) z_e(x) ze(x) 与选择的嵌入向量保持一致,从而减少量化误差。
3. 损失函数优化目标
VQ-VAE 的目标是最小化上述总损失函数,使得:
- 重建质量高:通过最小化 L rec L_{\text{rec}} Lrec,解码器能够更好地重建原始输入。
- 嵌入向量优化:通过最小化 L vq L_{\text{vq}} Lvq,嵌入向量 e k e_k ek 能够更好地表示数据分布。
- 编码器一致性:通过最小化 L commit L_{\text{commit}} Lcommit,编码器的输出与嵌入向量保持一致。
4. 梯度更新路径
- 解码器:
- 直接通过 L rec L_{\text{rec}} Lrec 的梯度优化解码器的参数。
- 编码器:
- 梯度通过红色线从 L rec L_{\text{rec}} Lrec 传递到编码器的输出 z e ( x ) z_e(x) ze(x),再更新编码器参数。
- 编码器同时受到 L commit L_{\text{commit}} Lcommit 的约束。
- 嵌入空间(Codebook):
- 通过 L vq L_{\text{vq}} Lvq 的梯度更新嵌入向量 e k e_k ek。
VQ-VAE 的训练过程
1. 前向传播
(1) 编码器
- 输入数据:
输入数据(例如图像、语音信号等) x x x 被送入编码器(通常是 CNN)。 - 生成连续潜在表示:
编码器提取输入数据的特征并输出连续的潜在表示 z e ( x ) z_e(x) ze(x),其维度通常为 B × H × W × D B \times H \times W \times D B×H×W×D,其中 B B B 是批次大小, H H H 和 W W W 是特征图的高度和宽度, D D D 是嵌入维度。
z e ( x ) = Encoder ( x ) z_e(x) = \text{Encoder}(x) ze(x)=Encoder(x)
(2) 向量量化
- 最近邻搜索:
通过计算 z e ( x ) z_e(x) ze(x) 与 Codebook 中每个嵌入向量 e k e_k ek 的欧几里得距离,找到最近的嵌入向量 e k e_k ek:
z q ( x ) = arg min e k ∥ z e ( x ) − e k ∥ 2 z_q(x) = \arg\min_{e_k} \| z_e(x) - e_k \|_2 zq(x)=argekmin∥ze(x)−ek∥2 - 量化:
使用最近邻搜索的结果,将 z e ( x ) z_e(x) ze(x) 离散化为嵌入向量 z q ( x ) z_q(x) zq(x),即:
z q ( x ) = e k z_q(x) = e_k zq(x)=ek
(3) 解码器
- 解码重建数据:
将量化后的离散表示 z q ( x ) z_q(x) zq(x) 传递到解码器,生成重建数据 x ^ \hat{x} x^:
x ^ = Decoder ( z q ( x ) ) \hat{x} = \text{Decoder}(z_q(x)) x^=Decoder(zq(x)) - 解码器学习如何将离散化的潜在表示映射回原始数据的空间。
(4) 损失函数计算
在前向传播结束后,计算以下损失:
-
重建损失 L rec L_{\text{rec}} Lrec:
- 衡量解码器重建的
x
^
\hat{x}
x^ 与原始输入
x
x
x 的相似性,通常使用均方误差(MSE):
L rec = ∥ x − x ^ ∥ 2 L_{\text{rec}} = \| x - \hat{x} \|^2 Lrec=∥x−x^∥2 - 目标是优化解码器,使重建数据尽可能接近原始数据。
- 衡量解码器重建的
x
^
\hat{x}
x^ 与原始输入
x
x
x 的相似性,通常使用均方误差(MSE):
-
量化损失 L vq L_{\text{vq}} Lvq:
- 约束嵌入向量
e
k
e_k
ek 更好地表示编码器输出
z
e
(
x
)
z_e(x)
ze(x):
L vq = ∥ sg [ z e ( x ) ] − e k ∥ 2 L_{\text{vq}} = \| \text{sg}[z_e(x)] - e_k \|^2 Lvq=∥sg[ze(x)]−ek∥2 - sg [ ⋅ ] \text{sg}[\cdot] sg[⋅] 表示停止梯度传播,确保只优化嵌入向量,而不影响编码器。
- 约束嵌入向量
e
k
e_k
ek 更好地表示编码器输出
z
e
(
x
)
z_e(x)
ze(x):
-
承诺损失 L commit L_{\text{commit}} Lcommit:
- 约束编码器输出
z
e
(
x
)
z_e(x)
ze(x) 更靠近嵌入向量
e
k
e_k
ek:
L commit = β ∥ z e ( x ) − sg [ e k ] ∥ 2 L_{\text{commit}} = \beta \| z_e(x) - \text{sg}[e_k] \|^2 Lcommit=β∥ze(x)−sg[ek]∥2 - β \beta β 是权衡参数,用于调节承诺损失的影响。
- 约束编码器输出
z
e
(
x
)
z_e(x)
ze(x) 更靠近嵌入向量
e
k
e_k
ek:
-
总损失:
- 将以上三种损失组合为最终的优化目标:
L = L rec + L vq + β L commit L = L_{\text{rec}} + L_{\text{vq}} + \beta L_{\text{commit}} L=Lrec+Lvq+βLcommit
- 将以上三种损失组合为最终的优化目标:
2. 反向传播与更新
在反向传播阶段,根据总损失 L L L 对模型中的各模块参数进行更新:
(1) 解码器更新
- 损失来源:
解码器的参数 ϕ \phi ϕ 直接受到 L rec L_{\text{rec}} Lrec 的影响。 - 目标:
通过反向传播优化解码器的参数,使其能够从 z q ( x ) z_q(x) zq(x) 中生成更准确的重建数据。
(2) 编码器更新
- 损失来源:
编码器的输出 z e ( x ) z_e(x) ze(x) 同时受到 L rec L_{\text{rec}} Lrec 和 L commit L_{\text{commit}} Lcommit 的影响。 - 梯度传递机制:
- L rec L_{\text{rec}} Lrec 的梯度通过 Stop Gradient 操作从 z q ( x ) z_q(x) zq(x) 传递到 z e ( x ) z_e(x) ze(x)。
- L commit L_{\text{commit}} Lcommit 直接作用在 z e ( x ) z_e(x) ze(x) 上,约束其靠近对应的嵌入向量。
- 目标:
优化编码器的参数 θ \theta θ,使其输出的 z e ( x ) z_e(x) ze(x) 更符合离散化的要求,同时更易于解码器使用。
(3) Codebook 更新
- 损失来源:
Codebook 的嵌入向量 e k e_k ek 通过 L vq L_{\text{vq}} Lvq 接受优化信号。 - 更新方式:
嵌入向量 e k e_k ek 会被调整,以更好地靠近编码器输出 z e ( x ) z_e(x) ze(x) 的分布。 - 注意:
由于 Stop Gradient 操作,嵌入向量的更新不会影响编码器输出。
3. 梯度流动路径
在 VQ-VAE 的反向传播中,梯度的流动路径如下:
- 解码器的梯度直接来源于 L rec L_{\text{rec}} Lrec,更新其参数。
- 编码器的梯度来自:
- L rec L_{\text{rec}} Lrec:通过 Stop Gradient 操作从解码器路径传递。
- L commit L_{\text{commit}} Lcommit:直接作用于编码器的输出。
- Codebook 的梯度来自 L vq L_{\text{vq}} Lvq,但编码器的梯度不会反向传递到 Codebook。
4. 训练过程的关键点
-
Stop Gradient 操作:
- 避免不可导的最近邻搜索影响模型的优化。
- 梯度通过解码器路径反向传播,直接作用于编码器的输出。
-
嵌入向量的更新:
- Codebook 是独立更新的,其优化过程完全由 L vq L_{\text{vq}} Lvq 控制。
-
协同优化:
- 编码器、解码器和 Codebook 在每一次反向传播中协同更新,最终实现高质量的重建和离散化表示。
5. 收敛目标
训练完成后,模型达到以下目标:
- 编码器 z e ( x ) z_e(x) ze(x) 的输出能够与 Codebook 中的嵌入向量 e k e_k ek 一一对应。
- 嵌入空间的离散表示 z q ( x ) z_q(x) zq(x) 能够准确表达输入数据。
- 解码器能够从 z q ( x ) z_q(x) zq(x) 生成高质量的重建数据 x ^ \hat{x} x^。
通过这样的训练流程,VQ-VAE 可以在保持生成质量的同时学习到紧凑的离散潜在表示,适用于图像生成、数据压缩等任务。
VQ-VAE 的优点
-
离散化表示学习:
- 将连续的潜在变量转化为离散值,更适合图像、语音等任务的压缩和生成。
-
避免 KL 散度问题:
- 与传统 VAE 中需要优化 KL 散度不同,VQ-VAE 使用代码簇直接进行离散化,训练更加稳定。
-
高质量生成:
- 模型能够生成清晰、高质量的图像或音频。
VQ-VAE 的应用场景
-
图像生成:
- 通过对潜在空间离散化,生成高分辨率的图像,适合 ImageNet 等数据集。
-
音频建模与压缩:
- 应用于语音信号的生成和压缩任务,例如 TTS(Text-to-Speech)。
-
离散特征学习:
- 用于生成对抗网络(GAN)中离散特征表示的学习。
VQ-VAE 的局限性
-
代码簇大小限制:
- 代码簇的数量有限,可能无法捕获复杂数据分布的全部细节。Codebook 大小对模型性能有较大影响,过小可能导致欠拟合,过大则增加计算开销。通常通过超参数搜索调整。
-
高维离散变量的优化:
- 离散化会带来训练的复杂性,尤其是在代码簇更新时可能出现收敛慢的问题。
总结
VQ-VAE 是离散表示学习的一次重要创新,它通过向量量化和离散潜在变量,在图像和语音生成等任务中取得了显著进展。相比传统的 VAE,其生成质量更高,训练过程更加稳定。结合论文内容,可以看出 VQ-VAE 的结构简单但功能强大,是生成模型领域的一个重要里程碑。
补充
VQ-VAE 本身并不是图像生成模型,可以很容易实现图像压缩、重建的目的,但是无法生成新的图像数据。它是一个表征学习模型,专注于将数据映射到离散潜在空间。然而,通过结合 PixelCNN 等生成模型,VQ-VAE 的潜在空间表示可以用于生成任务,从而实现高质量的图像生成。
VAE之所以把图片编码成符合正态分布的连续向量,就是为了能在图像生成时把编码器扔掉,让随机采样出的向量也能通过解码器变成图片。现在倒好,VQ-VAE把图片编码了一个离散向量,这个离散向量构成的空间是不好采样。VQ-VAE不是面临着和AE一样的问题嘛。
VQ-VAE根本不是一个图像生成模型。它和AE一样,只能很好地完成图像压缩,把图像变成一个短得多的向量,而不支持随机图像生成。VQ-VAE和AE的唯一区别,就是VQ-VAE会编码出离散向量,而AE会编码出连续向量。
此处我们不再对 PixelCNN 展开,只需要知道它是一个自回归生成模型,可以逐个像素的生成,因为其是自回归模型,所以每个位置都能看到之前位置的信息,这样生成的 latent code 能够更全面的考虑到空间信息,有助于提高模型生成图像的质量和多样性。
那么为什么不直接使用 PixelCNN 来生成图像呢,答案是肯定的,不过将 PixelCNN 和 VQ-VAE结合可以发挥各自的优势,比如提高训练效率和收敛性,对于 128x128x3 的图像,假设离散空间的大小为 32x32,那么 PixelCNN 不用生成 128x128x3 个像素(RGB),而只用生成 32x32 的离散 latent code 即可。
1. VQ-VAE 将图像编码为离散向量
- 编码器将输入图像 x x x 压缩为离散潜在向量 z q ( x ) z_q(x) zq(x),该潜在表示可以被看作一个“小图像”,例如:
- 输入图像:尺寸 256 × 256 × 3 256 \times 256 \times 3 256×256×3
- 编码后的小图像:尺寸 32 × 32 × D 32 \times 32 \times D 32×32×D(其中 D D D 是嵌入向量的维度)
- 离散潜在向量 z q ( x ) z_q(x) zq(x) 是由 Codebook 的嵌入向量 e k e_k ek 组成的。
2. 利用 PixelCNN 学习生成潜在向量
- PixelCNN 的特性:
- PixelCNN 是一种概率生成模型,能够学习图像像素的条件分布。
- 对于图像的每个像素值(如颜色通道取 0 ∼ 255 0 \sim 255 0∼255 的离散值),PixelCNN 可以输出其概率分布。
- PixelCNN 生成潜在向量:
- PixelCNN 被用于拟合离散潜在向量 z q ( x ) z_q(x) zq(x) 的分布,即学习如何从离散空间中生成潜在表示。
总结生成过程
第一步:训练 VQ-VAE
- 训练编码器和解码器,使得 VQ-VAE 能够将图像压缩为“小图像”(离散潜在向量),并能够从“小图像”还原到完整图像。
第二步:训练 PixelCNN
- 训练 PixelCNN,学习“小图像”的离散分布,即学习如何生成离散潜在向量。
第三步:随机生成图像
- 使用 PixelCNN 从离散潜在空间采样生成“小图像”。
- 将“小图像”输入 VQ-VAE 的解码器,解码生成完整图像。
VQ-VAE-2
论文:Generating Diverse High-Fidelity Images with VQ-VAE-2
VQ-VAE-2 是对 VQ-VAE 的改进模型,通过引入 多层次的潜在空间,显著提升了生成图像的质量和多样性。这篇论文提出了一种层次化的离散潜在变量建模方法,使模型能够捕获数据的全局结构和局部细节,广泛应用于高分辨率图像生成任务。
背景与动机
-
VQ-VAE 的局限性:
- 单层潜在空间的表示能力有限,无法捕捉数据中的复杂多尺度结构。
- 对高分辨率图像的生成表现不足,细节质量和全局一致性都受到限制。
-
多层潜在空间的需求:
- 图像生成需要同时捕捉全局信息(如物体的大致布局)和局部细节(如纹理和颜色)。
- 为此,VQ-VAE-2 引入分层潜在空间,逐步从高层到低层学习全局与局部特征。
VQ-VAE-2 的结构
(a) Overview of the architecture of our hierarchical VQ-VAE:
- 图像输入尺寸为 256 × 256 256 \times 256 256×256,通过编码器压缩为两个量化潜在图:
- 底层潜在图(Bottom Level):大小为 64 × 64 64 \times 64 64×64。
- 顶层潜在图(Top Level):大小为 32 × 32 32 \times 32 32×32。
- 解码器接收这两个潜在图,重建输入图像。
- 编码器和解码器由深度神经网络构成。
(b) Multi-stage image generation:
- 顶层的 PixelCNN 先验分布以类别标签作为条件。
- 底层的 PixelCNN 以类别标签和顶层潜在图为条件生成底层潜在图。
- 使用前馈解码器(feed-forward decoder),从离散潜在图到像素的映射速度很快。
- 生成的示例图像是一只鹦鹉,使用该模型生成。
(1) 模型架构
VQ-VAE-2 的核心思想是将潜在空间分为多个层次,每一层次负责学习不同粒度的特征:
-
顶层潜在空间(High-level Latent Space):
- 表示全局特征,例如图像的主要形状、布局和语义信息。
- 编码器将原始数据压缩到更小的潜在空间,形成高层次的离散表示。
-
底层潜在空间(Low-level Latent Space):
- 表示局部细节,例如纹理、边缘和颜色。
- 条件依赖于顶层潜在表示,对其进行细化。
-
分层关系:
- 编码器和解码器均是分层结构,高层次潜在变量为低层次潜在变量提供条件信息,逐步恢复图像。
(2) 编码器-解码器流程
编码器:
-
输入图像首先通过底层编码器(Bottom-Level Encoder),生成底层潜在表示(latent map),尺寸为 64 × 64 64 \times 64 64×64。
-
底层潜在表示会进一步传递给高层编码器(Top-Level Encoder),生成高层潜在表示,尺寸为 32 × 32 32 \times 32 32×32。
注意:
- 高层编码器的输入并不是直接从图像计算,而是以底层潜在表示为基础。
- 编码器是逐层提取特征的,底层特征用于补充高层的全局信息。
解码器:
-
高层潜在变量(Top-Level Latent Map) 是解码过程的起点,提供全局语义信息。
-
低层解码器(Bottom-Level Decoder) 接收低层潜在变量(Bottom-Level Latent Map)和高层解码输出的条件信息,逐步生成最终的图像。
总结:
- 高层潜在变量主要提供图像的全局结构。
- 底层潜在变量主要补充细节,如纹理、颜色和边缘信息。
多层潜在空间的特点
(1) 分层潜在空间的设计
- 每一层都有独立的 Codebook(代码簇),表示离散潜在空间:
z q ( l ) ( x ) = argmin e k ( l ) ∥ z e ( l ) ( x ) − e k ( l ) ∥ 2 z_q^{(l)}(x) = \text{argmin}_{e_k^{(l)}} \|z_e^{(l)}(x) - e_k^{(l)}\|_2 zq(l)(x)=argminek(l)∥ze(l)(x)−ek(l)∥2- l l l 表示层次的索引,高层为 l = 1 l=1 l=1,低层为 l > 1 l>1 l>1。
- e k ( l ) e_k^{(l)} ek(l) 是第 l l l 层 Codebook 中的嵌入向量。
(2) 潜在空间的上下文建模
- 高层:
- 高层潜在空间学习全局特征,表示数据的抽象信息。
- 低层:
- 条件依赖于高层,捕捉细节特征,弥补高层潜在变量的局部信息不足。
(3) 分层生成过程
- 高层潜在变量通过解码器解码为粗略的图像。
- 低层潜在变量根据高层特征进行局部细化,逐步生成高分辨率图像。
生成过程
(1) 潜在空间的采样
与 VQ-VAE 类似,VQ-VAE-2 的潜在空间难以直接采样。作者采用 PixelCNN(或增强版 PixelSNAIL)对每一层潜在空间进行分布建模,结合类别标签作为条件生成潜在变量:
-
高层潜在变量的建模:
- 高层潜在空间由 PixelCNN(或 PixelSNAIL)用来学习离散潜在变量的条件分布。
- 条件信息包括类别标签(如“鹦鹉”),PixelCNN 基于类别生成全局布局和主要结构。
- 高层潜在变量的生成为后续层次提供全局语义信息。
-
低层潜在变量的条件建模:
- PixelCNN(或 PixelSNAIL)学习低层潜在空间的条件分布,条件包括:
- 类别标签(如“鹦鹉”)。
- 已生成的高层潜在变量(用于捕捉上下文关系)。
- 低层潜在变量的生成补充了局部细节(如纹理和颜色)。
- PixelCNN(或 PixelSNAIL)学习低层潜在空间的条件分布,条件包括:
(2) 图像生成流程
-
顶层潜在空间的采样:
- 使用 PixelSNAIL 结合类别标签生成高层潜在变量。
- 高层潜在变量决定图像的全局结构,例如大致的形状和布局。
-
底层潜在空间的采样:
- 基于高层潜在变量,PixelSNAIL 进一步生成低层潜在变量。
- 低层潜在变量补充局部细节和纹理信息。
-
解码过程:
- 使用前馈解码器(Feed-forward Decoder),从生成的高层和低层潜在变量中逐步生成完整的高分辨率图像。
- 解码过程快速且高效,因为潜在变量已经包含了全局结构和局部细节的信息。
损失函数
VQ-VAE-2 的训练目标与 VQ-VAE 类似,但在多层次潜在空间上引入了条件依赖:
L
=
∑
l
=
1
L
(
L
rec
(
l
)
+
L
vq
(
l
)
+
β
L
commit
(
l
)
)
L = \sum_{l=1}^{L} \Big(L_{\text{rec}}^{(l)} + L_{\text{vq}}^{(l)} + \beta L_{\text{commit}}^{(l)}\Big)
L=l=1∑L(Lrec(l)+Lvq(l)+βLcommit(l))
其中:
- L rec ( l ) L_{\text{rec}}^{(l)} Lrec(l):第 l l l 层的重建损失。
- L vq ( l ) L_{\text{vq}}^{(l)} Lvq(l):第 l l l 层的量化损失。
- L commit ( l ) L_{\text{commit}}^{(l)} Lcommit(l):第 l l l 层的承诺损失,约束编码器输出靠近嵌入向量。
VQ-VAE vs. VQ-VAE-2
特点 | VQ-VAE | VQ-VAE-2 |
---|---|---|
潜在空间 | 单层离散潜在空间 | 多层次分层潜在空间 |
表示能力 | 全局与局部信息混合,难以细化 | 分层结构,分别捕捉全局特征与局部细节 |
生成能力 | 适合低分辨率图像生成 | 适合高分辨率、多样性图像生成 |
生成模块 | PixelCNN | PixelSNAIL |
总结
VQ-VAE-2 是对 VQ-VAE 的显著扩展,通过引入多层次的潜在空间和分层建模方法,实现了更高质量和多样性的生成。它在高分辨率图像生成任务中表现优异,为生成模型的发展提供了重要方向:
- 多层潜在空间: 捕捉全局与局部特征。
- 分层条件建模: 逐步细化生成结果。
- 高分辨率生成: 适用于高质量图像和复杂分布的建模任务。
dVAE
【TODO】
dVAE是VQ-VAE的一个变体,主要在DALL·E中使用,所以,dVAE会结合DALL·E一起讲。参考:
博文(DALL·E和dVAE的很多国内文章图片都来自下面的博文):
Understanding VQ-VAE (DALL-E Explained Pt. 1)
How is it so good ? (DALL-E Explained Pt. 2)
How OpenAI’s DALL-E works?
原文地址:https://blog.csdn.net/qq_42410605/article/details/144227012
免责声明:本站文章内容转载自网络资源,如本站内容侵犯了原著者的合法权益,可联系本站删除。更多内容请关注自学内容网(zxcms.com)!