阅读潜在扩散模型时了解到其pretrained autoencoders的架构是VQGAN,今天阅读了一下这篇文章跟大家分享一下,题目是Taming Transformers for High-Resolution Image Synthesis 驯服Taming Transformers 用于高分辨率图像合成。题目气的非常霸气。
了解VQGAN之前,还学习了VQVAE(Vector QuantisedVariational AutoEncoder))这篇论文Neural Discrete Representation Learning,
生成模型 VQVAE:Neural Discrete Representation Learning-CSDN博客
VQVAE的思想来源 vector quantisation (VQ) 向量矢量化
VQVAE的核心思想是:学习一个离散的潜在表征,用于规避VAE framework中出现的posterior collapse 后验坍塌问题
VQVAE和VAE 的不同点在于
1编码器的输出是离散编码的而不是连续编码
2 先验是学习的(一个离散的类别分布)而不是静态的(vae 的先验(p(z)是一个标准的正态分布)
摘要
Transformer 旨在学习序列数据上的远程交互,并继续在各种任务上展示最先进的结果。与 CNN 相比,它们不包含优先考虑局部交互的归纳偏差。这使得它们具有表现力,但对于长序列(例如高分辨率图像)来说在计算上也是不可行的。我们演示了如何将 CNN 归纳偏置的有效性与 Transformer 的表现力相结合,使它们能够建模并从而合成高分辨率图像。我们展示了如何 (i) 使用 CNN 学习图像成分的上下文丰富的词汇,然后 (ii) 利用转换器在高分辨率图像中有效地对其成分进行建模。我们的方法很容易应用于条件合成任务,其中非空间信息(例如对象类)和空间信息(例如分割)都可以控制生成的图像。特别是,我们首次提出了使用 Transformer 进行百万像素图像语义引导合成的结果
Transformers的特点
- 长距离交互:变换器被设计用来学习序列数据中的长距离交互,这在许多任务中取得了最先进的成果。
- 无局部交互偏置:与卷积神经网络(CNNs)不同,变换器不包含优先处理局部交互的归纳偏置。这种特性使得变换器非常灵活和表达力强
- 计算资源要求高:在处理长序列(如高分辨率图像)时,计算上可能不可行。
论文创新 结合思维
CNNs与Transformers的结合
- CNNs的归纳偏置:CNNs的局部连接和权值共享等特性使其非常适合捕捉图像的局部特征和空间层次结构,这种归纳偏置在处理图像时非常有效。
- 结合两者的优势:通过将CNNs的归纳偏置与transformers的表达能力相结合,可以构建一种既能捕捉图像局部特征又能建模长距离依赖的模型。具体来说,CNNs用于学习图像成分的丰富上下文词汇表(即特征表示),而transformers则用于高效地建模这些成分在高分辨率图像中的组合方式。
1. Introduction
我们获得有效且富有表现力的模型的关键见解是,卷积和 Transformer 架构结合在一起可以模拟我们视觉世界的组成性质
组合性建模:作者认为,卷积和变换器架构的结合能够有效地建模视觉世界的组合性质。这意味着,图像可以被视为由多个局部视觉部分(如边缘、纹理、形状等)组成,而这些部分又可以组合成更复杂的结构和场景。
方法概述
-
卷积方法学习视觉部分:首先,使用卷积方法高效地学习一个包含丰富上下文信息的视觉部分(或称为“视觉词汇”)的代码本(codebook)。这一步的目的是捕捉图像的局部结构和特征,这些局部结构是构建更复杂图像的基础。
-
变换器建模全局组合:接下来,利用变换器架构学习这些视觉部分之间的全局组合。由于变换器擅长建模长距离交互,因此它们能够捕捉视觉部分之间复杂的、非局部的关系,从而生成更加逼真和复杂的图像。
-
对抗性方法优化局部结构:同时,采用对抗性方法来确保局部部分的字典(即代码本)能够捕捉到感知上重要的局部结构。这样做可以减少变换器架构在建模低级统计信息方面的负担,使其能够更专注于其独特的优势——建模长距离关系。
优势与成果
-
高分辨率图像生成:通过让变换器专注于其强项(建模长距离关系),该模型能够生成高分辨率的图像,这是以前难以实现的。如图1所示,生成的图像在细节和逼真度上都达到了很高的水平。
-
条件控制:该模型还允许通过条件信息(如期望的对象类别或空间布局)直接控制生成的图像。这使得模型在条件合成任务中非常有用,可以根据用户的需求生成特定类型的图像。
-
性能超越:实验结果表明,该方法保留了变换器的优势,并在基于卷积架构的代码本方法中取得了最新的最佳性能。这表明将卷积和变换器结合起来是一种有效的策略,可以提高图像生成的质量和灵活性。
两阶段方法 Two-Stage Approaches
两阶段方法(Two-Stage Approaches)在图像生成和表示学习中是一个重要的概念,它主要分为两个阶段来逐步构建和优化模型。每个阶段都有其特定的学习目标和任务。
第一阶段:学习数据编码
在第一阶段,主要目标是学习数据的有效表示或编码。这种表示通常旨在捕捉数据中的关键信息,同时去除冗余或噪声。在这一阶段,常用的方法包括变分自动编码器(VAE)、矢量量化变分自编码器(VQVAE)等。这些方法通过训练编码器网络来将输入数据(如图像)映射到一个低维的潜在空间(latent space),同时训练一个解码器网络来从潜在表示中重建原始数据。
- VAE:通过最大化数据的对数似然的下界(即证据下界,ELBO)来训练,同时学习数据的潜在表示和潜在空间中的分布。
- VQVAE:进一步扩展了VAE,通过引入矢量量化层来学习离散的潜在表示。这种离散化表示对于使用自回归模型(如Transformer)进行后续建模特别有用。
第二阶段:学习编码的概率模型
在第二阶段,主要目标是在第一阶段学习到的编码(或潜在表示)的基础上,学习这些编码的概率模型。这通常涉及到对潜在空间中的编码进行建模,以便能够生成新的、未见过的数据表示。
- 使用VAE:在第一阶段使用VAE学习编码后,第二阶段可以再次使用VAE(或不同的VAE变体)来学习这些编码的分布。这有助于捕捉编码之间的统计关系,从而支持生成新的编码。
- 使用归一化流:另一种选择是使用无条件或条件归一化流(NF)来建模编码的分布。归一化流是一种可逆的变换,可以通过改变潜在空间的变量来精确地调整分布。
- 使用GAN:在自动编码器的表示上训练GAN,可以进一步改进生成过程,特别是在生成高分辨率图像时。GAN通过对抗性训练来学习生成器和判别器,生成器负责生成尽可能真实的图像,而判别器则负责区分真实图像和生成图像。
- 使用Transformer:对于自回归建模,如[8]所示,可以使用Transformer模型来学习编码的序列分布。Transformer特别擅长捕捉长距离依赖关系,这对于高分辨率图像中的复杂模式建模非常重要。然而,由于Transformer的计算成本较高,通常需要在较低分辨率的编码上应用它,或者使用VQVAE等方法先对图像进行编码。
提出的方法:
我们的目标是利用 Transformer 模型 [64] 极具前景的学习能力,并将其引入高达百万像素范围的高分辨率图像合成。之前的工作 [48, 8] 将 Transformer 应用于图像生成,对于尺寸高达 64 × 64 像素的图像展示了有希望的结果,但由于序列长度的成本呈二次方增加,无法简单地缩放到更高分辨率。高分辨率图像合成需要一个能够理解图像全局组成的模型,使其能够生成局部真实且全局一致的模式。因此,我们不是用像素来表示图像,而是将其表示为来自密码本的感知丰富的图像成分的组合。通过学习有效的代码,如第 2 节所述。 3.1 中,我们可以显着减少合成的描述长度,这使我们能够使用第 3 节中描述的转换器架构有效地对图像内的全局相互关系进行建模。 3.2.这种方法如图 2 所示,能够在无条件和有条件设置下生成逼真且一致的高分辨率图像。
3.1.学习用于 Transformers 的有效图像成分的码本 Effective Codebook
利用高度表达性的变换器(Transformer)架构进行图像合成的方法,但首先需要通过一个离散代码本(codebook)将图像的组成元素以序列的形式表示出来。这种方法避免了直接在像素级别上操作,而是通过学习到的离散表示来简化问题,提高了效率。以下是详细的理解:
图像表示的离散化
-
离散代码本:为了将图像中的复杂信息简化为可管理的形式,作者提出了使用一个离散代码本
Z
,其中包含了一系列的学习到的表示(或称为“码字”zk
),每个码字都是nz
维的向量。任何图像x ∈ RH×W×3
(其中H
和W
是图像的高度和宽度,3代表RGB三个颜色通道)都可以被表示为一个空间上的代码本条目集合zq ∈ Rh×w×nz
,其中h
和w
可能是经过某种下采样或划分后的尺寸,以减小计算量。 -
序列表示:等效地,这个代码本条目集合可以被视为一个序列,序列中的每个元素都是一个索引,指向代码本中的一个特定码字。这种表示方式使得图像合成问题可以转化为一个序列生成问题,从而可以利用变换器的优势。
学习离散空间代码本
和vae的公式表达不同
vqvae是
而vqgan 量化操作如下,本质都是去匹配一个码本的向量
a subsequent element-wise quantization q(·) of each spatial code ˆ zij ∈ Rnz onto its closest codebook entry zk:
-
卷积模型:为了有效地学习这样的离散空间代码本,作者提出了一个包含编码器
E
和解码器G
的卷积模型。编码器将输入图像x
编码为空间上的代码ˆz
,然后通过量化操作q(·)
将这些代码映射到最近的代码本条目上,得到zq
。解码器则使用这些量化的代码来重建原始图像ˆx
。 -
量化与梯度估计:由于量化操作
q(·)
是不可微的,直接进行反向传播是不可行的。因此,作者采用了“直通梯度估计器”(straight-through gradient estimator),即在反向传播时简单地将解码器的梯度复制到编码器上,从而绕过量化层。这样,整个模型(包括代码本)就可以通过损失函数进行端到端的训练。 -
损失函数:损失函数
LVQ(E, G, Z)
包含了三个部分:重建损失Lrec
(衡量重建图像ˆx
与原始图像x
之间的差异),“承诺损失”(commitment loss,衡量量化后的代码zq
与编码器输出E(x)
之间的差异,带有权重因子β
),以及一个用于正则化编码器输出的项(虽然在这段文字中没有直接提到,但通常在类似的方法中会包含这样的项来防止编码器输出过于发散)。
Learning a Perceptually Rich Codebook学习感知丰富的密码本 (感知不是但看某个像素,而是看一个区域的含义
- VQGAN在VQVAE的基础上进行了改进,特别是通过引入对抗性训练(adversarial training)和感知损失(perceptual loss),以在更高的压缩率下保持图像的感知质量。
-
- 传统方法可能依赖于像素级的L2损失来衡量重建质量,但这往往不足以保持图像的感知质量(如纹理、细节等)。
- VQGAN使用感知损失(如基于特征图的损失,通常从预训练的卷积神经网络中提取)来替代L2损失,这有助于更好地捕捉图像的语义和感知特征。所谓的感知损失,在一般请看下,就是把真实的图像,和解码器复原的图像,一起送给一个神经网络,比如VGG16,把这两张图像经过VGG16,都编码成特征向量,然后计算特征向量的差别
- 同时,引入了一个基于补丁的判别器D,用于区分真实图像和重建图像。这种对抗性训练机制可以进一步提高重建图像的视觉质量,使其更接近真实图像。
- 自适应权重λ:
- 在训练过程中,作者引入了一个自适应权重λ,用于平衡重建损失LGAN(对抗性训练损失)和感知重建损失Lrec。
- λ的计算基于两个损失关于解码器最后一层参数的梯度之间的比率,这有助于在训练过程中动态调整两者之间的权重,以优化整体性能。
- 单注意力层的应用:
- 为了在潜在代码的展开过程中聚合全局上下文信息,作者在最低分辨率特征图上应用了一个单一的注意力层。
- 这有助于变换器模型在处理潜在代码时能够捕捉到更广泛的上下文信息,从而提高其建模能力。
- 训练过程对序列长度的减少:
- 通过上述方法,VQGAN能够在保持图像质量的同时显著降低潜在代码的序列长度。
- 这使得更强大的变换器模型能够被应用于潜在代码的建模,因为它们不再受限于较长的序列长度。
VQGAN和VQVAE中的码表(Codebook)是这两个模型中的一个关键组件,它以一种形象化的方式实现了对图像特征的离散化编码。以下是对码表Codebook的形象化展示和解释:
3.2. Learning the Composition of Images with Transformers3.2.使用 Transformer 学习图像的构成
不仅要压缩,还要学到如何生成图片,对隐含编码进行建模
我们详细探讨了如何使用潜在变压器(基于向量量化的变分自编码器VQGAN)和自回归模型(如Transformer)来生成和条件控制图像。这里有几个关键点值得进一步解释:
1. 图像的量化和编码
- 量化编码:图像
x
首先通过编码器E
被编码成一个量化的表示zq
,其中zq
的每个元素都是码本Z
中的一个索引。这个索引表示了图像在特定位置上的视觉特征。 - 码本索引序列:量化后的表示
zq
可以被转换成一个序列s
,其中每个元素是码本Z
中的一个索引。这个序列s
描述了图像的整体结构和内容。
2. 自回归图像生成
- 自回归模型:使用Transformer这样的自回归模型,可以学习给定前面索引(
s<i
)的情况下,下一个索引(si
)的条件概率分布p(si|s<i)
。通过连乘这些条件概率,我们可以得到整个序列s
的联合概率分布p(s)
。 - 对数似然最大化:通过最大化数据表示的对数似然
log p(s)
,我们可以训练模型生成与训练数据相似的图像。
3. 条件图像合成
- 条件信息:在条件图像合成中,用户可以提供额外的信息
c
(如标签、另一个图像等)来控制生成过程。 - 条件VQGAN:对于具有空间范围的条件信息,可以训练另一个VQGAN来将其编码为索引序列
r
。 - 结合条件信息:在自回归模型中,条件信息
r
可以简单地前置到序列s
中,或者通过修改模型结构(如添加条件嵌入)来结合条件信息。 - 条件概率计算:在条件情况下,计算条件概率
p(si|s<i, r)
,这允许模型根据条件信息生成特定的图像。
4. 高分辨率图像生成
- 注意力机制限制:Transformer的注意力机制对输入序列的长度有限制,这限制了可以处理的图像分辨率。
- 分块处理:为了生成高分辨率图像,可以在训练时裁剪图像,并在采样时使用滑动窗口方法。
- 空间不变性:VQGAN和Transformer的结合依赖于数据的空间不变性或空间调节信息的可用性,以确保生成的图像质量。
5. 实际应用
- 无条件图像合成:在无条件图像合成中,可以通过简单的图像坐标条件化来应对空间不变性假设的违反。
- 多种应用:这种方法可以应用于多种图像生成和编辑任务,如文本到图像的生成、图像超分辨率、图像到图像的翻译等。
形象化展示
想象你有一个巨大的图书馆,这个图书馆里的每一本书都代表了一个特定的图像特征向量(或称为“码字”embedding)。这个图书馆就是我们的码表Codebook。每本书都有一个唯一的索引(或称为“码字索引”),通过这个索引,我们可以快速找到对应的书籍(即特征向量)。
详细解释
- 码表的结构:
- 码表是一个可学习的K×D的张量,其中K是码字(书籍)的数量,D是每个码字(书籍)的维度(可以理解为每本书的内容页数或章节数)。
- 在训练过程中,码表会根据输入图像的特征进行更新,以更好地表示图像中的关键信息。
- 码字的作用:
- 当一张图像通过编码器(Encoder)处理后,会得到一个特征图(feature map),这个特征图上的每个点都对应一个高维向量。
- 这些高维向量会与码表中的码字进行匹配,找到最接近的码字作为该点的离散化表示。这个过程类似于在图书馆中找到与当前查询最相关的书籍。
- 离散化编码的优势:
- 离散化编码有助于减少模型需要处理的信息量,提高计算效率。
- 同时,由于码表中的码字是有限的,这迫使模型学习到更加抽象和泛化的图像表示,从而提高生成模型的泛化能力。
- 训练过程:
- 在训练过程中,模型会尝试优化编码器、解码器和码表三个部分的参数。
- 编码器负责提取图像特征,并将其映射到码表中的码字上。
- 解码器则根据这些离散化的码字重构出原始图像。
- 码表则通过优化其内部的码字来更好地表示图像特征。
- 梯度传递问题:
- 由于离散化操作(即选择最接近的码字)是不可导的,因此在反向传播时无法直接计算梯度。
- 为了解决这个问题,VQGAN和VQVAE采用了一种类似“直通估计器”(straight-through estimator)的方法,在反向传播时直接将解码器梯度复制到编码器上,从而指导编码器的参数更新。
import numpy as np
# 假设 feature_map 是编码器的输出,形状为 (H, W, C),其中 H 是高度,W 是宽度,C 是通道数
# 假设 codebook 是码表,形状为 (K, C),其中 K 是码字的数量,C 是码字的维度(与特征向量的维度相同)
# 假设 feature_map 和 codebook 已经被正确加载和预处理
# 这里我们使用随机数组来模拟它们
H, W, C = 10, 10, 64 # 假设的特征图尺寸和通道数
K = 1024 # 假设的码字数量
np.random.seed(0)
feature_map = np.random.randn(H, W, C).astype(np.float32)
codebook = np.random.randn(K, C).astype(np.float32)
# 将特征图重塑为 (H*W, C) 以方便计算
feature_vectors = feature_map.reshape(-1, C)
# 计算每个特征向量与每个码字之间的欧氏距离
# 使用广播机制来避免显式循环
distances = np.sqrt(((feature_vectors[:, np.newaxis, :] - codebook[np.newaxis, :, :]) ** 2).sum(axis=2))
# 找到每个特征向量最近的码字的索引
nearest_indices = np.argmin(distances, axis=1)
# 如果需要,可以从码表中检索出对应的码字embedding
nearest_embeddings = codebook[nearest_indices]
# 注意:nearest_embeddings 的形状现在是 (H*W, C),与原始特征向量的形状相同(除了批次维度)
# 如果需要,可以将其重新塑形为原始特征图的形状
nearest_embeddings_reshaped = nearest_embeddings.reshape(H, W, C)
# 注意:在实际应用中,通常不需要将 nearest_embeddings_reshaped 存储下来,
# 因为我们更关心的是 nearest_indices,它用于在解码过程中从码表中检索码字
码表中的每个码字都是一个向量,其维度与特征向量的维度相同。在这个上下文中,每个码字的每个维度都有一个具体的数值。例如,如果特征向量有64个维度,那么码表中的每个码字也将有64个维度的值。
distances的输出
在上面的代码中,distances
是一个形状为(H*W, K)
的矩阵,其中H
是特征图的高度,W
是宽度,K
是码字的数量。这个矩阵的每个元素(i, j)
表示特征图上的第i
个特征向量(在将特征图重塑为(H*W, C)
之后)与码表中的第j
个码字之间的欧氏距离。
nearest_indices的输出
nearest_indices
是一个长度为H*W
的一维数组(在NumPy中通常表现为(H*W,)
的形状),它包含了特征图上每个特征向量对应的最近码字的索引。对于特征图上的第i
个特征向量(同样在重塑之后),nearest_indices[i]
给出了码表中与该特征向量最近的码字的索引。
举例
假设我们有一个非常小的特征图(例如1x1x2
,即只有一个特征向量,且该向量有2个维度)和一个包含3个码字的码表(每个码字也有2个维度)。那么:
- 特征向量可能是
[0.5, 1.5]
。 - 码表可能是
[[0, 0], [1, 1], [2, 2]]
(这是为了简化说明而选择的,实际中码字会更加复杂和多样)。 - 计算得到的
distances
可能是[1.58, 0.71, 2.24]
(这些值是通过计算特征向量与每个码字之间的欧氏距离得到的,并进行了四舍五入)。 nearest_indices
将是[1]
,因为码表中的第二个码字(索引为1)与特征向量的距离最短。
在量化过程中,特征图的每个位置(或更准确地说,是特征图上的每个特征向量)都会找到其在码表中对应的“最近”码字向量。这里的“最近”通常是通过计算特征向量与码字向量之间的某种距离度量(如欧氏距离)来确定的,然后选取距离最小的码字作为该特征向量的量化结果。
学习视频