近日，阿里雲機器學習平臺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會被增強爲
，其中
是實體的文本embedding，
是通過TransE算法得到的知識圖譜embedding。雖然我們是在中文場景下對CLIP模型進行知識增強，但是我們的增強方法同樣適用於其他的語言場景。在訓練特定領域的文圖生成模型時，爲了領域知識的對齊，我們將中文CLIP模型的文本編碼器設置爲可訓練。
隱空間噪聲預測器
在得到文本embedding
後,我們在隱空間下使用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
點擊立即免費試用雲產品 開啓雲上實踐之旅！
本文爲阿里雲原創內容，未經允許不得轉載。