Parti 水记

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 注释来将其自动分发到上千个设备中。

……看起来是我用不上的知识呢