Parti 水记
《Scaling Autoregressive Models for Content-Rich Text-to-Image Generation》
主页:https://parti.research.google/
论文:https://arxiv.org/abs/2206.10789
这篇也是做 文本生成图片 的工作,比 DALL-E-2 新,而且似乎更强,至少 Parti 能识字了
想看这篇的点在于,他终于抛弃了 diffusion 模型必备的 高斯分布上的自回归,转而使用一种更通用的自回归。
我一直觉得 高斯分布 不是必要的,应该用某种更自由的分布,或者某种训练出来的分布。终于,这样的模型出现了。
方法
与 DALL-E [2], CogView [3], Make-A-Scene [10] 相似, Parti 是一个两阶段模型。
第一阶段,训练一个 tokenizer ,将一张图片变成一串离散的视觉token,用于 训练 和 在推理时重建图像。
第二阶段,训练一个自回归 seq2seq 模型,从 文本token 生成 图像token。
Image Tokenizer
先前的工作 [2, 6, 10, 21] 用 离散VAE 在 图像patch 上学习 量化表征 来解决这个问题。
为了对第二阶段最有用, Image tonkenizer 需要学习出一个高效的视觉codebook,能够在广泛的图像中平衡 codebook 中每个条目的使用率。还必须能够支持从一串 视觉token 中重建出高质量的图片。
我们使用 ViT-VQGAN [21] 。其技术包括 L2-normalized code 和 factorized code ,能够有效提升训练稳定性、重建质量、codebook使用率。
(然后做了很多改动)
为了增强视觉效果,使用了一个简单的超分模型 WDSR [31]。这个模型就是一些卷积层和残差连接。把 256x256 的图像放大到 1024x1024 。
使用与 ViT-VQGAN 相同的损失函数来训练超分模型(感知损失、StyleGAN损失、L2损失,其权重与 ViT-VQGAN也相同)
我们注意到,超分模型也可以选择用diffusion模型,就像 DALL-E 2 [12] 和 Imagen [13] 所做的。可以有 text promot 也可以没有。
用于 文本到图像生成 的 编码器-解码器
如图中所示的,我们在二阶段训练一个标准的 encoder-decoder Transformer 。我们把 文本到图像 当作一个 seq2seq 来建模。
The model takes text as input and is trained using next-token prediction of rasterized image latent codes generated from the first stage image tokenizer.
模型输入文本 token,训练其预测下一个图像 token 。
对于文本 token ,我们建立了一个 sentence-piece model ,在我们自己采样的 image-text 数据集(Section 4.1) 上训练。对于图像 token,我们用 ViT-VQGAN。
在推理阶段,模型会自回归地采样,得到图像 token ,然后再用 ViT-VQGAN 解码器解码出像素。
文本token 数量最多为 128,图像token 数量固定为1024(=32*32)。我们训练了四种不同尺寸的模型,从350M个参数到20B个参数。
文本编码器的预训练
直观上,使用通用语言训练方法训练的编码器应该能更好的处理 prompt 。
我们用了两个数据集做预训练:在 C4 上用 BERT 目标,在我们的 image-text 数据集(Section 4.1)上用对比学习目标(不使用来自对比学习的图像编码器(啥意思?))。
在预训练后,我们继续训练编码器和解码器,使用 softmax 交叉熵损失,使用有8192个图像token的词汇表 。
经过预训练,该模型的表现在GLUE上与BERT相当,但是在text-to-image任务上经过完整的编码器解码器训练后,表现会下降。这个我们我们留做后续研究。尽管如此,在3B个参数的 Parti 模型上,text-encoder 预训练仍然会稍微提升一些 text-to-image 的损失,所以在更大的模型上我们默认使用预训练。
Classifier-Free Guidance and Reranking
Classifier-free guidance 是 diffusion 模型在没有预训练的分类器时提升内容质量的关键技术。
(此处介绍了diffusion 的方法和现有自回归模型的方法)
我们同样对 Parti 使用了 Classifier-free guidance ,我们发现这对文本-图像的对齐效果有显著提升,尤其是在有挑战性的文本prompt上。
对于从一个文本中采样的一批图像,在 DALL-E [2] 中使用 contrastive reranking 产生一个文本-图像对齐分数。我们将其应用到我们的工作上,我们发现,这是对 classifier-free guidance 的一个补充。
相比于 DALL-E 中使用 512 张图像,我们对每个 prompt 只采样16张图像。我们使用 Contrastive Captioners model (CoCa) [25] 来对 文本-图像的对齐程度打分。我们使用(Section 4.1)中表述的相同的数据集训练 CoCa 模型。
规模化
我们使用 Lingvo [36] 框架实现我们的模型,并且用 GSPMD [28] 把模型部署在云服务器上(包括训练和推理)。 GSPMD 是一个能将大型模型实现分布式部署的系统,可以将TPU集群视为单一虚拟设备,并且可以通过在一些 tensor上使用 sharding 注释来将其自动分发到上千个设备中。
……看起来是我用不上的知识呢