拳打DALL-E 2脚踢Imagen，谷歌最新Muse模型刷新文本图像合成排行榜
原文链接：https://www.techbeat.net/article-info?id=4501
作者：seven_
论文链接：
https://arxiv.org/abs/2301.00704
项目主页：
https://muse-model.github.io/
近期火爆AI社区的文本图像合成模型家族又添新成员了，之前在这一领域占据上风的是以DALL-E 2[1]和Imagen[2]为代表的扩散模型，以及以Parti[3]为代表的自回归模型。为了进一步提升文本图像合成任务的效率，近日，谷歌研究院再度发布全新基于生成式Transformer架构的Muse模型。不得不说，这一领域发展实在是太卷太快了。谷歌研究团队已经宣布，Muse模型目前已成为这一领域的新SOTA。Muse在训练过程中加入了Mask建模任务，首先给定从预训练的大型语言模型（LLMs）中提取的文本嵌入，Muse会根据文本内容来预测生成随机Mask掉的图像tokens。与DALL-E 2和Imagen等像素空间扩散模型相比，Muse基于离散tokens的训练模式大大减少了模型迭代次数，与Parti等自回归模型相比，作者设计的并行编码方式有效提高了模型的生成效率。此外，Muse可以从预训练的LLMs中获取细粒度的语言知识，从而转化为高精度的视觉概念识别能力。下面展示了一些Muse的文本图像生成效果。
Muse可以让”两个宇航员在埃菲尔铁塔前踢足球“，也可以让”两只喵喵一起做科学研究“。每张示例图像都展现了非常生动有趣的图像生成效果，并且每张图像都可以使用一张TPUv4在1.3秒内生成。
而对于AIGC应用的必备功能，图像编辑任务，Muse也可以胜任。用户可以指定一个Mask区域，随后输入编辑指令，如下图所示，Muse就可以直接将Mask区域更换为纽约或者巴黎。作者还强调，Muse无需对模型进行微调即可实现上述图像编辑功能，实现了”开箱即用“。
一、引言
以扩散模型为代表的文本图像合成模型的爆火，背后离不开新颖的深度学习训练范式，例如对语言和图像两种模态的Mask建模和重构任务等。但是这些工作需要极高的训练成本，而且在推理时仍然费事费力。本文提出的Muse模型在不损失图像生成效果的基础上，主打的就是模型推理效率。Muse凭借离散token编码和推理时的并行解码机制，在相同的硬件条件下（例如TPU-v4芯片）比Imagen-3B或Parti-3B模型的推理速度快10倍以上，比Stable Diffusion v1.4快3倍。此外，Muse为了同时兼顾生成图像的高层语义和低层图像细节，加入了双分辨率编码解码的机制，即同时对每张图像的256x256和512x512分辨率版本进行操作，尽管这样，Muse仍然相比其他扩散模型有明显的速度优势。虽然Muse主打模型效率，但是其图像合成质量仍然非常具有竞争力，Muse在目前公认的文本图像合成基准CC3M上以达到了6.06的FID分数，是目前业内的最高水平。
二、本文方法
之前的扩散模型或者自回归模型均需要庞大的参数量来存储从海量文本图像数据学习到的知识，这加重了模型的训练负担以及部署成本。Muse的提出，标志着另一种新型图像文本框架模式的诞生，作者将其概括为掩码生成式Transformer架构，即在训练时将输入图像的一部分像MAE那样随机Mask掉，然后设计预训练的文本任务来进行重构学习。下图为Muse的整体框架，其中内置了一个预训练文本编码器来对输入prompt进行编码，还包含一对VQGAN tokenizer模型，用来将输入图像编码为一系列离散tokens，两个VQGAN分别处理低分辨率（256x256）和高分辨率（512x512）。此外作者还设计了一个SuperRes Transformer，负责将低分辨率token转换为高分辨率token，实现超分辨率功能。下面我们将详细介绍这些模块。
2.1 预训练文本编码器
作者在模型构建过程中发现，使用预训练的LLMs非常有利于生成高质量的图像，从LLMs中可以提取到图像中有关对象（名词）、动作（动词）、视觉属性（形容词）、空间关系（介词）和其他属性的丰富信息。因此作者在这里内置了一个预训练的T5-XXL[4]编码器，并且冻结其参数，可以对输入文本prompt进行编码，产生维度为4096的语言嵌入向量，这些向量随后被投影到下游的Transformer模型中。
2.2 VQGAN生成离散tokens
基于VQGAN的跨分辨率编码模块是Muse模型的核心操作模块，每个编码模块由一个编码器和一个解码器构成，负责将图像输入转换为一系列特征tokens，为了能够适应不同分辨率的输入，作者将其设计为全卷积结构，给定大小为 H × W H \times W H×W 的图像，编码后的token大小为 H / f × W / f H / f \times W / f H/f×W/f，下采样率为 f f f。 两个VQGAN模块的下采样率分别为 f f f= 16 和 f f f=8。作者在256×256像素图像上使用 f f f=16，产生的token尺寸为16x16。在512x512像素图像上使用 f f f=8，产生的token尺寸为64x64。
2.3 基础模型
在将tokens随机mask后，联合文本嵌入一起送入到基础模型中，这里的基础模型是一个Masked Transformer结构，作者将Mask掉的token区域使用特殊的「MASK」进行标记，然后通过Transformer的自注意力层、交叉注意力层和MLP层来进行特征映射。在模型的输出层，使用MLP将每个Masked图像嵌入转换为一组logits（logits大小对应与VQGAN生成的离散tokens），这样设计就可以直接应用交叉熵损失函数，将原有tokens作为Ground-Truth进行优化。
2.4 超分辨率Transformer
作者发现直接预测512x512分辨率的tokens会导致模型更多的关注细节信息而忽略了对图像高级语义的学习，因此作者采用了双分辨率设计，即首先使用基础模型生成低分辨率token特征图（16x16分辨率，对应于256x256的输入图像），随后将其输入到超分辨率模型中上采样到64x64分辨率（对应于512x512的输入图像），整体流程如下图所示。
通过这种设计，模型可以实现跨分辨率的信息融合，低分辨率的token包含较多的图像语义，高分辨率的token又兼顾了图像细节，随后通过Transformer中的交叉注意力进行特征融合，并使用更高分辨率的VQGAN解码来得到最终的高质量高分辨率生成图像。
2.5 动态遮蔽率和迭代并行编码
作者对Muse的Mask进行了动态调整，使用基于余弦规则的调节方式来实现动态掩蔽率的训练，具体来说，对于每个训练样本，使用采样函数 p ( r ) = 2 π ( 1 − r 2 ) − 1 2 p(r)=\frac{2}{\pi}\left(1-r^{2}\right)^{-\frac{1}{2}} p(r)=π2(1−r2)−21 采样掩蔽率 r ∈ [ 0 , 1 ] r \in[0,1] r∈[0,1] 。与某些自回归方法中来学习条件分布 P ( x i ∣ x < i ) P(x_{i}\mid x_{< i}) P(xi∣x<i) 指导的固定token相比，动态遮蔽率允许模型实现推理阶段的迭代并行编码。迭代并行编码即模型在单个前向传播过程中来同时预测多个输出token，这一操作要求模型在解码时所处理的很多token都是条件独立的，由于这里的解码是通过余弦规则进行调度，该规则要求在当前预测时将最高置信度的区域进行标记，然后在后续步骤中对其进行屏蔽，这样就可以大大提高模型的解码效率，在基础模型中，模型仅通过24个解码步骤就可以完成256个token的推理，而在超分辨率模型中仅使用8个解码步骤就可以执行4096个token推理，这一设计有效提高了Muse模型的推理速度。下图展示了基础模型（左）和超分辨率模型（右）的迭代过程。
三、实验结果
Muse使用由460M规模的文本图像对构成的Imagen数据集进行训练，训练机器为512核的TPU-v4，batchsize设置为512，大概需要一周的时间进行训练。作者首先展示了Muse奇妙的图像生成效果，如下图所示，它可以根据不同属性的文本提示来精确的合成图像。下图左上角展示了Muse对数字概念的理解能力。下图右上角的例子展示了Muse对对象组合和相关性的理解，其不会简单的将对象放置在一个随机位置，而是会提取输入文本中的方位介词，例如”上与下，左与右“等。下图左下角的示例展示了Muse的跨风格图像生成能力，既可以描绘著名艺术家（例如伦勃朗，是欧洲巴洛克绘画艺术的代表画家之一）的绘画风格，又可以概括较为完整的整体艺术风格（例如中国水墨画）。下图右下角部分展示了Muse渲染单词和短语的能力。
此外，作者还将Muse与DALL-E 2和Imagen等明星模型进行了对比，对比效果如下图所示，可以看到，在某些文本条件情况下，Muse的生成效果与Imagen持平，图像质量相较于发布较早的DALL-E 2更好。
随后作者也进行了一系列的量化评估实验，评估数据集使用CC3M和COCO，评价指标为FID和CLIP分数，下面的表格分别展示了Muse在CC3M和COCO上的评估结果。其中Muse的632M版本在CC3M达到了SOTA效果，显着提高了FID和CLIP分数的最新水平。
四、总结
Muse模型的提出，再次证实了使用现有预训练大型语言模型作为训练文本图像合成模型模式的成功，这表明社区努力发展大模型的方向是正确的。此外，作者还表明，基于生成式Transformer架构的Muse模型也可以与现有流行的扩散模型和自回归模型相媲美，同时还具有更高的推理效率，并且还实现了精确的图像编辑功能。但是作者也呼吁社区，在开发生成模型时，也应该考虑模型的安全性，防止其被用于错误信息和危险信息的生成，因此，作者团队选择暂时不发布源代码，需要进行进一步的安全审查。
参考
[1] Ramesh, A., Dhariwal, P., Nichol, A., Chu, C., and Chen, M. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125, 2022.
[2] Saharia, C., Chan, W., Saxena, S., Li, L., Whang, J., Denton, E., Ghasemipour, S. K. S., Ayan, B. K., Mahdavi, S. S., Lopes, R. G., et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint arXiv:2205.11487, 2022.
[3] Yu, J., Xu, Y., Koh, J. Y., Luong, T., Baid, G., Wang, Z., Vasudevan, V., Ku, A., Yang, Y., Ayan, B. K., et al. Scaling autoregressive models for content-rich text-to-image generation. arXiv preprint arXiv:2206.10789, 2022.
[4] Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., Liu, P. J., et al. Exploring the limits of transfer learning with a unified text-to-text transformer. J. Mach. Learn. Res., 21(140):1–67, 2020.
Illustration by Delesign Graphics from IconScout
-The End-
关于我“门”
▼
将门是一家以专注于发掘、加速及投资技术驱动型创业公司的新型创投机构，旗下涵盖将门创新服务、将门-TechBeat技术社区以及将门创投基金。
将门成立于2015年底，创始团队由微软创投在中国的创始团队原班人马构建而成，曾为微软优选和深度孵化了126家创新的技术型创业公司。
如果您是技术领域的初创企业，不仅想获得投资，还希望获得一系列持续性、有价值的投后服务，欢迎发送或者推荐项目给我“门”：
[email protected]