近日，阿里云机器学习平台 PAI 与华南理工大学单干（阿里云与华南理工大学联结造就我的项目）在自然语言解决顶级会议 ACL2023 上发表了具备高效推理速度的中文畛域文图生成扩散模型和工具链 Rapid Diffusion。它是面向中文特定畛域的文图生成模型，采纳与 Stable Diffusion 一样的模型构造，在给定中文文本的状况下能够实现疾速的文图生成工作，此外，咱们提供了一键式部署性能，用户能够在集体数据上一键式的进行模型的 Finetune 与推理。
论文：
Bingyan Liu, Weifeng Lin, Zhongjie Duan, Chengyu Wang, Ziheng Wu, Zipeng Zhang, Kui Jia, Lianwen Jin, Cen Chen, Jun Huang. Rapid Diffusion: Building Domain-Specific Text-to-Image Synthesizers with Fast Inference Speed. ACL 2023（Industry Track)
背景
Text-to-Image Synthesis（TIS）是指依据文本输出生成图像的技术，给定一段文本指令，应用计算机程序生成合乎文本内容形容的图像。近年来，随着预训练大模型和扩散模型相干钻研的高速倒退，应用预训练文本编码器和基于扩散的图像生成器的文图生成模型曾经可能生成媲美人类画师的精美图像。然而，因为预训练语言模型不足特定畛域的实体常识且受限于扩散模型的推理速度，目前开源社区的风行文图生成模型难以反对特定工业畛域的利用。次要问题在于，基于扩散的办法须要应用预训练文本编码器对输出文本进行编码，而后作为扩散模型的 UNet 模型的条件输出。然而目前应用网上收集的文本图像对预训练的文本编码器模型不足特定实体概念的理解能力，难以捕捉特定实体常识，这对于生成真切的实体对象图片至关重要。同时，扩散模型的推理速度和计算成本也是须要思考的重要因素，而迭代逆扩散去噪过程的繁琐计算始终是扩散模型推理速度的瓶颈。
为了解决上述问题，钻研具备特定实体对象理解能力的扩散模型以生成具备合乎文本内容形容的高分辨率图像，并开发反对疾速在线推理的优化文图生成模型的框架是必要的。
算法概述
咱们提出了一个新框架：Rapid Diffusion，用于训练和部署文图生成扩散模型，模型架构图 1 所示。Rapid Diffusion 在 stable diffusion 模型的根底上进行了改良。为了晋升对特定实体的理解能力，咱们在 CLIP 的文本编码器中注入了丰盛的实体常识，应用常识图谱进行常识加强。文图生成模模型隐空间的噪声预测器与 stable diffusion，是一个带穿插注意力机制的 U -Net 网络。与开源 Stable Diffusion 间接利用大规模分层扩散模型不同，咱们在图像扩散模块之后集成了一个基于 ESRGAN 的网络，以进步生成图像的分辨率的同时无效解决了参数量爆炸和耗时长的问题。对于在线部署，咱们基于 FlashAttention 优化的神经架构设计了一个高效的推理流程。生成模型计算图的 Intermediate Representation（IR）通过端到端人工智能编译器 BladeDISC 进一步解决，以进步生成模型的推理速度。
常识加强的文本编码器
对于常识加强的文本编码器，咱们聚焦在中文场景下的文图生成问题，为了失去对中文文本以及中文实体常识更具理解能力的文本编码器，咱们应用了 1 亿的悟空图文对数据集作为咱们的文本编码器预训练数据。此外，对于实体常识，咱们应用了最新的中文常识图谱 OpenKG 数据集，蕴含了 1600 万的实体和 1.4 亿个三元关系对，来训练咱们的中文 CLIP 模型。在中文 CLIP 预训练阶段，悟空语料库句子中的实体 token 会被加强为 \(\bar{e}=\bar{e}_{txt}+\bar{e}_{kg} \)，其中 \(\bar{e}_{txt} \) 是实体的文本 embedding，\(\bar{e}_{kg} \) 是通过 TransE 算法失去的常识图谱 embedding。尽管咱们是在中文场景下对 CLIP 模型进行常识加强，然而咱们的加强办法同样实用于其余的语言场景。在训练特定畛域的文图生成模型时，为了畛域常识的对齐，咱们将中文 CLIP 模型的文本编码器设置为可训练。
隐空间噪声预测器
在失去文本 embedding \(\bar{e} \) 后, 咱们在隐空间下应用 latent diffusion 模型生成图像的潜在编码，latent diffusion 是一个带有能够捕获文本条件信息的穿插注意力机制的 Unet 模型，训练时图像重建的损失函数为：
图像生成的过程是扩散的逆过程，从随机采样的高斯噪声中基于条件文本信息生成图像，同样的，为了加强生成图像与文本信息之间的相关性，咱们在训练时应用 classifier-free guidance 的训练方法。为了缩小太大的采样步数带来的工夫开销，咱们应用了 PNDM 算法来缩小采样步数。在咱们框架中，咱们应用了悟空数据集预训练了 latent diffusion 模型，而后再在畛域场景数据下进行微调。
超分辨率网络
咱们的 latent diffusion 模型生成的图像的分辨率为 256*256, 为了失去更高分辨率的图像，与 stable diffusion 用另外一个扩散模型不同，咱们间接应用了训练好的 ESRGAN 模型来进步图像的分辨率，大大晋升了图像生成的速度。
推理减速设计
在推理减速设计这块，咱们剖析了原始 PyTorch 模型的推理速度，并察看到模型推理瓶颈次要在 U -Net 模型，其中 U -Net 中的穿插注意力计算在推理工夫中占主导地位。剖析后果见图 2。为了解决这个问题，咱们联合主动切片和编译优化技术来优化整个技术流水线，并引入 IO 感知注意力实现以进一步提高推理性能。
咱们的推理减速算法通过加强一组两头标识来创立一个残缺的动态图示意来实现的。对于内存拜访密集的操作，咱们充分利用共享内存设计了更大粒度的内核交融策略，无效缩小了 CPU/GPU 之间的切换。执行最佳图分区和内核实现抉择以获得最佳推理速度。
在主动编译优化的根底上，咱们进一步利用 FlashAttention 技术作为 U -Net 的穿插注意力算子，这是网络推理瓶颈的外围。该技术基于注意力 IO 个性，对注意力计算进行 tiling 操作，缩小内存读写计算量。咱们为计算设施和硬件架构以及动静输出的各种组合引入了不同的 FlashAttention 内核实现。通过减速穿插注意力计算，为 U-Net 模块带来 1.9 倍的减速。
算法精度评测
为了评测 Rapid Diffusion 模型，咱们在三个中文图文对数据集（电商，国画，美食）上进行测试，后果如表一所示：
表 1 Rapid Diffusion 与基线模型的性能比照 (FID 分数)。
后果能够证实，Rapid Diffusion 在这些数据集上取得不错的成果。从表 1 能够看出，Rapid Diffusion 在三个数据集上的体现优于所有交易对手，均匀 FID 得分为 21.90。结果表明，咱们针对特定畛域场景的常识加强模型能够更好地了解畛域常识，并且能够生成更真切和多样化的图像。
表 2 文本图像检索常识加强 CLIP 的性能。
因为 CLIP 模型旨在学习跨模态示意，咱们首先通过文本图像检索从实质上评估咱们的模型。咱们应用雷同的预训练文本图像语料库比拟中文 CLIP 模型和咱们的 Chinese Knowledge-enhanced CLIP (CKCLIP) 模型。表 2 报告了测试集上的文本到图像和文本到图像检索后果。咱们的 CKCLIP 模型显着进步了检索性能（尤其是对于 R@1 指标，显示了其学习跨模态示意的能力。
表 3 Rapid Diffusion 的推理减速后果。
在推理速度上，咱们应用了端到端人工智能编译器 BladeDISC 以及 FlashAttention 技术来进步模型的推理速度，同样的，试验后果表 3 所示，咱们的办法能够进步 1.73 倍的推理速度。尽管咱们是在 Rapid Diffusion 上进行试验，然而咱们提出的减速办法具备普适性，同样也实用于其余的 diffusion 模型，譬如 Stable Diffusion 和太乙 Diffusion 模型。咱们还将 Rapid Diffusion 与阿里云机器学习平台 PAI 进行集成，以展现其在理论利用中的实用价值。在阿里云机器学习平台 PAI 上，用户能够在本人的工作（数据）上一键式的进行训练，微调以及推理本人的模型。
在将来，咱们将扩大 Rapid Diffusion 的性能，并通过先进的编译优化技术进一步提高推理速度。为了更好地服务开源社区，咱们模型和源代码行将奉献在自然语言解决算法框架 EasyNLP 中，欢送 NLP 从业人员和研究者应用。
EasyNLP 开源框架：https://github.com/alibaba/EasyNLP
参考文献
- Chengyu Wang, Minghui Qiu, Taolin Zhang, Tingting Liu, Lei Li, Jianing Wang, Ming Wang, Jun Huang, Wei Lin. EasyNLP: A Comprehensive and Easy-to-use Toolkit for Natural Language Processing. EMNLP 2022
- Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. 2022. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10684–10695.
- Jonathan Ho, Ajay Jain, and Pieter Abbeel. 2020. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840– 6851.
- Jiaming Song, Chenlin Meng, and Stefano Ermon. 2021. Denoising diffusion implicit models. In International Conference on Learning Representations.
- Kai Zhu, WY Zhao, Zhen Zheng, TY Guo, PZ Zhao, JJ Bai, Jun Yang, XY Liu, LS Diao, and Wei Lin. 2021. Disc: A dynamic shape compiler for machine learning workloads. In Proceedings of the 1st Workshop on Machine Learning and Systems, pages 89– 95.
- Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. 2022. Flashattention: Fast and memory-efficient exact attention with io-awareness. CoRR, abs/2205.14135.
论文信息
论文题目：Rapid Diffusion: Building Domain-Specific Text-to-Image Synthesizers with Fast Inference Speed
论文作者：刘冰雁、林炜丰、段忠杰、汪诚愚、吴梓恒、张子鹏、贾奎、金连文、陈岑、黄俊
论文 PDF 链接：https://aclanthology.org/2023.acl-industry.28.pdf
收费支付 交互式建模 PAI-DSW、模型训练 PAI-DLC 5000CU* H 计算资源包，以及价值 500 元模型在线服务 PAI-EAS 抵扣包。