摘要

 

Text-to-Image模型不断刷新AI作画能力，这些模型可以大致分为两大派系：Autoregression Model和Diffusion Model。本文结合最新的两个模型，Parti和Imagen，来分析两大派系模型的基本工作流程及训练要点。

 

背景介绍

 

今年是Text-to-Image模型蓬勃发展的一年，新模型不断刷新AI作画能力，现在只要将所期望的场景通过文本描述输入到Text-to-Image模型中，就可以得到符合甚至超出自己预期的高分辨率图像。Text-to-Image模型大致分为两大派系：Autoregression Model和Diffusion Model，图1中列举了这两大派系的典型模型。从这些模型的图像生成样例来看，早期的模型更像是拼凑了训练数据集中的相关元素，且所处理的文本描述通常较为简单，而最新的模型已经可以处理很复杂的文本描述，并对文本含义有很精确的阐述。我们就以最新的两个模型，Parti[2]和Imagen[3]为例，来分析两大派系模型的基本工作流程及训练要点。

 

图1 Text-to-Image典型模型图像生成示例

 

Parti

Parti[2]是Google基于多模态AI架构Pathways[10]实现的Text-to-Image模型，其主要模块及工作流程如图2所示，左侧为Transformer Encoder和Transformer Decoder组成的Parti sequence-to-sequence autoregressive model (以下简称text encoder/decoder)，右侧为image tokenizer，使用ViT-VQGAN[11]实现，其基础结构也是transformer。

 

图2 Parti主要模块及工作流程[2]

 

Parti的训练主要分为两个阶段，第一阶段在图像数据集上单独训练image tokenizer，所使用的图像大小为256x256，模型配置为ViT-VQGAN-Small (8 blocks, 8 heads, model dimension 512, and hidden dimension 2048)，对应参数量为30M，codebook的token类别数为8192，即把图像数据映射到8192种不同的visual token上。

 

第二阶段会固定image tokenizer，首先在纯文本数据和图文数据上对text encoder进行预训练，令其掌握一定的文本理解能力，然后再同时训练text encoder和text decoder。比较有意思的一点是，text encoder在预训练结束后的文本理解能力与BERT[12]相当，但在随后text encoder-decoder同时训练的过程中，text encoder的文本理解能力有所下降，说明当前模型从纯文本数据中获取的自然语言表征与从图文数据中获取的表征存在一定差异，如何更好地理解和处理这些差异有待进一步研究。从测试结果来看，text encoder预训练有助于提升最终的text-to-image效果，故在训练参数量更大的Parti模型时还是沿用了text encoder预训练的做法。

 

在两个主要的训练阶段结束后，image tokenizer会再做一次调整，固定其encoder(即图2中Image Tokenizer)和codebook，在其decoder(即图2中Image Detokenizer)中使用尺寸更大的配置(32 blocks, 16 heads, model dimension 1280, and hidden dimension 5120)，并对decoder的参数进行微调。

 

上述步骤结束后便可将模型用于文本到图像生成，此时得到的图像大小为256x256，虽然已经能够呈现大部分的内容、结构、以及纹理信息，但其观感还是不如更高分辨率的图像。此时可以借助超分辨率模型，增加一个Super-Resolution Upsampler模块，将图像尺寸高清放大。用于图像放大的模块参数量与放大尺寸有关，放大到512x512时对应参数量为15M，放大到1024x1024时对应参数量为30M。超分辨率模块的训练方法如图3，训练时固定image tokenizer的参数，此时不需要text encoder/decoder参与。

 

图3 Super-Resolution Upsampler模块训练方法[2]

 

不同参数量的Parti模型具体配置及其文本到图像生成能力如图4，可以看到，随参数量的增加，模型对文本的理解更为准确，对文本中信息的提取更为完整，所生成的图像也更加合理、真实。

 

图4 不同参数量的Parti模型配置及文本到图像生成效果[2]

 

Imagen

Imagen[3]也是Google出品，但与Parti所采用的autoregressive model不同，Imagen使用diffusion model，其推理过程可以理解为图像连续加噪的逆过程，从白噪声图像开始，学习到数据分布的模型每次预测当前（小幅）噪声统计信息用于去除图像噪声，经过多个timestep迭代，得到生成图像，详细介绍可以参考我们之前的文章——深度生成网络新思路：扩散概率模型。

 

Imagen主要模块及工作流程如图5，输入的文本描述由Frozen Text Encoder编码成Text Embedding， Text-to-Image Diffusion Model根据Text Embedding从随机噪声图迭代产生一幅与输入描述匹配的64x64小图，小图由Super-Resolution模块根据Text Embedding放大为256x256的中等尺寸图像，中等尺寸图像由另一个Super-Resolution模块同样根据Text Embedding最终产生1024x1024的高分辨率图像。

 

图5 Imagen主要模块及工作流程[3]

 

Imagen所采用的text encoder为纯文本预训练好的T5模型[13]的encoder，测试结果显示T5-XXL效果最优，其对应参数量为4.6B。

 

Text-to-Image Diffusion Model采用U-Net结构[14]，如图6所示， Imagen在U-Net的基础上插入了一些注意力层，这样能更好地利用文本信息。

 

图6 Imagen主要模块及工作流程[15]

 

Text embedding的信息以pooled embedding vector的形式提供给模型，实现方法如图7所示。

 

图7 Pooled embedding vector示意图[15]

 

Super-Resolution模块同样使用U-Net结构实现，为减少显存占用、加速收敛、提升推理速度，Imagen对其进行了优化，称为Efficient U-Net，其整体结构如图8。主要改进措施为：

1. 在低分辨率部分使用更多ResNetBlock，使得模型参数量分布偏向低分辨率部分；

2. 将skip connections缩放一个系数， 1/sqrt(2)，有助于模型更快收敛 ；

3. DBlock模块先降采样再卷积，UBlock模块先卷积再上采样，有助于提升模型推理速度。

 

图8 Efficient U-Net结构示意图 (左侧：整体结构；右侧：蓝色框内为DBlock结构，红色框内为UBlock结构，金色框内为ResNetBlock结构)[3]

 

Imagen会分别训练3个Diffusion Model (Text-to-Image 模块和两个Super-Resolustion模块)，整个训练过程有很多细节上的设计，感兴趣的读者可以阅读Imagen论文[3]，这里主要介绍预测值的thresholding策略。Imagen训练使用classifier-free guidance[16]，可以通过一个guidance weight值来调整文本信息注入的强度，一些实验结果表明，这个值越大时，生成的图像与文本相关度越高，但图像质量会下降。Imagen测试过程中发现这是训练和推理时预测值的范围不一致导致的，据此提出了两种thresholding的方案：static thresholding (将超出预测范围的值截断)和dynamic thresholding (将预测值整体除以一定系数)，测试结果如图9，可以看到dynamic thresholding配合较大的guidance weight值可以生成质量更高的图像。

 

图9 Thresholding策略比较(从上到下guidance weight值逐步增大)[3]

 

总结与思考

本文结合Parti[2]和Imagen[3]介绍了Text-to-Image两类主流模型，Autoregression Model和Diffusion Model的基本工作流程和训练要点，两类模型在大量数据驱动下都能达到很好的文本到图像生成效果。除了使用大量的数据，Parti和Imagen对于模型结构和训练过程也做了很多细节上的测试和调整，对于模型最终文本到图像生成效果以及模型训练效率的提升都很有帮助。另外，对于文本处理模块，Parti和Imagen做出了不同的选择，Parti使用图文数据来调整文本处理模块的可训练参数，而Imagen使用固定的纯文本预训练的文本处理模块，不会在图文数据上对其进行进一步调整。在Parti的实验结果中，可以看到经过图文数据调整后，文本处理模块的纯文本理解能力下降。如何更好地处理文本表征和图像表征之间的关系仍待进一步探索。

 

参考文献

[1] Browne, Kieran. "Who (or what) is an AI Artist?." Leonardo 55.2 (2022): 130-134.

[2] Yu, Jiahui, et al. "Scaling autoregressive models for content-rich text-to-image generation." arXiv preprint arXiv:2206.10789 (2022).

[3] Saharia, Chitwan, et al. "Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding." arXiv preprint arXiv:2205.11487 (2022).

[4] Ramesh, Aditya, et al. "Zero-shot text-to-image generation." International Conference on Machine Learning. PMLR, 2021.

[5] Ding, Ming, et al. "Cogview: Mastering text-to-image generation via transformers." Advances in Neural Information Processing Systems 34 (2021): 19822-19835.

[6] Gafni, Oran, et al. "Make-a-scene: Scene-based text-to-image generation with human priors." arXiv preprint arXiv:2203.13131 (2022).

[7] Nichol, Alex, et al. "Glide: Towards photorealistic image generation and editing with text-guided diffusion models." arXiv preprint arXiv:2112.10741 (2021).

[8] Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.

[9] Ramesh, Aditya, et al. "Hierarchical text-conditional image generation with clip latents." arXiv preprint arXiv:2204.06125 (2022).

[10] Introducing Pathways: A next-generation AI architecture. https://blog.google/technology/ai/introducing-pathways-next-generation-ai-architecture/

[11] Jiahui Yu, et al. “Vector-quantized image modeling with improved vqgan.” arXiv preprint arXiv:2110.04627, 2021.

[12] Jacob Devlin, et al. “Bert: Pre-training of deep bidirectional transformers for language understanding.” arXiv preprint arXiv:1810.04805, 2018.

[13] Colin Raffel, et al. “Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer.” JMLR, 21(140), 2020.

[14] Alex Nichol, et al. Improved denoising diffusion probabilistic models. arXiv preprint arXiv:2102.09672, 2021.

[15] How Imagen Works: A Deep Dive. https://www.assemblyai.com/blog/how-imagen-actually-works/#how-imagen-works-a-deep-dive

[16] Jonathan Ho et al. “Classifier-free diffusion guidance.” In NeurIPS 2021 Workshop on Deep Generative Models and Downstream Applications, 2021.