基于DiT，支持4K圖像生成，華為諾亞0.6B文生圖模型PixArt-Σ來了

　　眾所周知，開發(fā)頂級的文生圖(T2I)模型需要大量資源，因此資源有限的個人研究者基本都不可能承擔得起，這也成為了 AIGC(人工智能內(nèi)容生成)社區(qū)創(chuàng)新的一大阻礙。同時隨著時間的推移，AIGC 社區(qū)又能獲得持續(xù)更新的、更高質(zhì)量的數(shù)據(jù)集和更先進的算法。

　　于是關鍵的問題來了:我們能以怎樣的方式將這些新元素高效地整合進現(xiàn)有模型，依托有限的資源讓模型變得更強大?

　　為了探索這個問題，華為諾亞方舟實驗室等研究機構的一個研究團隊提出一種新的訓練方法:由弱到強式訓練(weak-to-strong training)。

　　他們的研究基于他們?nèi)ツ晔�月提出的一種高效的文生圖訓練方法 PixArt-α，參閱機器之心報道《超低訓練成本文生圖模型 PixArt 來了，效果媲美 MJ，只需 SD10% 訓練時間》。PixArt-α 是 DiT(擴散 Transformer)框架的一種早期嘗試。而現(xiàn)在，隨著 Sora 登上熱搜以及 Stable Diffusion 層出不窮的應用，DiT 架構的有效性得到了研究社區(qū)越來越多工作的驗證，例如 PixArt， Dit-3D， GenTron 等「1」。

　　該團隊使用 PixArt-α 的預訓練基礎模型，通過整合高級元素以促進其持續(xù)提升，最終得到了一個更加強大的模型 PixArt-Σ。圖1展示了一些生成結果示例。

　　PixArt-Σ 如何煉成?

　　具體來說，為了實現(xiàn)由弱到強式訓練，造出 PixArt-Σ，該團隊采用了以下改進措施。

　　更高質(zhì)量的訓練數(shù)據(jù)

　　該團隊收集了一個高質(zhì)量數(shù)據(jù)集 Internal-Σ，其主要關注兩個方面:

　　(1) 高質(zhì)量圖像:該數(shù)據(jù)集包含3300萬張來自互聯(lián)網(wǎng)的高分辨率圖像，全都超過1K 分辨率，包括230萬張分辨率大約為4K 的圖像。這些圖像的主要特點是美觀度高并且涵蓋廣泛的藝術風格。

　　(2) 密集且準確的描述:為了給上述圖像提供更精準和詳細的描述，該團隊將 PixArt-α 中使用的 LLaVA 替換成了一種更強大的圖像描述器 Share-Captioner。

　　不僅如此，為了提升模型對齊文本概念和視覺概念的能力，該團隊將文本編碼器(即 Flan-T5)的 token 長度擴展到了大約300詞。他們觀察到，這些改進可以有效消除模型產(chǎn)生幻覺的傾向，實現(xiàn)更高質(zhì)量的文本 - 圖像對齊。

　　下表1展示了不同數(shù)據(jù)集的統(tǒng)計數(shù)據(jù)。

　　高效的 token 壓縮

　　為了增強 PixArt-α，該團隊將其生成分辨率從1K 提升到了4K。為了生成超高分辨率(如2K/4K)的圖像，token 數(shù)量會大幅增長，這就會導致計算需求大幅增長。

　　為了解決這一難題，他們引入了一種專門針對 DiT 框架調(diào)整過的自注意力模塊，其中使用了鍵和值 token 壓縮。具體來說，他們使用了步長為2的分組卷積來執(zhí)行鍵和值的局部聚合，如下圖7所示。

　　此外，該團隊還采用了一種專門設計的權重初始化方案，可在不使用 KV(鍵 - 值)壓縮的前提下從預訓練模型實現(xiàn)平滑適應。這一設計可有效將高分辨率圖像生成的訓練和推理時間降低大約34%。

　　由弱到強式訓練策略

　　該團隊提出了多種微調(diào)技術，可快速高效地將弱模型調(diào)整為強模型。其中包括:

　　(1) 替換使用了一種更強大的變分自動編碼器(VAE):將 PixArt-α 的 VAE 替換成了 SDXL 的 VAE。

　　(2) 從低分辨率到高分辨率擴展，這個過程為了應對性能下降的問題，他們使用了位置嵌入(PE)插值方法。

　　(3) 從不使用 KV 壓縮的模型演進為使用 KV 壓縮的模型。

　　實驗結果驗證了由弱到強式訓練方法的可行性和有效性。

　　通過上述改進，PixArt-Σ 能以盡可能低的訓練成本和盡可能少的模型參數(shù)生成高質(zhì)量的4K 分辨率圖像。

　　具體來說，通過從一個已經(jīng)預訓練的模型開始微調(diào)，該團隊僅額外使用 PixArt-α 所需的9% 的 GPU 時間，就得到了能生成1K 高分辨率圖像的模型。如此表現(xiàn)非常出色，因為其中還替換使用了新的訓練數(shù)據(jù)和更強大的 VAE。

　　此外，PixArt-Σ 的參數(shù)量也只有0.6B，相較之下，SDXL 和 SD Cascade 的參數(shù)量分別為2.6B 和5.1B。

　　PixArt-Σ 生成的圖像的美觀程度足以比肩當前最頂級的文生圖產(chǎn)品，比如 DALL・E3和 MJV6。此外，PixArt-Σ 還展現(xiàn)出了與文本 prompt 細粒度對齊的卓越能力。

　　圖2展示了一張 PixArt-Σ 生成4K 高分辨率圖像的結果，可以看到生成結果很好地遵從了復雜且信息密集的文本指令。

　　實驗

　　實現(xiàn)細節(jié)

　　訓練細節(jié):對于執(zhí)行條件特征提取的文本編碼器，該團隊按照 Imagen 和 PixArt-α 的做法使用了 T5的編碼器(即 Flan-T5-XXL)�；�礎擴散模型就是 PixArt-α。不同于大多數(shù)研究提取固定的77個文本 token 的做法，這里將文本 token 的長度從 PixArt-α 的120提升到了300，因為 Internal-Σ 中整理的描述信息更加密集，可以提供高細粒度的細節(jié)。另外 VAE 使用了來自 SDXL 的已預訓練的凍結版 VAE。其它實現(xiàn)細節(jié)與 PixArt-α 一樣。