VQ-Diffusion

矢量量化扩散（VQ-Diffusion）是由中国科学技术大学和微软开发的一种条件潜在扩散模型。与大多数常见研究的扩散模型不同，VQ-Diffusion 的加噪和去噪过程在量化潜在空间中进行，即潜在空间由一组离散向量组成。离散扩散模型比其连续对应模型探索较少，并提供了与自回归（AR）模型有趣的比较点。

• Hugging Face 模型卡片
• Hugging Face Spaces
• 原始实现
• 论文

演示

🧨 Diffusers 允许您仅用几行代码运行 VQ-Diffusion。

安装依赖项

pip install 'diffusers[torch]' transformers ftfy

加载管道

from diffusers import VQDiffusionPipeline

pipe = VQDiffusionPipeline.from_pretrained("microsoft/vq-diffusion-ithq")

如果您想使用 FP16 权重

from diffusers import VQDiffusionPipeline
import torch

pipe = VQDiffusionPipeline.from_pretrained("microsoft/vq-diffusion-ithq", torch_dtype=torch.float16, revision="fp16")

移至 GPU

pipe.to("cuda")

运行管道！

prompt = "A teddy bear playing in the pool."

image = pipe(prompt).images[0]

架构

VQ-VAE

图像使用 VQ-VAE 编码器编码成一组离散的“tokens”或嵌入向量。为此，图像被分割成小块，然后每个小块被替换为代码本中词汇量固定大小的最接近的条目。这减少了输入像素空间的维度。VQ-Diffusion 使用 Taming Transformers 中的 VQGAN 变体。这篇博客文章是更好地理解 VQ-VAE 的一个很好的资源。

VQ-Diffusion 使用了一个预训练的 VQ-VAE，在扩散训练过程中被冻结。

前向过程

在前向扩散过程中，每个潜在 token 可以保持不变，被重新采样到不同的潜在向量（每个都有相等的概率），或者被遮蔽。一旦一个潜在 token 被遮蔽，它将保持遮蔽状态。$\alpha_t$、$\beta_t$ 和 $\gamma_t$ 是控制从步骤 $t-1$ 到步骤 $t$ 的前向扩散过程的超参数。$\gamma_t$ 是未被遮蔽的 token 变为遮蔽的概率。$\alpha_t + \beta_t$ 是未被遮蔽的 token 保持不变的概率。该 token 可以以 $\beta_t$ 的概率转换为任何单个未被遮蔽的潜在向量。换句话说，$\alpha_t + K \beta_t + \gamma_t = 1$，其中 $K$ 是未被遮蔽的潜在向量的数量。更多细节请参见论文第 4.1 节。

近似逆向过程

编码器-解码器 Transformer 近似未加噪的潜在变量 $x_0$ 的类别，并以提示 $y$ 为条件。编码器是一个 CLIP 文本编码器，其权重被冻结。解码器 Transformer 为所有潜在像素提供未遮蔽的全局注意力，并输出向量嵌入的分类分布的对数概率。解码器 Transformer 在一次前向传播中预测未加噪潜在变量的整个分布，为 $x_t$ 提供全局自注意力。将问题框定为离散值上的条件序列到序列问题，为编码器-解码器 Transformer 为何如此合适提供了一些直观的解释。

AR 模型部分提供了 VQ-Diffusion 架构与基于 AR Transformer 模型的比较的额外背景信息。

Taming Transformers 对将原始像素转换为压缩潜在空间中的离散 token 进行了很好的讨论，从而使 Transformer 在图像数据上变得计算可行。

VQ-Diffusion 在上下文中

扩散模型

当代扩散模型大多是连续的。在前向过程中，连续扩散模型迭代地添加高斯噪声。逆向过程通过 $p_{\theta}(x_{t-1} | x_t) = N(x_{t-1}; \mu_{\theta}(x_t, t), \Sigma_{\theta}(x_t, t))$ 来近似。在更简单的 DDPM 情况下，协方差矩阵是固定的，训练一个 U-Net 来预测 $x_t$ 中的噪声，并且 $x_{t-1}$ 由噪声推导。

近似逆向过程在结构上类似于离散逆向过程。然而在离散情况下，预测 $x_t$ 中的噪声没有明确的类比，直接预测 $x_0$ 的分布是一个更明确的目标。

关于离散扩散模型的文献量少于连续扩散模型。使用非平衡热力学的深度无监督学习引入了基于二项分布的扩散模型。Argmax Flows and Multinomial Diffusion 将离散扩散扩展到多项分布，并训练一个 Transformer 来预测用于语言建模任务的无噪声分布。离散状态空间中的结构化去噪扩散模型 通过替代加噪过程（均匀、吸收、离散化高斯和 token 嵌入距离）推广了多项扩散。替代加噪过程在连续扩散模型中也是可能的，但如论文所述，只有加性高斯噪声受到了 significant attention。

自回归模型

将 VQ-Diffusion 与 AR 模型进行比较可能更有趣，因为它们更常以 Transformer 进行离散分布预测为特征。虽然 Transformer 在 AR 建模中取得了成功，但它们仍然面临图像分辨率增加导致推理速度线性下降、错误累积和方向性偏差的问题。VQ-Diffusion 在所有三个痛点上都有所改进。

AR 图像生成模型通过将图像概率因子化来表征，使得每个像素都以前一个像素以光栅扫描顺序（从左到右，从上到下）为条件，即 $p(x) = \prod_i p(x_i | x_{i-1}, x_{i-2}, ... x_{2}, x_{1})$。因此，可以通过直接最大化对数似然来训练模型。此外，作用于实际像素（非潜在）值的 AR 模型，从离散多项分布中预测通道值，即首先从 256 路 softmax 中采样红色通道值，然后绿色通道预测以红色通道值为条件。

AR 图像生成模型在架构上不断发展，许多工作致力于使 Transformer 在计算上可行。在基于 Transformer 的模型之前，PixelRNN、PixelCNN 和 PixelCNN++ 是最先进的模型。

Image Transformer 很好地讨论了非 Transformer 模型以及向基于 Transformer 模型过渡的过程（省略了引文，请参阅论文）。

训练循环神经网络以顺序预测即使是小图像的每个像素在计算上也非常具有挑战性。因此，使用卷积神经网络的可并行化模型，如 PixelCNN，最近受到了更多关注，并且现在在质量上已经超越了 PixelRNN。

与 RNN 相比，CNN 的一个缺点是其感受野通常相当有限。这会对其建模图像中常见的长距离现象（如对称性和遮挡）的能力产生不利影响，尤其是在层数较少的情况下。已证明扩大感受野可以显著提高质量（Salimans 等人）。然而，这样做会显著增加参数数量，从而影响计算性能，并使此类模型的训练更具挑战性。

... 自注意力可以在必然顺序的 PixelRNN 的几乎无限感受野与更具并行性的 PixelCNN 及其各种扩展的有限感受野之间的权衡中实现更好的平衡。

Image Transformer 通过限制局部像素邻域的自注意力来使用 Transformer。

Taming Transformers 和 DALL-E 1 结合了卷积和 Transformer。两者都训练一个 VQ-VAE 来学习离散潜在空间，然后在压缩的潜在空间中训练一个 Transformer。Transformer 的上下文是全局的但被遮蔽，因为注意力提供给所有先前预测的潜在像素，但模型仍然是 AR，因此不能提供注意力给尚未预测的像素。

ImageBART 结合了卷积、Transformer 和扩散过程。它学习了一个离散潜在空间，并通过一个短多项扩散过程进一步压缩。然后训练单独的编码器-解码器 Transformer 来逆转扩散过程中的每个步骤。编码器 Transformer 提供 $x_t$ 上的全局上下文，而解码器 Transformer 自回归地预测 $x_{t-1}$ 中的潜在像素。因此，每个像素都会在噪声更大的图像上获得全局交叉注意力。对于更复杂的数据集，使用了 2-5 个扩散步骤。

尽管取得了巨大的进步，AR 模型仍然面临图像分辨率增加导致推理速度线性下降、错误累积和方向性偏差的问题。对于同等大小的 AR Transformer 模型，只要扩散步骤的数量小于潜在像素的数量，VQ-Diffusion 的推理大 O 复杂度就更优。对于 ITHQ 数据集，潜在分辨率为 32x32，模型训练到 100 个扩散步骤，实现了约 10 倍的大 O 改进。在实践中，VQ-Diffusion“可以比 AR 方法快 15 倍，同时实现更好的图像质量”（更多细节请参阅论文）。此外，VQ-Diffusion 不需要教师强制，而是学习纠正错误预测的 token。在训练期间，加噪图像既被遮蔽，其潜在像素也被随机 token 替换。VQ-Diffusion 还能够提供 $x_t$ 上的全局上下文，同时预测 $x_{t-1}$。

VQ-Diffusion 和 🧨 Diffusers 的后续步骤

目前，我们只移植了在 ITHQ 数据集上训练的 VQ-Diffusion 模型。还有已发布的 VQ-Diffusion 模型在 CUB-200、Oxford-102、MSCOCO、Conceptual Captions、LAION-400M 和 ImageNet 上进行了训练。

VQ-Diffusion 还支持更快的推理策略。网络重参数化依赖于以未加噪图像为条件的扩散过程的后验是可处理的。使用时间步长 $\Delta t$ 时，也会应用类似的公式，该时间步长跳过若干个逆向扩散步骤，$p_\theta (x_{t - \Delta t } | x_t, y) = \sum_{\tilde{x}_0=1}^{K}{q(x_{t - \Delta t} | x_t, \tilde{x}_0)} p_\theta(\tilde{x}_0 | x_t, y)$。

改进的矢量量化扩散模型 通过离散无分类器指导和替代推理策略来解决“联合分布问题”，从而提高了 VQ-Diffusion 的样本质量——更多细节请参见第 3.2 节。离散无分类器指导已合并到 Diffusers 中，但替代推理策略尚未添加。

欢迎贡献！