深度学习笔记——VQ-VAE和VQ-VAE-2

本文详细介绍VQ-VAE和VQ-VAE-2的原理和训练过程，为后面的dVAE在DALLE中的使用打下坚实的基础。

文章目录

• AE和VAE
• VQ-VAE
• • 传统 VAE 的问题
  • VQ-VAE 与 VAE 的对比
  • VQ-VAE 的主要改进
  • VQ-VAE 的核心思想
  • VQ-VAE 模型结构
  • • 1. 编码器 (Encoder)
    • 2. 向量量化层 (Vector Quantization Layer)
    • 3. 解码器 (Decoder)
    • 4. 嵌入空间的动态调整
  • 损失函数
  • • 1. 总损失函数
    • 2. 各部分损失的定义
    • • (1) 重建损失（$L_{\text{rec}}$）
      • (2) 量化损失（$L_{\text{vq}}$）
      • (3) 承诺损失（$L_{\text{commit}}$）
    • 3. 损失函数优化目标
    • 4. 梯度更新路径
  • VQ-VAE 的训练过程
  • • 1. 前向传播
    • • (1) 编码器
      • (2) 向量量化
      • (3) 解码器
      • (4) 损失函数计算
    • 2. 反向传播与更新
    • • (1) 解码器更新
      • (2) 编码器更新
      • (3) Codebook 更新
    • 3. 梯度流动路径
    • 4. 训练过程的关键点
    • 5. 收敛目标
  • VQ-VAE 的优点
  • VQ-VAE 的应用场景
  • VQ-VAE 的局限性
  • 总结
  • 补充
• VQ-VAE-2
• • 背景与动机
  • VQ-VAE-2 的结构
  • • (1) 模型架构
    • (2) 编码器-解码器流程
  • 多层潜在空间的特点
  • • (1) 分层潜在空间的设计
    • (2) 潜在空间的上下文建模
    • (3) 分层生成过程
  • 生成过程
  • • (1) 潜在空间的采样
    • (2) 图像生成流程
  • 损失函数
  • VQ-VAE vs. VQ-VAE-2
  • 总结
• dVAE
• 历史文章
• • 机器学习
  • 深度学习

AE和VAE

参考：深度学习——AE、VAE

VQ-VAE

论文：Neural Discrete Representation Learning

VQ-VAE（Vector Quantized Variational AutoEncoder，向量量化变分自编码器） 主要是将连续潜在空间的点映射到最近的一组离散的向量（即codebook中的向量）。模型引入了离散潜在空间的思想，克服了传统 VAE 中连续潜在空间表示的局限性，能够有效学习高质量的离散特征表示。

传统 VAE 的问题

1. 连续潜在空间的限制：

   • VAE 的潜在变量 $z$ 是连续值，这会导致模型生成的表示较为分散、不够紧凑，无法高效捕获复杂数据的离散结构（如图像中的清晰边缘、重复纹理，或离散的语音特征）。

2. 后验坍塌问题：

   • 潜在变量的表示能力未被充分利用。指编码器生成的潜在表示 $z$ 对解码器的输出贡献非常小，可能部分或完全被忽略。

• 当 KL 散度正则化过强时，编码器可能输出接近于先验分布（如 $\mathcal{N}(0,1)$），导致潜在变量 $z$ 的信息丢失。

VQ-VAE 与 VAE 的对比

特点	VAE	VQ-VAE
潜在空间	连续空间	离散空间
潜在变量 $z$	每一维是连续的实数值，包括所有的有理数（如整数、小数和分数）以及无理数	每一维是离散的整数
潜在分布建模	高斯分布	离散分布（通过 codebook 表示）

VQ-VAE 的主要改进

1. 解决后验坍塌问题：

   • 通过 codebook 的离散化潜在变量，解码器被迫使用潜在变量的所有信息，避免后验坍塌。

2. 增强生成质量：

   • VAE 生成的图像质量相对有限，而 VQ-VAE 可以生成更高分辨率和更清晰的图像。

3. 为后续模型奠定基础：

   • VQ-VAE 的离散表示为 DALL·E 和 VQ-GAN 等模型的开发提供了关键技术支持。

VQ-VAE 的核心思想

关键点：

• VQ-VAE 的最大特点是引入了 向量量化（Vector Quantization），将连续的潜在变量 $z$ 离散化。这有助于模型捕捉和表示更加丰富和复杂的数据分布。
• 将编码器输出的连续潜在向量 $z_e(x)$ 映射到离散的代码簇（codebook）中最近的离散向量 $e_k$，编码范围更加可控。

Codebook 是嵌入空间中一组可学习的向量集合，用来将连续的编码器输出 $z_e(x)$ 映射为离散表示 $z_q(x)$。这些嵌入向量也称为 离散嵌入向量，它们定义了潜在空间中离散值的分布。

高效的重建与生成：

• 编码器将输入映射为潜在空间的连续值。
• 向量量化将这些连续值投影到离散的代码簇中。
• 解码器根据离散化的潜在变量重建输入。

---

VQ-VAE 模型结构

VQ-VAE 与 VAE 的结构非常相似，只是中间部分不是学习概率分布，而是换成 VQ 来学习 Codebook。

VAE架构图

VQ-VAE 架构图

VQ-VAE 的整体结构如下：

Figure 1：

• 左侧：描述 VQ-VAE 的图示。编码器 $z(x)$ 的输出被映射到最近的嵌入点 $e_2$。梯度 ${\nabla }_{z}L$（红色箭头）将推动编码器改变输出，从而在下一次前向传播中调整配置。
• 右侧：嵌入空间的可视化。编码器输出 $z_e(x)$ 被映射到最近的离散点 $e_2$。

VQ-VAE 的核心思想是输入的图片通过编码器、向量量化模块和解码器的协同作用，将连续潜在变量映射为离散的潜在表示，并最终重建输入数据，得到新图像。以下是模型的主要流程：

1. 编码器 (Encoder)

• 输入：左侧图片（例如狗的图像）被输入到 卷积神经网络 (CNN) 编码器中。
• 输出：编码器将输入数据 $x$ 转化为连续的潜在变量 $z_e(x)$【图中绿色】：
  $$z_e(x)=\text{Encoder}(x)$$
  • $z_e(x)$ 是一个连续高维向量，表示输入数据的特征。
  • 从图中可以看到，输出 $z_e(x)$ 是一个 $D$ 维特征张量。

2. 向量量化层 (Vector Quantization Layer)

• 嵌入空间：
  • 嵌入空间由一个 代码簇 (Codebook) 构成【紫色】，包含 $K$ 个离散的嵌入向量 { e 1 , e 2 , … , e k } \{e_1, e_2, \dots, e_k\} {e1​,e2​,…,ek​}，每个向量是可学习的。
  • 嵌入空间用于将编码器的连续表示 $z_e(x)$ 离散化。

Codebook 是嵌入空间中一组可学习的向量集合，用来将连续的编码器输出 $z_e(x)$ 映射为离散表示 $z_q(x)$。这些嵌入向量也称为 离散嵌入向量，它们定义了潜在空间中离散值的分布。
Codebook 的初始化：使用均匀分布初始化嵌入向量，范围为 $[-\frac{1}{\text{num\_embeddings}},\frac{1}{\text{num\_embeddings}}]$，保证稳定性和覆盖性。Codebook 大小对模型性能有较大影响，过小可能导致欠拟合，过大则增加计算开销。VQ-VAE 的论文中，num_embeddings 的典型值在 512 到 4096 之间，具体选择取决于任务和数据集。

• 量化过程：

  1. 对编码器输出 $z_e(x)$ 中的每一个特征向量，找到嵌入空间中与之最近的嵌入向量 $e_k$【紫色】。
     $$z_q(x)=\arg \underset{e_k}{\min }\Vert z_e(x)-e_k{\Vert }_2$$
  2. 得到量化后的潜在表示 $z_q(x)$【紫色】，每个元素对应一个离散的代码【青色】。

• 梯度反向传播：

  • 红色箭头 ${\nabla }_{z}L$ 表示通过量化后的误差对编码器进行反向传播，以调整编码器的输出。
  • 代码簇本身也会更新，以更好地表示数据。

最近邻搜索通过计算嵌入空间 { e 1 , e 2 , … , e k } \{e_1, e_2, \dots, e_k\} {e1​,e2​,…,ek​} 中每个嵌入向量与 $z_e(x)$ 的欧几里得距离，选择最近的嵌入向量 $e_k$：
$$z_q(x)=\arg \underset{e_k}{\min }\Vert z_e(x)-e_k{\Vert }_2$$ 最近邻搜索是一个不可导操作，因为$\arg \min$ 或 $\arg \max$ 是一个离散操作，它选择一个代码簇索引 $k$，这个操作本身不具有梯度，因此编码器的参数 $\theta$， 包括 CNN的所有权重（weights）和偏置（biases），无法直接通过反向传播更新。

为了绕过最近邻搜索的不可导问题，VQ-VAE 使用了 Stop Gradient 操作。具体步骤如下：

1. 停止梯度传播（Stop Gradient）

• 当计算嵌入向量 $z_q(x)$ 时，停止对最近邻搜索的梯度传播。
• 操作：
  • $z_q(x)$ 直接从 codebook 中选取最近的嵌入 $e_k$，但不对 $e_k$ 的选择过程传播梯度。
  • 换句话说，$z_q(x)$ 的选择被认为是固定的。

2. 梯度的复制

• 解码器的损失对 $z_q(x)$ 的梯度 ${\nabla }_{z}L$ 会通过 红色线 被直接复制到编码器的输出 $z_e(x)$ 上。
• 这样，虽然最近邻搜索不可导，但解码器提供的损失信号仍然能够通过 $z_q(x)$ 的路径更新编码器参数。

3. 对嵌入向量的优化

• 嵌入空间中的代码簇 { e 1 , e 2 , … , e k } \{e_1, e_2, \dots, e_k\} {e1​,e2​,…,ek​} 是可学习的。
• 使用以下损失更新嵌入向量： $$L_{\text{vq}}=\Vert \text{sg}[z_e(x)]-e_k{\Vert }^2$$
  • 这里 $\text{sg}[z_e(x)]$ 停止梯度传播，嵌入向量 $e_k$ 会被直接更新。

---

红色线的意义 在图中，红色线的主要作用是通过停止梯度传播的技巧，让解码器的梯度绕过不可导的最近邻搜索，直接作用在编码器的输出 $z_e(x)$ 上，从而实现编码器的参数更新：

1. 梯度传播路径：

   • 从解码器计算的重建损失 $\nabla L_{\text{rec}}$ 沿着 $z_q(x)$ 的路径，跳过最近邻搜索的不可导部分，直接作用在 $z_e(x)$ 上。
   • 通过这种方式，编码器仍然能够接收有效的梯度信号，从而调整其参数。

2. 编码器更新：

   • $z_e(x)$ 被优化为更好地靠近其最近的嵌入向量 $e_k$，减少量化误差。

总结

• ${\nabla }_{z}L$ 是解码器的损失对 $z_q(x)$ 的梯度，但由于使用 Stop Gradient 操作，它被传递到编码器的输出 $z_e(x)$ 上。
• 它是解码器反向传播过程中提供的梯度信号，用于间接优化编码器的输出 $z_e(x)$ 和参数 $\theta$。
• ${\nabla }_{z}L$ 的最终目标是减少重建误差 $L_{\text{rec}}$ 和量化误差 $\Vert z_e(x)-e_k{\Vert }^2$，提升模型的重建能力和表示质量。

---

3. 解码器 (Decoder)

• 输入：量化后的离散潜在表示 $z_q(x)$ 被传递给解码器。

• 解码器使用 CNN 将离散的潜在变量 $z_q(x)$ 转化为重建图像 $\hat{x}$：
  $$\hat{x}=\text{Decoder}(z_q(x))$$

• 输出：最终生成的重建图像 $\hat{x}$ 与输入图像 $x$ 尽可能相似。

---

4. 嵌入空间的动态调整

• 嵌入空间的作用：
  • 从右侧的嵌入空间可视化可以看到，编码器输出 $z_e(x)$ 被映射到离嵌入点最近的 $e_2$。
  • 梯度 ${\nabla }_{z}L$ 将推动编码器的输出朝向最近的嵌入点 $e_2$ 移动，以减少量化误差。

---

损失函数

1. 总损失函数

VQ-VAE 的总损失函数为：
$$L=L_{\text{rec}}+L_{\text{vq}}+\beta L_{\text{commit}}$$
其中：

• $L_{\text{rec}}$：重建损失，优化解码器，衡量输入 $x$ 和重建 $\hat{x}$ 的差异，提升重建质量。
• $L_{\text{vq}}$：量化损失，用于优化嵌入空间中的代码簇（codebook），减少量化误差。
• $L_{\text{commit}}$：承诺损失，约束编码器的输出靠近选择的嵌入向量。
• $\beta$：权衡参数，用于调节 $L_{\text{commit}}$ 的影响。

---

2. 各部分损失的定义

(1) 重建损失（$L_{\text{rec}}$）

• 衡量模型的输入数据 $x$ 与解码器生成的重建 $\hat{x}$ 之间的差异。
• 通常采用均方误差（MSE）作为损失：
  $$L_{\text{rec}}=\Vert x-\hat{x}{\Vert }^2$$
• 作用：
  • 直接优化解码器，使其生成与输入 $x$ 尽可能一致的输出。

---

(2) 量化损失（$L_{\text{vq}}$）

• 用于优化嵌入空间中的代码簇（codebook）。

• 量化损失是为了将编码器的连续输出 $z_e(x)$ 静态地绑定到最近邻的嵌入向量 $e_k$：
  $$L_{\text{vq}}=\Vert \text{sg}[z_e(x)]-e_k{\Vert }^2$$

• 其中：

  • $\text{sg}[\cdot ]$ 表示停止梯度传播，确保该损失只优化嵌入向量 $e_k$，不影响编码器。
  • $z_e(x)$：编码器的连续输出。
  • $e_k$：嵌入空间中与 $z_e(x)$ 最近的向量。

• 作用：

  • 确保嵌入空间中的向量 $e_k$ 能够准确表示编码器的输出 $z_e(x)$。

---

(3) 承诺损失（$L_{\text{commit}}$）

• 约束编码器的输出 $z_e(x)$ 靠近所选的嵌入向量 $e_k$，避免输出偏离。
  $$L_{\text{commit}}=\Vert z_e(x)-\text{sg}[e_k]{\Vert }^2$$

• 其中：

  • $\text{sg}[e_k]$ 表示停止梯度传播，确保嵌入向量 $e_k$ 不被优化，只作用在编码器上。

• $\beta$：用于平衡承诺损失的权重。

• 作用：

  • 确保编码器的输出 $z_e(x)$ 与选择的嵌入向量保持一致，从而减少量化误差。

---

3. 损失函数优化目标

VQ-VAE 的目标是最小化上述总损失函数，使得：

1. 重建质量高：通过最小化 $L_{\text{rec}}$，解码器能够更好地重建原始输入。
2. 嵌入向量优化：通过最小化 $L_{\text{vq}}$，嵌入向量 $e_k$ 能够更好地表示数据分布。
3. 编码器一致性：通过最小化 $L_{\text{commit}}$，编码器的输出与嵌入向量保持一致。

---

4. 梯度更新路径

• 解码器：
  • 直接通过 $L_{\text{rec}}$ 的梯度优化解码器的参数。
• 编码器：
  • 梯度通过红色线从 $L_{\text{rec}}$ 传递到编码器的输出 $z_e(x)$，再更新编码器参数。
  • 编码器同时受到 $L_{\text{commit}}$ 的约束。
• 嵌入空间（Codebook）：
  • 通过 $L_{\text{vq}}$ 的梯度更新嵌入向量 $e_k$。

---

VQ-VAE 的训练过程

1. 前向传播

(1) 编码器

• 输入数据：
  输入数据（例如图像、语音信号等）$x$ 被送入编码器（通常是 CNN）。
• 生成连续潜在表示：
  编码器提取输入数据的特征并输出连续的潜在表示 $z_e(x)$，其维度通常为 $B\times H\times W\times D$，其中 $B$ 是批次大小，$H$ 和 $W$ 是特征图的高度和宽度，$D$ 是嵌入维度。
  $$z_e(x)=\text{Encoder}(x)$$

---

(2) 向量量化

• 最近邻搜索：
  通过计算 $z_e(x)$ 与 Codebook 中每个嵌入向量 $e_k$ 的欧几里得距离，找到最近的嵌入向量 $e_k$：
  $$z_q(x)=\arg \underset{e_k}{\min }\Vert z_e(x)-e_k{\Vert }_2$$
• 量化：
  使用最近邻搜索的结果，将 $z_e(x)$ 离散化为嵌入向量 $z_q(x)$，即：
  $$z_q(x)=e_k$$

---

(3) 解码器

• 解码重建数据：
  将量化后的离散表示 $z_q(x)$ 传递到解码器，生成重建数据 $\hat{x}$：
  $$\hat{x}=\text{Decoder}(z_q(x))$$
• 解码器学习如何将离散化的潜在表示映射回原始数据的空间。

---

(4) 损失函数计算

在前向传播结束后，计算以下损失：

1. 重建损失 $L_{\text{rec}}$：

   • 衡量解码器重建的 $\hat{x}$ 与原始输入 $x$ 的相似性，通常使用均方误差（MSE）：
     $$L_{\text{rec}}=\Vert x-\hat{x}{\Vert }^2$$
   • 目标是优化解码器，使重建数据尽可能接近原始数据。

2. 量化损失 $L_{\text{vq}}$：

   • 约束嵌入向量 $e_k$ 更好地表示编码器输出 $z_e(x)$：
     $$L_{\text{vq}}=\Vert \text{sg}[z_e(x)]-e_k{\Vert }^2$$
   • $\text{sg}[\cdot ]$ 表示停止梯度传播，确保只优化嵌入向量，而不影响编码器。

3. 承诺损失 $L_{\text{commit}}$：

   • 约束编码器输出 $z_e(x)$ 更靠近嵌入向量 $e_k$：
     $$L_{\text{commit}}=\beta \Vert z_e(x)-\text{sg}[e_k]{\Vert }^2$$
   • $\beta$ 是权衡参数，用于调节承诺损失的影响。

4. 总损失：

   • 将以上三种损失组合为最终的优化目标：
     $$L=L_{\text{rec}}+L_{\text{vq}}+\beta L_{\text{commit}}$$

---

2. 反向传播与更新

在反向传播阶段，根据总损失 $L$ 对模型中的各模块参数进行更新：

(1) 解码器更新

• 损失来源：
  解码器的参数 $\varphi$ 直接受到 $L_{\text{rec}}$ 的影响。
• 目标：
  通过反向传播优化解码器的参数，使其能够从 $z_q(x)$ 中生成更准确的重建数据。

---

(2) 编码器更新

• 损失来源：
  编码器的输出 $z_e(x)$ 同时受到 $L_{\text{rec}}$ 和 $L_{\text{commit}}$ 的影响。
• 梯度传递机制：
  1. $L_{\text{rec}}$ 的梯度通过 Stop Gradient 操作从 $z_q(x)$ 传递到 $z_e(x)$。
  2. $L_{\text{commit}}$ 直接作用在 $z_e(x)$ 上，约束其靠近对应的嵌入向量。
• 目标：
  优化编码器的参数 $\theta$，使其输出的 $z_e(x)$ 更符合离散化的要求，同时更易于解码器使用。

---

(3) Codebook 更新

• 损失来源：
  Codebook 的嵌入向量 $e_k$ 通过 $L_{\text{vq}}$ 接受优化信号。
• 更新方式：
  嵌入向量 $e_k$ 会被调整，以更好地靠近编码器输出 $z_e(x)$ 的分布。
• 注意：
  由于 Stop Gradient 操作，嵌入向量的更新不会影响编码器输出。

---

3. 梯度流动路径

在 VQ-VAE 的反向传播中，梯度的流动路径如下：

1. 解码器的梯度直接来源于 $L_{\text{rec}}$，更新其参数。
2. 编码器的梯度来自：
   • $L_{\text{rec}}$：通过 Stop Gradient 操作从解码器路径传递。
   • $L_{\text{commit}}$：直接作用于编码器的输出。
3. Codebook 的梯度来自 $L_{\text{vq}}$，但编码器的梯度不会反向传递到 Codebook。

---

4. 训练过程的关键点

1. Stop Gradient 操作：

   • 避免不可导的最近邻搜索影响模型的优化。
   • 梯度通过解码器路径反向传播，直接作用于编码器的输出。

2. 嵌入向量的更新：

   • Codebook 是独立更新的，其优化过程完全由 $L_{\text{vq}}$ 控制。

3. 协同优化：

   • 编码器、解码器和 Codebook 在每一次反向传播中协同更新，最终实现高质量的重建和离散化表示。

---

5. 收敛目标

训练完成后，模型达到以下目标：

1. 编码器 $z_e(x)$ 的输出能够与 Codebook 中的嵌入向量 $e_k$ 一一对应。
2. 嵌入空间的离散表示 $z_q(x)$ 能够准确表达输入数据。
3. 解码器能够从 $z_q(x)$ 生成高质量的重建数据 $\hat{x}$。

通过这样的训练流程，VQ-VAE 可以在保持生成质量的同时学习到紧凑的离散潜在表示，适用于图像生成、数据压缩等任务。

---

VQ-VAE 的优点

1. 离散化表示学习：

   • 将连续的潜在变量转化为离散值，更适合图像、语音等任务的压缩和生成。

2. 避免 KL 散度问题：

   • 与传统 VAE 中需要优化 KL 散度不同，VQ-VAE 使用代码簇直接进行离散化，训练更加稳定。

3. 高质量生成：

   • 模型能够生成清晰、高质量的图像或音频。

---

VQ-VAE 的应用场景

1. 图像生成：

   • 通过对潜在空间离散化，生成高分辨率的图像，适合 ImageNet 等数据集。

2. 音频建模与压缩：

   • 应用于语音信号的生成和压缩任务，例如 TTS（Text-to-Speech）。

3. 离散特征学习：

   • 用于生成对抗网络（GAN）中离散特征表示的学习。

---

VQ-VAE 的局限性

1. 代码簇大小限制：

   • 代码簇的数量有限，可能无法捕获复杂数据分布的全部细节。Codebook 大小对模型性能有较大影响，过小可能导致欠拟合，过大则增加计算开销。通常通过超参数搜索调整。

2. 高维离散变量的优化：

   • 离散化会带来训练的复杂性，尤其是在代码簇更新时可能出现收敛慢的问题。

---

总结

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$ 压缩为离散潜在向量 $z_q(x)$，该潜在表示可以被看作一个“小图像”，例如：
  • 输入图像：尺寸 $256\times 256\times 3$
  • 编码后的小图像：尺寸 $32\times 32\times D$（其中 $D$ 是嵌入向量的维度）
• 离散潜在向量 $z_q(x)$ 是由 Codebook 的嵌入向量 $e_k$ 组成的。

2. 利用 PixelCNN 学习生成潜在向量

• PixelCNN 的特性：
  • PixelCNN 是一种概率生成模型，能够学习图像像素的条件分布。
  • 对于图像的每个像素值（如颜色通道取 $0\sim 255$ 的离散值），PixelCNN 可以输出其概率分布。
• PixelCNN 生成潜在向量：
  • PixelCNN 被用于拟合离散潜在向量 $z_q(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 的改进模型，通过引入 多层次的潜在空间，显著提升了生成图像的质量和多样性。这篇论文提出了一种层次化的离散潜在变量建模方法，使模型能够捕获数据的全局结构和局部细节，广泛应用于高分辨率图像生成任务。

---

背景与动机

1. VQ-VAE 的局限性：

   • 单层潜在空间的表示能力有限，无法捕捉数据中的复杂多尺度结构。
   • 对高分辨率图像的生成表现不足，细节质量和全局一致性都受到限制。

2. 多层潜在空间的需求：

   • 图像生成需要同时捕捉全局信息（如物体的大致布局）和局部细节（如纹理和颜色）。
   • 为此，VQ-VAE-2 引入分层潜在空间，逐步从高层到低层学习全局与局部特征。

---

VQ-VAE-2 的结构

(a) Overview of the architecture of our hierarchical VQ-VAE:

• 图像输入尺寸为 $256\times 256$，通过编码器压缩为两个量化潜在图：
  • 底层潜在图（Bottom Level）：大小为 $64\times 64$。
  • 顶层潜在图（Top Level）：大小为 $32\times 32$。
• 解码器接收这两个潜在图，重建输入图像。
• 编码器和解码器由深度神经网络构成。

(b) Multi-stage image generation:

• 顶层的 PixelCNN 先验分布以类别标签作为条件。
• 底层的 PixelCNN 以类别标签和顶层潜在图为条件生成底层潜在图。
• 使用前馈解码器（feed-forward decoder），从离散潜在图到像素的映射速度很快。
• 生成的示例图像是一只鹦鹉，使用该模型生成。

(1) 模型架构

VQ-VAE-2 的核心思想是将潜在空间分为多个层次，每一层次负责学习不同粒度的特征：

1. 顶层潜在空间（High-level Latent Space）：

   • 表示全局特征，例如图像的主要形状、布局和语义信息。
   • 编码器将原始数据压缩到更小的潜在空间，形成高层次的离散表示。

2. 底层潜在空间（Low-level Latent Space）：

   • 表示局部细节，例如纹理、边缘和颜色。
   • 条件依赖于顶层潜在表示，对其进行细化。

3. 分层关系：

   • 编码器和解码器均是分层结构，高层次潜在变量为低层次潜在变量提供条件信息，逐步恢复图像。

(2) 编码器-解码器流程

编码器：

1. 输入图像首先通过底层编码器（Bottom-Level Encoder），生成底层潜在表示（latent map），尺寸为 $64\times 64$。

2. 底层潜在表示会进一步传递给高层编码器（Top-Level Encoder），生成高层潜在表示，尺寸为 $32\times 32$。

   注意：

   • 高层编码器的输入并不是直接从图像计算，而是以底层潜在表示为基础。
   • 编码器是逐层提取特征的，底层特征用于补充高层的全局信息。

解码器：

1. 高层潜在变量（Top-Level Latent Map） 是解码过程的起点，提供全局语义信息。

2. 低层解码器（Bottom-Level Decoder） 接收低层潜在变量（Bottom-Level Latent Map）和高层解码输出的条件信息，逐步生成最终的图像。

   总结：

   • 高层潜在变量主要提供图像的全局结构。
   • 底层潜在变量主要补充细节，如纹理、颜色和边缘信息。

多层潜在空间的特点

(1) 分层潜在空间的设计

• 每一层都有独立的 Codebook（代码簇），表示离散潜在空间：
  $$z_q^{(l)}(x)={\text{argmin}}_{e_k^{(l)}}\Vert z_e^{(l)}(x)-e_k^{(l)}{\Vert }_2$$
  • $l$ 表示层次的索引，高层为 $l=1$，低层为 $l>1$。
  • $e_k^{(l)}$ 是第 $l$ 层 Codebook 中的嵌入向量。

(2) 潜在空间的上下文建模

• 高层：
  • 高层潜在空间学习全局特征，表示数据的抽象信息。
• 低层：
  • 条件依赖于高层，捕捉细节特征，弥补高层潜在变量的局部信息不足。

(3) 分层生成过程

• 高层潜在变量通过解码器解码为粗略的图像。
• 低层潜在变量根据高层特征进行局部细化，逐步生成高分辨率图像。

---

生成过程

(1) 潜在空间的采样

与 VQ-VAE 类似，VQ-VAE-2 的潜在空间难以直接采样。作者采用 PixelCNN（或增强版 PixelSNAIL）对每一层潜在空间进行分布建模，结合类别标签作为条件生成潜在变量：

1. 高层潜在变量的建模：

   • 高层潜在空间由 PixelCNN（或 PixelSNAIL）用来学习离散潜在变量的条件分布。
   • 条件信息包括类别标签（如“鹦鹉”），PixelCNN 基于类别生成全局布局和主要结构。
   • 高层潜在变量的生成为后续层次提供全局语义信息。

2. 低层潜在变量的条件建模：

   • PixelCNN（或 PixelSNAIL）学习低层潜在空间的条件分布，条件包括：
     • 类别标签（如“鹦鹉”）。
     • 已生成的高层潜在变量（用于捕捉上下文关系）。
   • 低层潜在变量的生成补充了局部细节（如纹理和颜色）。

(2) 图像生成流程

1. 顶层潜在空间的采样：

   • 使用 PixelSNAIL 结合类别标签生成高层潜在变量。
   • 高层潜在变量决定图像的全局结构，例如大致的形状和布局。

2. 底层潜在空间的采样：

   • 基于高层潜在变量，PixelSNAIL 进一步生成低层潜在变量。
   • 低层潜在变量补充局部细节和纹理信息。

3. 解码过程：

   • 使用前馈解码器（Feed-forward Decoder），从生成的高层和低层潜在变量中逐步生成完整的高分辨率图像。
   • 解码过程快速且高效，因为潜在变量已经包含了全局结构和局部细节的信息。

---

损失函数

VQ-VAE-2 的训练目标与 VQ-VAE 类似，但在多层次潜在空间上引入了条件依赖：
$$L=\sum _{l=1}^L(L_{\text{rec}}^{(l)}+L_{\text{vq}}^{(l)}+\beta L_{\text{commit}}^{(l)})$$
其中：

• $L_{\text{rec}}^{(l)}$：第 $l$ 层的重建损失。
• $L_{\text{vq}}^{(l)}$：第 $l$ 层的量化损失。
• $L_{\text{commit}}^{(l)}$：第 $l$ 层的承诺损失，约束编码器输出靠近嵌入向量。

---

VQ-VAE vs. VQ-VAE-2

特点	VQ-VAE	VQ-VAE-2
潜在空间	单层离散潜在空间	多层次分层潜在空间
表示能力	全局与局部信息混合，难以细化	分层结构，分别捕捉全局特征与局部细节
生成能力	适合低分辨率图像生成	适合高分辨率、多样性图像生成
生成模块	PixelCNN	PixelSNAIL

---

总结

VQ-VAE-2 是对 VQ-VAE 的显著扩展，通过引入多层次的潜在空间和分层建模方法，实现了更高质量和多样性的生成。它在高分辨率图像生成任务中表现优异，为生成模型的发展提供了重要方向：

1. 多层潜在空间： 捕捉全局与局部特征。
2. 分层条件建模： 逐步细化生成结果。
3. 高分辨率生成： 适用于高质量图像和复杂分布的建模任务。

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?

历史文章

机器学习

机器学习笔记——损失函数、代价函数和KL散度
机器学习笔记——特征工程、正则化、强化学习
机器学习笔记——30种常见机器学习算法简要汇总
机器学习笔记——感知机、多层感知机(MLP)、支持向量机(SVM)
机器学习笔记——KNN（K-Nearest Neighbors，K 近邻算法）
机器学习笔记——朴素贝叶斯算法
机器学习笔记——决策树
机器学习笔记——集成学习、Bagging（随机森林）、Boosting（AdaBoost、GBDT、XGBoost、LightGBM）、Stacking
机器学习笔记——Boosting中常用算法（GBDT、XGBoost、LightGBM）迭代路径
机器学习笔记——聚类算法（Kmeans、GMM-使用EM优化）
机器学习笔记——降维

深度学习

深度学习笔记——优化算法、激活函数
深度学习——归一化、正则化
深度学习——权重初始化、评估指标、梯度消失和梯度爆炸
深度学习笔记——前向传播与反向传播、神经网络（前馈神经网络与反馈神经网络）、常见算法概要汇总
深度学习笔记——卷积神经网络CNN
深度学习笔记——循环神经网络RNN、LSTM、GRU、Bi-RNN
深度学习笔记——Transformer
深度学习笔记——3种常见的Transformer位置编码
深度学习笔记——GPT、BERT、T5
深度学习笔记——ViT、ViLT
深度学习笔记——DiT（Diffusion Transformer）
深度学习笔记——CLIP、BLIP
深度学习笔记——AE、VAE
深度学习笔记——生成对抗网络GAN
深度学习笔记——模型训练工具（DeepSpeed、Accelerate）
深度学习笔记——模型压缩和优化技术（蒸馏、剪枝、量化）
深度学习笔记——模型部署