2024-03-11 13:18
基于DiT，支持4K圖像生成，華為諾亞0.6B文生圖模型PixArt-Σ來了
本文轉載自機器之心
眾所周知，開發(fā)頂級的文生圖(T2I)模型需要大量資源，因此資源有限的個人研究者基本都不可能承擔得起，這也成為了 AIGC(人工智能內容生成)社區(qū)創(chuàng)新的一大阻礙。同時隨著時間的推移，AIGC 社區(qū)又能獲得持續(xù)更新的、更高質量的數據集和更先進的算法。
于是關鍵的問題來了：我們能以怎樣的方式將這些新元素高效地整合進現(xiàn)有模型，依托有限的資源讓模型變得更強大?
為了探索這個問題，華為諾亞方舟實驗室等研究機構的一個研究團隊提出一種新的訓練方法：由弱到強式訓練(weak-to-strong training)。
論文標題：PixArt-Σ: Weak-to-Strong Training of Diffusion Transformer for 4K Text-to-Image Generation
論文地址：https://arxiv.org/pdf/2403.04692.pdf
項目頁面：https://pixart-alpha.github.io/PixArt-sigma-project/
他們的研究基于他們去年十月提出的一種高效的文生圖訓練方法 PixArt-α，參閱機器之心報道《超低訓練成本文生圖模型 PixArt 來了，效果媲美 MJ，只需 SD 10% 訓練時間》。PixArt-α 是 DiT(擴散 Transformer)框架的一種早期嘗試。而現(xiàn)在，隨著 Sora 登上熱搜以及 Stable Diffusion 層出不窮的應用，DiT 架構的有效性得到了研究社區(qū)越來越多工作的驗證，例如 PixArt, Dit-3D, GenTron 等「1」。
該團隊使用 PixArt-α 的預訓練基礎模型，通過整合高級元素以促進其持續(xù)提升，最終得到了一個更加強大的模型 PixArt-Σ。圖 1 展示了一些生成結果示例。
具體來說，為了實現(xiàn)由弱到強式訓練，造出 PixArt-Σ，該團隊采用了以下改進措施。
更高質量的訓練數據
該團隊收集了一個高質量數據集 Internal-Σ，其主要關注兩個方面：
(1) 高質量圖像：該數據集包含 3300 萬張來自互聯(lián)網的高分辨率圖像，全都超過 1K 分辨率，包括 230 萬張分辨率大約為 4K 的圖像。這些圖像的主要特點是美觀度高并且涵蓋廣泛的藝術風格。
(2) 密集且準確的描述：為了給上述圖像提供更精準和詳細的描述，該團隊將 PixArt-α 中使用的 LLaVA 替換成了一種更強大的圖像描述器 Share-Captioner。
不僅如此，為了提升模型對齊文本概念和視覺概念的能力，該團隊將文本編碼器(即 Flan-T5)的 token 長度擴展到了大約 300 詞。他們觀察到，這些改進可以有效消除模型產生幻覺的傾向，實現(xiàn)更高質量的文本 - 圖像對齊。
下表 1 展示了不同數據集的統(tǒng)計數據。
高效的 token 壓縮
為了增強 PixArt-α，該團隊將其生成分辨率從 1K 提升到了 4K。為了生成超高分辨率(如 2K/4K)的圖像，token 數量會大幅增長，這就會導致計算需求大幅增長。
為了解決這一難題，他們引入了一種專門針對 DiT 框架調整過的自注意力模塊，其中使用了鍵和值 token 壓縮。具體來說，他們使用了步長為 2 的分組卷積來執(zhí)行鍵和值的局部聚合，如下圖 7 所示。
此外，該團隊還采用了一種專門設計的權重初始化方案，可在不使用 KV(鍵 - 值)壓縮的前提下從預訓練模型實現(xiàn)平滑適應。這一設計可有效將高分辨率圖像生成的訓練和推理時間降低大約 34%。
由弱到強式訓練策略
該團隊提出了多種微調技術，可快速高效地將弱模型調整為強模型。其中包括：
(1) 替換使用了一種更強大的變分自動編碼器(VAE)：將 PixArt-α 的 VAE 替換成了 SDXL 的 VAE。
(2) 從低分辨率到高分辨率擴展，這個過程為了應對性能下降的問題，他們使用了位置嵌入(PE)插值方法。
(3) 從不使用 KV 壓縮的模型演進為使用 KV 壓縮的模型。
實驗結果驗證了由弱到強式訓練方法的可行性和有效性。
通過上述改進，PixArt-Σ 能以盡可能低的訓練成本和盡可能少的模型參數生成高質量的 4K 分辨率圖像。
具體來說，通過從一個已經預訓練的模型開始微調，該團隊僅額外使用 PixArt-α 所需的 9% 的 GPU 時間，就得到了能生成 1K 高分辨率圖像的模型。如此表現(xiàn)非常出色，因為其中還替換使用了新的訓練數據和更強大的 VAE。
此外，PixArt-Σ 的參數量也只有 0.6B，相較之下，SDXL 和 SD Cascade 的參數量分別為 2.6B 和 5.1B。
PixArt-Σ 生成的圖像的美觀程度足以比肩當前最頂級的文生圖產品，比如 DALL?E 3 和 MJV6.此外，PixArt-Σ 還展現(xiàn)出了與文本 prompt 細粒度對齊的卓越能力。
圖 2 展示了一張 PixArt-Σ 生成 4K 高分辨率圖像的結果，可以看到生成結果很好地遵從了復雜且信息密集的文本指令。
實現(xiàn)細節(jié)
訓練細節(jié)：對于執(zhí)行條件特征提取的文本編碼器，該團隊按照 Imagen 和 PixArt-α 的做法使用了 T5 的編碼器(即 Flan-T5-XXL)?；A擴散模型就是 PixArt-α。不同于大多數研究提取固定的 77 個文本 token 的做法，這里將文本 token 的長度從 PixArt-α 的 120 提升到了 300.因為 Internal-Σ 中整理的描述信息更加密集，可以提供高細粒度的細節(jié)。另外 VAE 使用了來自 SDXL 的已預訓練的凍結版 VAE。其它實現(xiàn)細節(jié)與 PixArt-α 一樣。
模型是基于 PixArt-α 的 256px 預訓練檢查點開始微調的，并使用了位置嵌入插值技術。
最終的模型(包括 1K 分辨率)是在 32 塊 V100 GPU 上訓練的。他們還額外使用了 16 塊 A100 GPU 來訓練 2K 和 4K 圖像生成模型。
評估指標：為了更好地展示美觀度和語義能力，該團隊收集了 3 萬對高質量文本 - 圖像，以對最強大的文生圖模型進行基準評估。這里主要是通過人類和 AI 偏好來評估 PixArt-Σ，因為 FID 指標可能無法適當地反映生成質量。
性能比較
圖像質量評估：該團隊定性地比較了 PixArt-Σ 與閉源文生圖(T2I)產品和開源模型的生成質量。如圖 3 所示，相比于開源模型 SDXL 和該團隊之前的 PixArt-α，PixArt-Σ 生成的人像的真實感更高，并且也有更好的語義分析能力。與 SDXL 相比，PixArt-Σ 能更好地遵從用戶指令。
PixArt-Σ 不僅優(yōu)于開源模型，而且與當前的閉源產品相比也頗具競爭力，如圖 4 所示。
生成高分辨率圖像：新方法可以直接生成 4K 分辨率的圖像，而無需任何后處理。此外，PixArt-Σ 也能準確遵從用戶提供的復雜和詳細的長文本。因此，用戶無需費心去設計 prompt 也能得到讓人滿意的結果。
人類 / AI(GPT-4V)偏好研究：該團隊也研究了人類和 AI 對生成結果的偏好。他們收集了 6 個開源模型的生成結果，包括 PixArt-α、PixArt-Σ、SD1.5、Stable Turbo、Stable XL、Stable Cascade 和 Playground-V2.0.他們開發(fā)了一個網站，可通過展現(xiàn) prompt 和對應的圖像來收集人類偏好反饋。
人類評估者可根據生成質量以及與 prompt 的匹配程度來給圖像排名。結果見圖 9 的藍色條形圖。
可以看出人類評估者對 PixArt-Σ 的喜愛勝過其它 6 個生成器。相比于之前的文生圖擴散模型，如 SDXL(2.6B 參數)和 SD Cascade(5.1B 參數)，PixArt-Σ 能以少得多的參數(0.6B)生成質量更高且更符合用戶 prompt 的圖像。
此外，該團隊還使用了先進的多模態(tài)模型 GPT-4 Vision 來執(zhí)行 AI 偏好研究。他們的做法是給 GPT-4 Vision 提供兩張圖像，讓它基于圖像質量和圖像 - 文本對齊程度進行投票。結果見圖 9 中的橙色和綠色條形圖，可以看到情況與人類評估基本一致。
該團隊也進行了消融研究來驗證各種改進措施的有效性。更多詳情，請訪問原論文。
相關文章
數字生命卡茲克
yaoyao
AI設計帆哥
