基於DiT，支援4K影像生成，華為諾亞0.6B文生圖模型PixArt-Σ來了
机器之心發表於2024-03-11
眾所周知，開發頂級的文生圖（T2I）模型需要大量資源，因此資源有限的個人研究者基本都不可能承擔得起，這也成為了 AIGC（人工智慧內容生成）社群創新的一大阻礙。同時隨著時間的推移，AIGC 社群又能獲得持續更新的、更高質量的資料集和更先進的演算法。於是關鍵的問題來了：我們能以怎樣的方式將這些新元素高效地整合進現有模型，依託有限的資源讓模型變得更強大？為了探索這個問題，華為諾亞方舟實驗室等研究機構的一個研究團隊提出一種新的訓練方法：由弱到強式訓練（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）框架的一種早期嘗試。而現在，隨著 Sora 登上熱搜以及 Stable Diffusion 層出不窮的應用，DiT 架構的有效性得到了研究社群越來越多工作的驗證，例如 PixArt, Dit-3D, GenTron 等「1」。該團隊使用 PixArt-α 的預訓練基礎模型，透過整合高階元素以促進其持續提升，最終得到了一個更加強大的模型 PixArt-Σ。圖 1 展示了一些生成結果示例。
具體來說，為了實現由弱到強式訓練，造出 PixArt-Σ，該團隊採用了以下改進措施。該團隊收集了一個高質量資料集 Internal-Σ，其主要關注兩個方面：(1) 高質量影像：該資料集包含 3300 萬張來自網際網路的高解析度影像，全都超過 1K 解析度，包括 230 萬張解析度大約為 4K 的影像。這些影像的主要特點是美觀度高並且涵蓋廣泛的藝術風格。(2) 密集且準確的描述：為了給上述影像提供更精準和詳細的描述，該團隊將 PixArt-α 中使用的 LLaVA 替換成了一種更強大的影像描述器 Share-Captioner。不僅如此，為了提升模型對齊文字概念和視覺概念的能力，該團隊將文字編碼器（即 Flan-T5）的 token 長度擴充套件到了大約 300 詞。他們觀察到，這些改進可以有效消除模型產生幻覺的傾向，實現更高質量的文字 - 影像對齊。
為了增強 PixArt-α，該團隊將其生成解析度從 1K 提升到了 4K。為了生成超高解析度（如 2K/4K）的影像，token 數量會大幅增長，這就會導致計算需求大幅增長。為了解決這一難題，他們引入了一種專門針對 DiT 框架調整過的自注意力模組，其中使用了鍵和值 token 壓縮。具體來說，他們使用了步長為 2 的分組卷積來執行鍵和值的區域性聚合，如下圖 7 所示。
此外，該團隊還採用了一種專門設計的權重初始化方案，可在不使用 KV（鍵 - 值）壓縮的前提下從預訓練模型實現平滑適應。這一設計可有效將高解析度影像生成的訓練和推理時間降低大約 34%。該團隊提出了多種微調技術，可快速高效地將弱模型調整為強模型。其中包括：(1) 替換使用了一種更強大的變分自動編碼器（VAE）：將 PixArt-α 的 VAE 替換成了 SDXL 的 VAE。(2) 從低解析度到高解析度擴充套件，這個過程為了應對效能下降的問題，他們使用了位置嵌入（PE）插值方法。(3) 從不使用 KV 壓縮的模型演進為使用 KV 壓縮的模型。實驗結果驗證了由弱到強式訓練方法的可行性和有效性。透過上述改進，PixArt-Σ 能以儘可能低的訓練成本和儘可能少的模型引數生成高質量的 4K 解析度影像。具體來說，透過從一個已經預訓練的模型開始微調，該團隊僅額外使用 PixArt-α 所需的 9% 的 GPU 時間，就得到了能生成 1K 高解析度影像的模型。如此表現非常出色，因為其中還替換使用了新的訓練資料和更強大的 VAE。此外，PixArt-Σ 的引數量也只有 0.6B，相較之下，SDXL 和 SD Cascade 的引數量分別為 2.6B 和 5.1B。PixArt-Σ 生成的影像的美觀程度足以比肩當前最頂級的文生圖產品，比如 DALL・E 3 和 MJV6。此外，PixArt-Σ 還展現出了與文字 prompt 細粒度對齊的卓越能力。
圖 2 展示了一張 PixArt-Σ 生成 4K 高解析度影像的結果，可以看到生成結果很好地遵從了複雜且資訊密集的文字指令。
訓練細節：對於執行條件特徵提取的文字編碼器，該團隊按照 Imagen 和 PixArt-α 的做法使用了 T5 的編碼器（即 Flan-T5-XXL）。基礎擴散模型就是 PixArt-α。不同於大多數研究提取固定的 77 個文字 token 的做法，這裡將文字 token 的長度從 PixArt-α 的 120 提升到了 300，因為 Internal-Σ 中整理的描述資訊更加密集，可以提供高細粒度的細節。另外 VAE 使用了來自 SDXL 的已預訓練的凍結版 VAE。其它實現細節與 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-Σ 不僅優於開源模型，而且與當前的閉源產品相比也頗具競爭力，如圖 4 所示。
生成高解析度影像：新方法可以直接生成 4K 解析度的影像，而無需任何後處理。此外，PixArt-Σ 也能準確遵從使用者提供的複雜和詳細的長文字。因此，使用者無需費心去設計 prompt 也能得到讓人滿意的結果。人類 / AI（GPT-4V）偏好研究：該團隊也研究了人類和 AI 對生成結果的偏好。他們收集了 6 個開源模型的生成結果，包括 PixArt-α、PixArt-Σ、SD1.5、Stable Turbo、Stable XL、Stable Cascade 和 Playground-V2.0。他們開發了一個網站，可透過展現 prompt 和對應的影像來收集人類偏好反饋。人類評估者可根據生成質量以及與 prompt 的匹配程度來給影像排名。結果見圖 9 的藍色條形圖。可以看出人類評估者對 PixArt-Σ 的喜愛勝過其它 6 個生成器。相比於之前的文生圖擴散模型，如 SDXL（2.6B 引數）和 SD Cascade（5.1B 引數），PixArt-Σ 能以少得多的引數（0.6B）生成質量更高且更符合使用者 prompt 的影像。
此外，該團隊還使用了先進的多模態模型 GPT-4 Vision 來執行 AI 偏好研究。他們的做法是給 GPT-4 Vision 提供兩張影像，讓它基於影像質量和影像 - 文字對齊程度進行投票。結果見圖 9 中的橙色和綠色條形圖，可以看到情況與人類評估基本一致。該團隊也進行了消融研究來驗證各種改進措施的有效性。更多詳情，請訪問原論文。參考文章：1.https://www.shoufachen.com/Awesome-Diffusion-Transformers/