谷歌开始卷自己，AI架构Pathways加持，推出200亿生成模型

机器之心

发布于 2022-06-27 12:42:36

4640

发布于 2022-06-27 12:42:36

文章被收录于专栏：机器之心

机器之心报道

机器之心编辑部

继 Imagen 后，谷歌又推出了文本 - 图像生成模型 Parti。

你见过一只小狗破壳而出吗？或者用飞艇俯瞰蒸汽朋克中的城市？又或者两个机器人在电影院像人类一样看电影…… 这些听起来可能有些天马行空，但一种名为「文本到图像生成」的新型机器学习技术使这些成为可能。

谷歌研究院的科学家和工程师一直致力于探索使用各种 AI 技术生成文本到图像的方法。

今年 5 月底，谷歌推出 AI 创作神器 Imagen，它结合了 Transformer 语言模型和高保真扩散模型的强大功能，在文本到图像的合成中提供前所未有的逼真度和语言理解能力。与仅使用图像 - 文本数据进行模型训练的先前工作相比，Imagen 的关键突破在于：谷歌的研究者发现在纯文本语料库上预训练的大型 LM 的文本嵌入对文本到图像的合成显著有效。Imagen 的文本到图像生成可谓天马行空，能生成多种奇幻却逼真的有趣图像。

Imagen 生成效果是这样的，比如正在户外享受骑行的柴犬（下图左）以及狗狗照镜子发现自己是只猫（下图右）：

时隔没多久，谷歌又推出了 Parti（Pathways Autoregressive Text-to-Image），该模型最高可扩展至 200 亿参数，并且随着可使用参数数量的增长，其输出的图像也能够更加逼真。

值得一提的是，这是谷歌大牛 Jeff Dean 提出的多任务 AI 大模型蓝图 Pathways 的一部分。

我们先来看下 Parti 效果，袋熊在瀑布旁，背着书包，拄着拐杖眺望着远方：

埃及阿努比斯肖像，在洛杉矶背景下，戴着飞行员护目镜，穿着白色 t 恤和黑色皮夹克：

一只熊猫戴着一顶巫师帽骑在马上：

下面我们介绍一下 Parti 的实现原理。

Parti 模型

与 DALL-E、CogView 和 Make-A-Scene 类似，Parti 是一个两阶段模型，由图像 tokenizer 和自回归模型组成，如下图 3 所示。第一阶段训练一个 tokenizer，该 tokenizer 可以将图像转换为一系列离散的视觉 token，用于训练并在推理时重建图像。第二阶段训练从文本 token 生成图像 token 的自回归序列到序列模型。

图像 Tokenizer

首先，该研究训练了一个 ViT-VQGAN-Small 模型（8 个块，8 个头，模型维度 512，隐藏维度 2048，总参数约为 30M），并且学习了 8192 张图像 token 类别用于代码本。

为了进一步提高第二阶段编码器 - 解码器训练后重建图像的视觉灵敏度，该研究冻结了 tokenizer 的编码器和代码本，并微调更大尺寸的 tokenizer 解码器（32 个块，16 个头，模型维度 1280，隐藏维度 5120，总参数约 600M）。图像 tokenizer 的输入和输出使用 256×256 分辨率。

最后，虽然分辨率为 256×256 的图像捕获了大部分内容、结构和纹理，但更高分辨率的图像具有更大的视觉冲击力。为此，该研究在图像 tokenizer 上采用了一个简单的超分辨率模块，如下图 4 所示。

文本到图像生成的编码器 - 解码器架构

如上图 3 所示，该研究第二阶段训练了一种标准的编码器 - 解码器 Transformer 模型，将文本到图像视为序列到序列建模问题。该模型将文本作为输入，并使用从第一阶段图像 tokenizer 生成的光栅化图像潜在代码的下一个 token 预测进行训练。对于文本编码，该研究构建了一个 sentence-piece 模型，词汇量为 16000。在推理时，模型对图像 token 进行自回归采样，随后使用 ViT-VQGAN 解码器将其解码为像素。

该研究使用的文本 token 最大为 128，图像 token 的长度固定为 1024。所有模型都使用 conv-shaped 掩码稀疏注意力。该研究训练了四种变体，参数量从 3.5 亿到 200 亿不等，如下表 1 所示。

以下为对 Parti 模型四种大小比较结果，可以观察到：模型性能和输出图像质量在持续地提高；20B 模型尤其擅长于那些抽象的、需要世界知识的、特定视角的、或符号渲染的 prompt。

在悉尼歌剧院前的草地上，一只袋鼠穿着橙色卫衣，戴着蓝色墨镜，胸前挂着「欢迎朋友」的牌子。

松鼠把苹果送给了小鸟。

文本编码器预训练

该研究在两个数据集上预训练文本编码器：具有 BERT [36] 预训练目标的 Colossal Clean Crawled Corpus (C4) [35]，以及具有对比学习目标的图像文本数据。预训练后，该研究继续训练编码器和解码器，在 8192 个离散图像 token 的词汇表上使用 softmax 交叉熵损失生成文本到图像。

预训练后的文本编码器在 GLUE 上的性能与 BERT 相当；然而，在文本到图像生成的完整编码器 - 解码器训练过程之后，文本编码器会降级。

扩展

该研究在 Lingvo 上来实现模型，并在 CloudTPUv4 硬件上使用 GSPMD 进行扩展，以用于训练和推理。GSPMD 是一个基于 XLA 编译器的模型分布系统，它允许将 TPU 集群视为单个虚拟设备，并在几个张量上使用 sharding annotations 来指示编译器自动分发数据并在数千个设备上进行计算。

该研究用数据并行性训练 350M 和 750M 模型。对于 3B 模型，该研究使用 4 路内层模型并行（参见下图 5）和 128 路数据并行。