VQ‑Diffusion：离散空间中的条件潜在扩散

VQ-Diffusion 是由中国科学技术大学和微软联合开发的条件潜在扩散模型。不同于大多数在连续空间中工作的扩散模型，VQ-Diffusion 在由离散向量集合构成的量化潜在空间上执行加噪和去噪过程。

架构与技术实现

VQ-Diffusion 结合了用于降维的 VQ‑VAE 与用于扩散过程的编码器‑解码器 Transformer。

VQ‑VAE 潜在空间

图像首先通过 VQ‑VAE 编码器被编码为离散 token 或嵌入向量。该过程包括将图像划分为若干块，并用固定大小词表码本中最接近的条目替换每个块。VQ‑Diffusion 特别采用了 Taming Transformers 中的 VQGAN 变体，使用预训练且在扩散训练期间保持冻结的 VQ‑VAE。

前向扩散过程

在前向过程里，潜在 token 从干净状态转变为加噪状态。每个 token 可以：

保持不变。
以相等概率重新采样为另一个潜在向量。
被掩码。

一旦 token 被掩码，它将保持掩码状态。该过程受超参数 $\alpha_{t}$、$\beta_{t}$、$\gamma_{t}$ 控制，其中 $\gamma_{t}$ 是 token 变为掩码的概率，$\alpha_{t}+\beta_{t}$ 是它保持不变的概率。转移到任意非掩码潜在向量的概率为 $\beta_{t}$。

近似逆过程

为了去噪并生成图像，使用编码器‑解码器 Transformer 来近似在给定提示 ($y$) 条件下的未加噪潜在 ($x_{0}$) 类别分布。

Encoder（编码器）： 使用冻结权重的 CLIP 文本编码器。
Decoder（解码器）： 一个 Transformer，向所有潜在像素提供未掩码的全局注意力，并输出向量嵌入的分类分布的对数概率。

由于解码器在一次前向传播中预测了未加噪潜在的完整分布，它能够对当前状态 $x_{t}$ 实现全局自注意力。

VQ‑Diffusion 与其他生成模型的比较

VQ‑Diffusion 为连续扩散模型和自回归（AR）模型提供了一种技术上的替代方案。

与连续扩散的比较

大多数当代扩散模型是连续的，迭代地添加高斯噪声，并通过 U‑Net 预测噪声以近似逆过程。相比之下，VQ‑Diffusion 在离散值上操作，直接预测 $x_{0}$ 的分布比预测噪声更为明确。

与自回归（AR）模型的比较

VQ‑Diffusion 解决了基于 Transformer 的图像 AR 模型的三个主要痛点：随分辨率提升而线性下降的推理速度、误差累积以及方向性偏差。

推理速度： AR 模型将图像概率因式分解，使每个像素在光栅扫描顺序上依赖于前面的像素。只要扩散步数少于潜在像素数量，VQ‑Diffusion 的推理复杂度就更优。以 ITHQ 数据集（32×32 潜在分辨率）为例，VQ‑Diffusion 训练至 100 步，约实现 10 倍的大 O 改进。实际使用中，其推理速度可比 AR 方法快至 15 倍，同时获得更好的图像质量。
错误校正： 与需要教师强制（teacher‑forcing）的 AR 模型不同，VQ‑Diffusion 通过在训练中同时进行掩码和随机 token 替换，学习纠正错误预测的 token。
全局上下文： VQ‑Diffusion 在预测 $x_{t-1}$ 时能够利用 $x_{t}$ 的全局上下文，而 AR 模型受限于掩码注意力，无法看到未来像素。

实现与进一步优化

VQ‑Diffusion 已通过 VQDiffusionPipeline 集成到 Hugging Face diffusers 库中。模型已在包括 ITHQ、CUB‑200、Oxford‑102、MSCOCO、Conceptual Captions、LAION‑400M 与 ImageNet 在内的多种数据集上进行训练。

进一步的优化包括使用时间步长 $\Delta t$ 跳过逆扩散步骤的更快推理策略。此外，Improved Vector Quantized Diffusion Models 引入了离散的 classifier‑free guidance 以提升采样质量，并提供了另一种推理策略来解决联合分布问题。