VQ-扩散

VQ-Diffusion

Vector Quantized 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 = "一个泳池里玩的泰迪熊。"

image = pipe(prompt).images[0]

架构

VQ-VAE

使用VQ-VAE编码器将图像编码为一组离散的“令牌”或嵌入向量。为此,图像被分割成补丁,然后每个补丁都被固定大小词汇表中最接近的条目替换。这降低了输入像素空间的维度。VQ-Diffusion使用了Taming Transformers中的VQGAN变种。这篇博文是更好地理解VQ-VAE的资源。

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

前向过程

在前向扩散过程中,每个潜在令牌可以保持不变,被重新采样为不同的潜在向量(每个潜在向量的概率相等),或者被屏蔽。一旦一个潜在令牌被屏蔽,它将保持屏蔽状态。α t \alpha_t α t ​ ,β t \beta_t β t ​ 和 γ t \gamma_t γ t ​ 是控制从步骤 t − 1 t-1 t − 1 到步骤 t t t 的前向扩散过程的超参数。 γ t \gamma_t γ t ​ 是未屏蔽令牌变为屏蔽的概率。α t + β t \alpha_t + \beta_t α t ​ + β t ​ 是未屏蔽令牌保持不变的概率。该令牌可以以 β t \beta_t β t ​ 的概率过渡到任何单个非屏蔽潜在向量。换句话说,α t + K β t + γ t = 1 \alpha_t + K \beta_t + \gamma_t = 1 α t ​ + K β t ​ + γ t ​ = 1 ,其中 K K K 是非屏蔽潜在向量的数量。更多细节请参阅论文的第4.1节。

近似逆过程

编码器-解码器变换器在给定提示 y y y 的条件下,近似了未加噪音的潜在向量类别 x 0 x_0 x 0 ​ 。编码器是一个带有冻结权重的CLIP文本编码器。解码器变换器对所有潜在像素提供非屏蔽的全局注意力,并输出分类分布上向量嵌入的对数概率。解码器变换器在一次前向传递中预测了整个未加噪音潜在向量的分布,对 x t x_t x t ​ 进行全局自注意。将问题框架化为基于离散值的条件序列到序列问题,可以从一定程度上解释为什么编码器-解码器变换器是一个良好的选择。

与AR模型部分提供了关于VQ-Diffusion架构的其他上下文,以及与基于AR变换器的模型的比较。

《驯服Transformer》对将原始像素转换为离散标记在压缩潜在空间中进行了很好的讨论,以使得Transformer在处理图像数据时变得可行。

VQ-Diffusion在上下文中

扩散模型

当代扩散模型大多是连续的。在正向过程中,连续扩散模型迭代地添加高斯噪声。反向过程通过近似计算得到,其中 p θ ( x t − 1 ∣ x t ) = N ( x t − 1 ; μ θ ( x t , t ) , Σ θ ( x t , t ) ) p_{\theta}(x_{t-1} | x_t) = N(x_{t-1}; \mu_{\theta}(x_t, t), \Sigma_{\theta}(x_t, t)) p θ ​ ( x t − 1 ​ ∣ x t ​ ) = N ( x t − 1 ​ ; μ θ ​ ( x t ​ , t ) , Σ θ ​ ( x t ​ , t ) ) 。在DDPM的简化情况下,协方差矩阵是固定的,使用U-Net来训练预测 x t x_t x t ​ 中噪声的模型,并通过噪声得到 x t − 1 x_{t-1} x t − 1 ​ 。

近似的反向过程在结构上与离散反向过程类似。然而,在离散情况下,没有明确的方法来预测 x t x_t x t ​ 中的噪声,因此直接预测 x 0 x_0 x 0 ​ 的分布更为明确。

关于离散扩散模型的文献比连续扩散模型的文献少。《使用非平衡热力学进行深度无监督学习》介绍了一个基于二项分布的扩散模型。《Argmax Flows和Multinomial Diffusion》将离散扩散推广到多项分布,并训练了一个Transformer来预测语言建模任务中未添加噪声的分布。《离散状态空间中的结构去噪扩散模型》将多项式扩散与替代噪声过程(均匀分布、吸收、离散化高斯分布和标记嵌入距离)进行了推广。连续扩散模型也可以使用替代噪声过程,但正如文章中所指出的,只有加性高斯噪声得到了重点关注。

自回归模型

与VQ-Diffusion相比,与AR模型进行比较可能更有意思,因为AR模型更频繁地使用Transformer对离散分布进行预测。尽管Transformer在AR建模中取得了成功,但它们仍然在推理速度下降、图像分辨率增加、误差积累和方向偏差方面存在问题。VQ-Diffusion改善了这三个痛点。

AR图像生成模型的特点是将图像概率进行因式分解,使得每个像素都在光栅扫描顺序(从左到右,从上到下)中依赖于前面的像素,即 p ( x ) = ∏ i p ( x i ∣ x i − 1 , x i − 2 , . . . x 2 , x 1 ) p(x) = \prod_i p(x_i | x_{i-1}, x_{i-2}, … x_{2}, x_{1}) p ( x ) = ∏ i ​ p ( x i ​ ∣ x i − 1 ​ , x i − 2 ​ , . . . x 2 ​ , x 1 ​ ) 。因此,模型可以通过直接最大化对数似然来进行训练。此外,AR模型在操作实际像素(非潜在)值时,从离散多项式分布中预测通道值,即首先从256种方式的softmax中对红色通道值进行采样,然后在红色通道值的条件下预测绿色通道。

AR图像生成模型在体系结构上已经发展了很多工作,以使得Transformer的计算变得可行。在基于Transformer的模型出现之前,PixelRNN、PixelCNN和PixelCNN++是最先进的技术。

《图像Transformer》对非Transformer模型和向Transformer模型的过渡进行了很好的讨论(有关省略的引文,请参见论文)。

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

CNN相对于RNN的一个缺点是它们通常具有相对有限的感受野。这可能会对它们对图像中常见的长程现象(如对称性和遮挡)建模的能力产生不利影响,尤其是在层数较少的情况下。扩大感受野已经被证明可以显著提高质量(Salimans等)。然而,这样做会带来大量参数和计算性能的显著成本,并且可能使训练此类模型更具挑战性。

……自注意机制可以在必要的顺序PixelRNN的几乎无限感受野和更容易并行化的PixelCNN及其各种扩展之间实现更好的平衡。

图像Transformer通过限制自注意力在像素的局部邻域上使用transformer。

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

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

尽管取得了巨大的进步,自回归模型在增加图像分辨率、错误积累和方向偏差方面仍然遭受线性降低的推理速度。对于相同大小的自回归transformer模型,VQ-Diffusion的推理的复杂度比较好,只要扩散步骤的数量小于潜在像素的数量。对于ITHQ数据集,潜在分辨率为32×32,模型训练了100个扩散步骤,实现了大约10倍的复杂度改进。在实践中,VQ-Diffusion”可以比自回归方法快15倍,并获得更好的图像质量”(详见论文以获取更多细节)。此外,VQ-Diffusion不需要教师强制,而是学习纠正错误预测的标记。在训练过程中,噪声图像既被屏蔽又被随机标记替换。VQ-Diffusion还能够在预测x t − 1 x_{t-1} x t − 1 ​时在x t x_t x t 上提供全局上下文。

进一步使用VQ-Diffusion和🧨 Diffusers

到目前为止,我们只移植了在ITHQ数据集上训练的VQ-Diffusion模型。还有在CUB-200、Oxford-102、MSCOCO、Conceptual Captions、LAION-400M和ImageNet上训练的VQ-Diffusion模型。

VQ-Diffusion还支持更快的推理策略。网络重参数化依赖于在未加噪声的图像条件下扩散过程的后验是可计算的。当使用时间步长Δ t \Delta t Δ t 跳过一些逆扩散步骤时,类似的公式适用,p θ ( x t − Δ t ∣ x t , y ) = ∑ x ~ 0 = 1 K q ( x t − Δ t ∣ x t , x ~ 0 ) p θ ( x ~ 0 ∣ x t , y ) 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) p θ ​ ( x t − Δ t ​ ∣ x t ​ , y ) = ∑ x ~ 0 ​ = 1 K ​ q ( x t − Δ t ​ ∣ x t ​ , x ~ 0 ​ ) p θ ​ ( x ~ 0 ​ ∣ x t ​ , y ) 。

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

欢迎贡献!