大家好,这里是Goodnote(好评笔记)。本文详细介绍CogView模型的设计及其论文解读。CogView 的基本架构与 DALL·E 类似,都使用了 VQ-VAE(Vector Quantized Variational Autoencoder)进行图像编码,并结合 Transformer 进行生成。
论文
CogView: Mastering Text-to-Image Generation via Transformers
CogView 的基本架构与 DALL·E 类似,使用 VQ-VAE(Vector Quantized Variational Autoencoder)将图像编码为离散 token,并结合 Transformer 对文本和图像 token 的联合分布进行建模,生成的 token 序列最终通过 VQ-VAE 解码器重构为图像。
背景介绍(Introduction)
在本节,作者主要介绍了在CV领域中的下一个技术亮点是T2I(Text-To-Image),以及早期的生成式模型(如GAN等,还有VAE)的缺点,之后又介绍了自回归生成模型(包括NLP领域的GPT和CV领域的PixelCNN等),并且提出了像素级自回归模型的大规模图像生成任务中表现出的计算量过大的缺点。
此后,作者提出了使用VQ-VAE框架解决这些问题,并且指出了CogView的架构:结合VQ-VAE框架和Transformer架构(VQ-VAE没有生成能力,需要搭配其他生成模型,这里使用Transformer),也说明了CogView的创新点:精度瓶颈松弛和三明治层归一化(PB-relaxation和Sandwich-LN)。
1.背景和动机:
- T2I(Text-to-Image)技术:文本到图像生成是计算机视觉领域的下一个技术亮点,旨在从图像像素中解耦出形状、颜色、手势等特征,并理解输入文本与图像之间的关系。
- 早期生成式模型:如GAN(生成对抗网络)和VAE(变分自编码器)等,虽然在一些特定任务上取得进展,但在复杂和通用数据集(如MS COCO)上效果有限。
- 自回归生成模型:包括NLP领域的GPT和CV领域的PixelCNN等,它们将图像生成任务转化为像素级的自回归模型,但在大规模图像生成中,计算量过大成为其主要缺点。
2.VQ-VAE框架:
- VQ-VAE:通过将图像压缩到低维离散潜在空间来缓解大规模生成模型计算量过大的问题。VQ-VAE框架使用两阶段训练:
- 第一阶段将图像编码为低维潜在表示;
- 第二阶段使用像PixelCNN这样的自回归模型来拟合这些隐藏变量的先验。
- CogView中的应用:CogView结合VQ-VAE框架和Transformer架构,使用Transformer来拟合图像的潜在变量先验。
3.CogView 的创新:
- 大规模预训练:CogView 使用 3000 万中文文本图像对进行预训练,并采用 40 亿参数的 Transformer 模型,显著增强了文本到图像生成的能力。
- 精度瓶颈松弛和三明治层归一化(PB-relaxation 和 Sandwich-LN):通过优化 LayerNorm 和注意力机制的计算方式,以及在每个残差分支后引入附加 LayerNorm,解决了训练大规模 Transformer 时的数值稳定性问题,从而为生成质量提供了强大的训练保障。
- 中文领域的创新:CogView 是中文领域首个大型文本到图像 Transformer 模型,在中文文本生成任务中表现优异,但其整体性能与 DALL-E 各有所长。
方法
理论(Theory)
背景与推导
整个理论部分
将从 VAE 推导出 CogView 的理论。【这里其实有点歧义,CogView是基于VQ-VAE框架的但是这里从 VAE 推导。其实从后文来看,这里是想说:本部分将从 VAE 推导出VQ-VAE,再进而推导出 CogView 的理论
】。CogView的目标【和DALL·E一样】是优化图像和文本联合似然的证据下界(ELBO),使得它能够有效地处理文本到图像的生成任务。在没有文本信息时,这一推导实际上就变成了VQ-VAE的重新解释。
CogView 采用了 离散潜在空间(而不是传统的连续潜在空间),所以其理论基础实际上是建立在 VQ-VAE 上。VQ-VAE 通过向量量化技术将潜在变量离散化,从而避免了传统 VAE 中的后验塌陷问题(posterior collapse)。
假设数据集 ( X , T ) = { x i , t i } i = 1 N (X, T) = \{x_i, t_i\}_{i=1}^N (X,T)={xi,ti}i=1N由 N N N个图像和文本对组成。图像 x i x_i xi 通过潜在变量 z i z_i zi 生成。具体的生成过程如下:
- 文本 t i t_i ti 从先验 p ( t ; θ ) p(t; \theta) p(t;θ) 生成。
- 潜在变量 z i z_i zi 从条件分布 p ( z ∣ t = t i ; θ ) p(z | t = t_i; \theta) p(z∣t=ti;θ) 生成(潜在变量 z i z_i zi 是 离散)。
- 图像 x i x_i xi 最终从条件分布 p ( x ∣ z = z i ; ψ ) p(x | z = z_i; \psi) p(x∣z=zi;ψ) 生成。
这个过程描述了如何从文本输入生成图像输出。首先,模型理解文本的意义并生成相应的潜在变量(这相当于将文本信息转换为潜在空间中的一个表示),然后通过该潜在变量生成与文本描述一致的图像。这个过程符合文本到图像生成(Text-to-Image Generation) 的范式,在许多现代生成模型(如CogView、DALL·E等)中都有类似的结构。
对数似然和ELBO
因为CogView涉及到文本到图像的双向生成任务,所以其对数似然和ELBO被分解成文本生成部分和图像生成部分。
本小节是关于变分自编码器(VAE)框架中的对数似然(log-likelihood)和证据下界(ELBO)的公式推导。它展示了在使用变分推断时,如何通过最大化 ELBO 来优化模型。下面逐步解释每个部分。
1. 对数似然(Log-Likelihood)
首先,我们有目标函数 log p ( X , T ; θ , ψ ) \log p(X, T; \theta, \psi) logp(X,T;θ,ψ),它表示数据集 ( X , T ) (X, T) (X,T) 的对数似然,其中:
- X = { x 1 , x 2 , … , x N } X = \{ x_1, x_2, \dots, x_N \} X={x1,x2,…,xN} 是图像数据集,
- T = { t 1 , t 2 , … , t N } T = \{ t_1, t_2, \dots, t_N \} T={t1,t2,…,tN} 是与每张图像对应的文本描述。
对数似然的形式可以分解成两部分:
log p ( X , T ; θ , ψ ) = ∑ i = 1 N log p ( t i ; θ ) + ∑ i = 1 N log p ( x i ∣ t i ; θ , ψ ) (1) \log p(X, T; \theta, \psi) = \sum_{i=1}^{N} \log p(t_i; \theta) + \sum_{i=1}^{N} \log p(x_i | t_i; \theta, \psi) \tag{1} logp(X,T;θ,ψ)=i=1∑Nlogp(ti;θ)+i=1∑Nlogp(xi∣ti;θ,ψ)(1)
- 第一项 ∑ i = 1 N log p ( t i ; θ ) \sum_{i=1}^{N} \log p(t_i; \theta) ∑i=1Nlogp(ti;θ) :是生成文本数据 t i t_i ti 的对数似然,模型通过参数 θ \theta θ 来生成文本。这个项通常与文本生成的部分有关,它衡量模型生成正确文本的能力。
- 第二项 ∑ i = 1 N log p ( x i ∣ t i ; θ , ψ ) \sum_{i=1}^{N} \log p(x_i | t_i; \theta, \psi) ∑i=1Nlogp(xi∣ti;θ,ψ) :是生成图像 x i x_i xi 的对数似然,给定相应的文本描述 t i t_i ti,模型通过参数 θ \theta θ 和 ψ \psi ψ 来生成图像 x i x_i xi。这个项衡量了模型从文本描述生成图像的能力。
2. 证据下界(ELBO)
这部分和DALL·E是一样的,详情可以参考:DALL·E模型及其论文详解
由于直接计算对数似然 log p ( x i ∣ t i ; θ , ψ ) \log p(x_i | t_i; \theta, \psi) logp(xi∣ti;θ,ψ) 可能非常复杂(因为涉及后验分布的计算),我们通过变分推断引入一个 变分分布 q ( z ∣ x i ; ϕ ) q(z|x_i; \phi) q(z∣xi;ϕ) 来近似潜在变量 z z z 的真实后验分布 p ( z ∣ x i ) p(z | x_i) p(z∣xi)。变分分布的作用是用来近似计算 p ( x ∣ t ) p(x | t) p(x∣t) 时遇到的困难。
ELBO 是对数似然的下界,它是通过引入变分分布来优化的。ELBO 的推导基于变分推断的 Jensen 不等式,使得对数似然可以转化为一个优化目标。具体的形式是:
log p ( X , T ; θ , ψ ) ≥ − ∑ i = 1 N ( − log p ( t i ; θ ) + E z i ∼ q ( z ∣ x i , ϕ ) [ − log p ( x i ∣ z i ; ψ ) ] + KL ( q ( z ∣ x i ; ϕ ) ∥ p ( z ∣ t i ; θ ) ) ) (2) \log p(X, T; \theta, \psi) \geq - \sum_{i=1}^{N} \left( - \log p(t_i; \theta) + \mathbb{E}_{z_i \sim q(z | x_i, \phi)} \left[ - \log p(x_i | z_i; \psi) \right] + \text{KL}(q(z | x_i; \phi) \parallel p(z | t_i; \theta)) \right) \tag{2} logp(X,T;θ,ψ)≥−i=1∑N(−logp(ti;θ)+Ezi∼q(z∣xi,ϕ)[−logp(xi∣zi;ψ)]+KL(q(z∣xi;ϕ)∥p(z∣ti;θ)))(2)
其中:
- 第一项 − log p ( t i ; θ ) -\log p(t_i; \theta) −logp(ti;θ) 是文本 t i t_i ti 的负对数似然(NLL loss),它是一个与文本生成相关的损失项,用于训练模型中的文本生成部分。
- 第二项 E z i ∼ q ( z ∣ x i , ϕ ) [ − log p ( x i ∣ z i ; ψ ) ] \mathbb{E}_{z_i \sim q(z | x_i, \phi)} \left[ - \log p(x_i | z_i; \psi) \right] Ezi∼q(z∣xi,ϕ)[−logp(xi∣zi;ψ)] 是重建损失(reconstruction loss),它表示在给定潜在变量 z i z_i zi 的情况下,生成图像 x i x_i xi 与真实图像的差异。这个项是变分自编码器的核心部分,用来确保编码器和解码器能够有效地将图像映射到潜在空间,并将其重建回来。
- 第三项 KL ( q ( z ∣ x i ; ϕ ) ∥ p ( z ∣ t i ; θ ) ) \text{KL}(q(z | x_i; \phi) \parallel p(z | t_i; \theta)) KL(q(z∣xi;ϕ)∥p(z∣ti;θ)) 是 KL 散度(Kullback-Leibler divergence),表示变分分布 q ( z ∣ x i ; ϕ ) q(z | x_i; \phi) q(z∣xi;ϕ) 和基于文本 t i t_i ti 的先验分布 p ( z ∣ t i ; θ ) p(z | t_i; \theta) p(z∣ti;θ) 之间的差异。这个项的作用是正则化变分分布,确保它不偏离文本条件的先验分布。KL 散度帮助避免后验崩溃(posterior collapse)问题。
- ELBO(证据下界)中的
第一部分与文本生成相关
,第二和第三部分则与图像生成相关
。
总结
- 对数似然 由两部分组成:文本生成部分和图像生成部分。
- ELBO 是对数似然的下界,包含了文本生成的负对数似然、图像重建损失和变分分布与先验分布之间的 KL 散度。
- 通过优化 ELBO,可以有效地训练模型,使其在生成图像时能够捕捉到文本描述的含义,并且通过 KL 散度避免模型出现不合理的潜在表示。
之前我们在DALL·E模型系列:DALL·E模型及其论文详解
中的第一阶段:Learning the Visual Codebook,介绍DALL·E模型通过证据下界(ELBO)优化视觉,原理和此处基本一致。
VQ-VAE框架
强烈推荐阅读历史文章多模态论文笔记——VQ-VAE和VQ-VAE-2
,不然下面关于VQ-VAE的损失函数和公式含义很难读懂。
若只想大致了解CogView论文【因为VQ-VAE这种结构在后面的CogView系列模型迭代中,会被优化掉,CogView 3 和 DALL·E 2中就舍弃了VQ-VAE架构,转而拥抱扩散模型】,只需要知道VQ-VAE与 VAE 的损失函数的区别:
- VAE 的损失函数优化的是输入数据的重建质量,同时通过 KL 散度确保潜在空间的分布与先验分布一致。这是 VAE 的核心思想,基于连续潜在空间。
- VQ-VAE 的损失函数优化三个目标:重建质量、量化误差(编码器输出与量化后的向量之间的差异),以及承诺损失(编码器输出与离散潜在空间的码本之间的差异)
- KL散度:
- VAE 的 KL 散度项通常是直接衡量近似后验分布 q ( z ∣ x ) q(z | x) q(z∣x) 与先验分布 p ( z ) p(z) p(z) 之间的差异,鼓励编码器学习一个与先验分布接近的潜在空间。
- VQ-VAE 没有直接的 KL 散度项,而是通过 量化损失 和 承诺损失 来引导潜在空间的优化,使得潜在变量 z z z 是离散的并且与码本一致。
1. 传统 VAE 的 KL 散度处理
- 在传统的 VAE 中,潜在变量 z z z 的先验 p ( z ) p(z) p(z) 通常被假设为标准正态分布 N ( 0 , I ) \mathcal{N}(0, I) N(0,I)。
- 编码器 q ( z ∣ x i ) q(z | x_i) q(z∣xi) 学习根据输入数据 x i x_i xi 来估计潜在变量的后验分布。然而,这种方法可能会导致 posterior collapse(后验崩溃),即编码器的输出在多个输入样本上变得相似,导致潜在空间没有得到有效的学习。
2. VQ-VAE 的处理方式
- VQ-VAE 的框架有所不同,主要体现在编码器 ϕ \phi ϕ 固定的方式。VQ-VAE 中,编码器的输出并不直接生成潜在变量的后验分布,而是使用离散化的技术将潜在空间量化。
- 在 VQ-VAE 中,先验分布 p ( z ∣ t i ; θ ) p(z | t_i; \theta) p(z∣ti;θ) 是由另一个模型(由参数 θ \theta θ 控制)拟合的,而不是直接假设为标准正态分布。这避免了后验崩溃问题,因为编码器 ϕ \phi ϕ 不再直接学习潜在空间的后验分布,而是通过量化步骤生成离散的潜在表示。
3. 解决后验崩溃
- 由于编码器只更新以优化 重建损失,不再直接影响潜在空间的后验分布,posterior collapse 问题得以避免。
- 然而,这意味着潜在变量的后验分布 q ( z ∣ x i ) q(z | x_i) q(z∣xi) 在不同输入数据 x i x_i xi 下可能会有很大的变化。因此,需要一个非常强大的 先验模型 p ( z ∣ t i ; θ ) p(z | t_i; \theta) p(z∣ti;θ),来确保能够有效地最小化 KL 散度,避免潜在空间的退化。
4. KL 散度
- VQ-VAE 的关键技术在于使用离散潜在变量和量化技术来优化 KL 散度。KL 散度项被用来衡量潜在变量分布与先验分布之间的差异,并在训练过程中进行最小化。
- 为了消除KL散度,需要一个强大的先验模型。目前,最强大的生成式模型Transforme(GPT),是通过在离散字典上生成tokens序列来使用的。为了使用它,我们使得 z ∈ { 0 , . . . , ∣ V ∣ − 1 } h × w z \in \{0, ..., |V| - 1\}^{h \times w} z∈{0,...,∣V∣−1}h×w, ∣ V ∣ |V| ∣V∣是码本大小, h × w h \times w h×w是潜在变量 z i z_i zi的维度。
- 潜在变量 z i z_i zi可以从变分分布 q ( z ∣ x i ; ϕ ) q(z | x_i; \phi) q(z∣xi;ϕ)中采样,或直接选择最可能的 z i = arg max z q ( z ∣ x i ; ϕ ) z_i = \arg \max_z q(z | x_i; \phi) zi=argmaxzq(z∣xi;ϕ)。
最终ELBO公式可以重写为:
− ∑ i = 1 N ( E z i ∼ q ( z ∣ x i , ϕ ) [ − log p ( x i ∣ z i ; ψ ) ] − log p ( t i ; θ ) − log p ( z i ∣ t i ; θ ) ) - \sum_{i=1}^{N} \left( \mathbb{E}_{z_i \sim q(z | x_i, \phi)} \left[ -\log p(x_i | z_i; \psi) \right] - \log p(t_i; \theta) - \log p(z_i | t_i; \theta) \right) −i=1∑N(Ezi∼q(z∣xi,ϕ)[−logp(xi∣zi;ψ)]−logp(ti;θ)−logp(zi∣ti;θ))
- 重建损失 E z i ∼ q ( z ∣ x i , ϕ ) [ − log p ( x i ∣ z i ; ψ ) ] \mathbb{E}_{z_i \sim q(z | x_i, \phi)} \left[ -\log p(x_i | z_i; \psi) \right] Ezi∼q(z∣xi,ϕ)[−logp(xi∣zi;ψ)]:使得模型能够根据潜在变量准确重建图像。
- 文本生成损失 − log p ( t i ; θ ) -\log p(t_i; \theta) −logp(ti;θ):确保模型根据图像生成合理的文本描述。
- 潜在变量与文本之间的一致性 − log p ( z i ∣ t i ; θ ) -\log p(z_i | t_i; \theta) −logp(zi∣ti;θ):通过最小化潜在变量的负对数似然损失,确保潜在表示 z i z_i zi 与文本描述 t i t_i ti 保持一致。
“在没有文本信息时,这一推导实际上就变成了VQ-VAE的重新解释。”
如果我们去掉文本信息 t i t_i ti,则推导仅涉及图像和其潜在表示 z i z_i zi,此时的目标变为图像的重建任务。此时,公式中的损失项只包含图像的重建损失和潜在变量的量化损失。
- 在这种情况下,最终推导公式变成了:
− ∑ i = 1 N E z i ∼ q ( z ∣ x i , ϕ ) [ − log p ( x i ∣ z i ; ψ ) ] -\sum_{i=1}^{N} \mathbb{E}_{z_i \sim q(z | x_i, \phi)} \left[ -\log p(x_i | z_i; \psi) \right] −i=1∑NEzi∼q(z∣xi,ϕ)[−logp(xi∣zi;ψ)]
这个公式只剩下 重建损失,关注的是通过离散潜在变量 z i z_i zi 重建图像 x i x_i xi。
可能会有疑问:VQ-VAE的量化误差和承诺损失怎么没有在上文看到?
解释:CogView 和 VQ-VAE一样 都使用了 离散潜在空间,但 CogView 采用了 Transformer 模型(如 GPT) 来处理潜在表示和文本之间的关系。这种方式 隐式 解决了传统 VQ-VAE 中的量化问题,因此没有显式地定义量化误差和承诺损失。CogView 的目标是通过最小化重建损失和潜在变量与文本描述的匹配损失来优化模型,而不需要显式的量化和承诺损失。
架构和模型学习过程
架构如下:
这个架构和DALL·E架构相似, DALL·E的架构如下:
学习过程分为两个阶段:
- 第一阶段:
- 使用 VQ-VAE 框架将输入图像 x i x_i xi编码为离散的潜在变量 z i z_i zi(token 序列),并通过解码器重建图像。
- 此阶段的目标是将图像转换为离散的 token 序列,以便后续建模。
- 第二阶段:使用 GPT 模型对输入文本 t i t_i ti 和离散潜在变量 z i z_i zi 建模,学习生成图像 token 的概率分布。
这里其实少了一步,且架构图中也没有明显的标注。那就是第三阶段:生成阶段。
- GPT 生成的离散 token 序列会通过 VQ-VAE 的解码器解码为最终的图像。
- GPT 也可以直接生成文本序列,作为模型的辅助输出或自监督的训练目标。
总结
CogView的创新之处在于将VQ-VAE框架与GPT结合,前者负责图像的离散化和编码,而后者则通过强大的语言建模能力来处理生成任务。通过分阶段的训练,CogView在图像生成的精度和效率上取得了显著提升。
标记化(Tokenization)
在本节中,介绍有关文本和图像 tokenizer 的详细信息。
文本和图像的 Tokenizer
在 CogView 中,文本和图像的 tokenization 方法分别为:
-
文本 Tokenizer:
- CogView 使用了 SentencePiece 作为文本的 tokenizer,利用 BPE(Byte Pair Encoding) 或类似的算法来将文本分割为 50,000 个 tokens。
-
图像 Tokenizer:
- CogView 的图像 tokenizer 是一个 离散自编码器,与 VQ-VAE 或 D-VAE 的第1阶段类似。具体来说,它将图像通过编码器映射到潜在空间后,使用离散的码本(codebook)进行量化。
图像 Tokenizer 的工作原理
图像 x x x 的处理过程如下:
- 输入的图像 x x x 的形状为 H × W × 3 H \times W \times 3 H×W×3,即标准 RGB 图像。
- 通过编码器 Enc ϕ ( x ) \text{Enc}_{\phi}(x) Encϕ(x) 将图像映射到形状为 h × w × d h \times w \times d h×w×d 的潜在空间,其中 h × w h \times w h×w 是编码后的特征图的空间维度, d d d 是潜在变量的维度。
- 然后,编码后的每个 d d d 维向量会被 量化 为可学习字典中的某个邻近嵌入。这个字典的大小为 ∣ V ∣ = 8192 |V| = 8192 ∣V∣=8192,嵌入的维度 d = 256 d = 256 d=256。
- 量化后的结果可以表示为一个 h × w h \times w h×w 的嵌入索引矩阵 z z z,其值来自字典中的索引 { 0 , 1 , . . . , ∣ V ∣ − 1 } \{0, 1, ..., |V|-1\} {0,1,...,∣V∣−1}。
- 解码器 ψ \psi ψ 将量化后的向量映射回一个模糊的图像,以重建输入图像。
在 CogView 的 4B 参数模型中,配置为:
- 字典大小: ∣ V ∣ = 8192 |V| = 8192 ∣V∣=8192
- 嵌入维度: d = 256 d = 256 d=256
- 图像尺寸: H = W = 256 H = W = 256 H=W=256
- 量化后空间维度: h = w = 32 h = w = 32 h=w=32
图像 Tokenizer 的训练方法
在CogView中,图像Tokenizer的训练由于涉及离散选择而具有一定的挑战性。为了解决这个问题,CogView 提出了四种不同的图像Tokenizer训练方法:
- 最近邻映射与直通估计器(原始VQ-VAE方法)。
- Gumbel采样与直通估计器(基于VAE的重参数化技术,使用温度缩放方法,核心思想和 d-VAE 相同)。
- 最近邻映射与移动平均(根据最近映射的向量更新码本嵌入)。
- 最近邻映射与固定码本(初始化后固定码本)。
下面是四种方法的详细介绍:
-
最近邻映射,直通估计器(The nearest-neighbor mapping, straight-through estimator):
- 方法:这是原始VQ-VAE提出的方法。
- 核心思想:通过最近邻映射将每个图像映射到最接近的码本条目,并使用直通估计器来允许梯度通过离散选择过程传递。
- 挑战:一个常见的担忧是,当码本较大且初始化不当时,可能会因为维度灾难而导致仅使用少量的嵌入向量。然而,在CogView的实验中,并没有观察到这种现象。
- 参考:多模态论文笔记——VQ-VAE和VQ-VAE-2
-
Gumbel采样,直通估计器(Gumbel sampling, straight-through estimator):
- 方法:这种方法通过距离计算重参数化潜在变量 z z z 的类别分布,类似于VAE中的方法。
- 核心公式:
p ( z i × w + j = v k ∣ x ) = e − ∥ v k − Enc ϕ ( x ) i j ∥ 2 τ ∑ k = 0 ∣ V ∣ − 1 e − ∥ v k − Enc ϕ ( x ) i j ∥ 2 τ p(z_{i \times w + j} = v_k | x) = \frac{e^{-\frac{\|v_k - \text{Enc}_{\phi}(x)_{ij}\|^2}{\tau}}}{\sum_{k=0}^{|V|-1} e^{-\frac{\|v_k - \text{Enc}_{\phi}(x)_{ij}\|^2}{\tau}}} p(zi×w+j=vk∣x)=∑k=0∣V∣−1e−τ∥vk−Encϕ(x)ij∥2e−τ∥vk−Encϕ(x)ij∥2- 采样策略:使用无偏采样策略选择码本嵌入:
z i × w + j = arg max k ( g k − ∥ v k − Enc ϕ ( x ) i j ∥ 2 τ ) , g k ∼ Gumbel ( 0 , 1 ) z_{i \times w + j} = \arg\max_k \left( g_k - \frac{\|v_k - \text{Enc}_{\phi}(x)_{ij}\|^2}{\tau} \right), \quad g_k \sim \text{Gumbel}(0, 1) zi×w+j=argkmax(gk−τ∥vk−Encϕ(x)ij∥2),gk∼Gumbel(0,1) - 温度 τ \tau τ:温度 τ \tau τ 逐渐减小至0,以使分布更加集中。
- Softmax近似:可以使用softmax函数来近似argmax操作,使过程可微分。
- 采样策略:使用无偏采样策略选择码本嵌入:
- 稳定性:这种方法也被DALL-E采用,并且需要其他技巧来稳定训练。
- 参考:多模态论文笔记——dVAE(DALL·E的核心部件)
-
最近邻映射,移动平均(The nearest-neighbor mapping, moving average):
- 方法:在训练过程中,周期性地将每个嵌入更新为最近映射到该嵌入的向量的均值。
- 核心思想:与使用固定码本不同,这种方法动态更新码本中的嵌入。
-
最近邻映射,固定码本(The nearest-neighbor mapping, fixed codebook):
- 方法:在这种方法中,初始化后码本固定,训练过程中不再更新。
- 核心思想:一旦初始化完成,码本保持不变,整个训练过程不做修改。
实验结果与选择
为了比较这四种训练图像Tokenizer的方法,CogView 团队使用相同的架构、相同的数据集和随机种子来训练四个不同的图像Tokenizer,并展示了图2中的损失曲线。实验结果显示,四种方法的表现基本相当。这意味着,只要初始化得当,码本中嵌入向量的学习并不是特别重要。
在预训练阶段,CogView选择了最近邻映射,移动平均方法(即第三种方法)。
自回归(Auto-regressive)Transformer
CogView 的主干采用单向 Transformer(类似 GPT)。
- Transformer 架构:
- 层数:48层
- 隐藏层大小:2560
- 注意力头:40个
- 总参数量:40亿个参数
序列处理:
- 四个特殊标记被添加到每个输入序列中,以指示文本和图像的边界:
- [ROI1]:图像的参考文本
- [BASE]
- [BOI1]:图像开始
- [EOI1]:图像结束
- 所有序列被裁剪或填充至长度为 1088。
预训练任务(Pretext Task):
- 预训练任务为从左到右的标记预测,也叫语言建模(Language Modeling)。
图像和文本标记的处理:
- 图像和文本标记受到同等的对待。
- DALL-E [39] 推荐降低文本标记的损失权重,但在小规模实验中,CogView 发现文本建模是文本到图像预训练成功的关键。
- 如果文本标记的损失权重设置为零,模型无法找到文本和图像之间的联系,并且会生成与输入文本完全无关的图像。
- 通过文本建模,模型能抽象出隐藏层中的知识,这些知识在后续图像建模中能被有效利用。
稳定训练(Stabilization of training)
在预训练的大型模型(>2B参数)中,通常依赖16位精度来节省GPU存储并加速计算。许多框架,如 DeepSpeed ZeRO,甚至只支持 FP16 参数。然而,text-to-image 任务的预训练在16位精度下非常不稳定。稳定训练是CogView最具挑战性的部分,它与DALL-E保持一致。
1. 数值问题解决方案
- 混合精度框架:由于不同层之间的数值和梯度变化剧烈,DALL-E 提出了一个混合精度框架,per-resblock loss scaling,将所有增益、偏差、嵌入和非嵌入以32位精度存储。这种方案较为复杂,消耗额外的时间和内存,并且当前框架不完全支持。
- CogView的数值正规化:为了应对数值不稳定性,CogView 通过数值正规化来解决。
- 不稳定性:观察到的两种不稳定性包括:
- 溢出(表现为NaN loss)
- 下溢(表现为loss不收敛)
- 不稳定性:观察到的两种不稳定性包括:
2. 技术解决方案
1.Precision Bottleneck Relaxation (PB-Relax):
- 分析模型训练后发现,溢出问题常发生在两个瓶颈操作中:final LayerNorm 和 attention。
- LayerNorm 中的输出可能会爆炸,导致溢出。为了解决这个问题,CogView 提出了通过对输入值进行归一化,即 LayerNorm(x) = LayerNorm(x / max(x)),来松弛瓶颈。
- 注意力分数
Q
T
K
d
\frac{Q^TK}{\sqrt{d}}
dQTK 可能会明显大于输入元素,导致溢出。通过将计算顺序调整为
Q
T
(
K
/
d
)
Q^T (K / \sqrt{d})
QT(K/d) 可以缓解该问题。进一步优化的公式为:
softmax ( Q T K d ) = softmax ( ( Q T α d K ) − max ( Q T α d K ) × α ) \text{softmax}\left(\frac{Q^TK}{\sqrt{d}}\right) = \text{softmax}\left(\left(\frac{Q^T}{\alpha \sqrt{d}} K\right) - \max\left(\frac{Q^T}{\alpha \sqrt{d}} K\right) \times \alpha \right) softmax(dQTK)=softmax((αdQTK)−max(αdQTK)×α)
其中 α \alpha α 是较大的常数(例如 α \alpha α = 32),通过这种方式可以防止溢出。
2.Sandwich LayerNorm (Sandwich-LN):
Transformer中不同LayerNorm结构的图示如下:
- Post-LN 源自原始 Transformer 论文,LayerNorm 位于残差连接之后;
- Pre-LN 是目前最流行的结构,尤其在深度模型(如 GPT)中使用,LayerNorm 位于残差连接之前。因其能提高训练稳定性 ;
- Sandwich-LN 是我们提出的一种创新结构,通过在残差块输入和输出两侧添加 LayerNorm,显著增强了大规模 Transformer 在复杂任务中的训练稳定性。
性能对比:
- Transformer 中的 LayerNorm 对于稳定训练至关重要。Pre-LN 在许多工作中被证明比传统的 Post-LN 更快收敛且更稳定,已成为 Transformer 层的默认结构。
- 然而,在 text-to-image 预训练中,Pre-LN 仍不够稳定。因此,提出了 Sandwich-LayerNorm,即在每个残差分支结束时添加一个新的 LayerNorm,以确保每层的输入值保持在合理范围内,避免了深层模型中常见的 value explosion 问题。
- 实验结果:实验表明,在训练500M模型时,Sandwich-LayerNorm 对收敛的影响几乎可以忽略不计。
微调(Finetuning)
CogView 比 DALL-E 进行了更多的微调,尤其在 text-to-image 生成、超分辨率图像和 self-reranking 方面。
风格学习(Style Learning)
- CogView在四种风格上进行微调:中国画、油画、草图、动画。
- 每种风格用约1000张图像进行微调,文本描述为“An image of {style} style”。
- 生成时,模型根据格式“A {object} of {style} style”生成目标图像。
超分辨率(Super-resolution)
- 将CogView微调为超分辨率模型,处理从16×16图像到32×32图像的转换。
- 使用中央滑动窗策略将32×32图像进一步放大到64×64图像。
- 数据预处理包括裁剪2百万图像并下采样至128×128,tokenization后形成32×32和16×16图像对进行微调。
图像标题生成与自我排序(Image Captioning and Self-reranking)
- 微调CogView进行图像标题生成,通过交换输入序列中的文本和图像字符顺序实现。
- 主要目标是self-reranking,通过标题评分(CAP)评估图像和文本的匹配度。
- 使用 CapS(x, t) 评分公式,根据图像和文本的对应关系调整生成,最终选择具有最高CAPS评分的图像。
工业时尚设计(Industrial Fashion Design)
- 对单一领域生成任务进行优化,简化文本复杂度。
- 使用 VQGAN 替代 VQVAE,降低参数数量并增加序列长度以提高分辨率。
- 在1000万时尚标题数据对上训练3B参数模型,并生成800×800像素的时尚图像,成功应用于 Alibaba Rhino fashion production。