自回歸模型優于擴散：Llama用于可擴展的圖像生成

paper是香港大學發布在Arxiv2024的工作

paper title:Autoregressive Model Beats Diffusion: Llama for Scalable Image Generation

Code:鏈接

Abstract

我們介紹了LlamaGen，一種新的圖像生成模型系列，它將大型語言模型的原始“下一個token預測”范式應用于視覺生成領域。這是對一個問題的肯定回答：如果適當擴展，純粹的自回歸模型，例如Llama，沒有針對視覺信號的歸納偏置，能否實現最先進的圖像生成性能。我們重新審視了圖像標記器的設計空間、圖像生成模型的可擴展性特性以及它們的訓練數據質量。這個探索的結果包括：（1）一個具有16的下采樣比率、0.94 rFID重建質量和97%代碼本使用率的圖像標記器，在ImageNet基準測試中表現出色。（2）一系列從111M到3.1B參數的類條件圖像生成模型，在ImageNet 256×256基準測試中取得了2.18的FID，超過了流行的擴散模型，如LDM、DiT。（3）一個具有775M參數的文本條件圖像生成模型，經過在LAION-COCO和高美學質量圖像上的兩階段訓練，展示了競爭性的視覺質量和文本對齊性能。（4）我們驗證了LLM服務框架在優化圖像生成模型推理速度方面的有效性，取得了326% - 414%的加速效果。我們發布了所有模型和代碼，以促進視覺生成和多模態基礎模型的開源社區發展。

圖1：使用純粹自回歸模型進行圖像生成。我們展示了來自我們類條件圖像生成模型（上排）和文本條件圖像生成模型（下排）的樣本。

1 Introduction

基于自回歸模型，大型語言模型（LLMs）[Vaswani et al 2017; Devlin et al 2018; Radford et al 2018; Raffel et al 2020; Radford et al 2019; Brown et al 2020; Zhang et al 2022]通過預測序列中的下一個token生成文本。這種“下一個token預測”范式在以類人對話方式解決語言任務方面展示了前所未有的能力[Ouyang et al 2022; OpenAI 2022, 2023b; Google 2023; Anthropic 2023; Workshop et al 2022; Touvron et al 2023a,b; Bai et al 2023a; Yang et al 2023; Team 2023; Bi et al 2024]，并展現了驚人的可擴展性[Kaplan et al 2020; Henighan et al 2020; Hoffmann et al 2022; Wei et al 2022; Alabdulmohsin et al 2022; Chowdhery et al 2023; Anil et al 2023]，展示了通用人工智能模型的有希望的發展路徑。

見證了自回歸模型在大型語言模型中的可擴展性，開創性工作嘗試在圖像生成中探索自回歸模型，例如VQVAE [Van Den Oord et al 2017; Razavi et al 2019]、VQGAN [Esser et al 2021; Lee et al 2022]、DALL-E [Ramesh et al 2021]、Parti [Yu et al 2021, 2022]。它們引入了圖像標記器，將連續圖像轉換為離散標記，并應用自回歸模型以下一個token預測的方式生成圖像標記。它們在2022年前展示了在同類中的強大表現[Brock et al 2018; Ho et al 2020; Dhariwal & Nichol 2021]。然而，它們的開源社區發展不足，這在很大程度上限制了它們的進一步改進。

在同一時期，另一種圖像生成方法——擴散模型[Song & Ermon 2019; Ho et al 2020; Song et al 2020; Dhariwal & Nichol 2021; Nichol et al 2021; Lu et al 2022a; Ho et al 2022a; Ho & Salimans 2022; Rombach et al 2022; Ramesh et al 2022; Saharia et al 2022; Rombach et al 2022]快速發展。隨著它們的開源社區的擴展，它們主導了今天的視覺生成領域。然而，擴散模型與自回歸語言模型有著不同的范式，這使得構建一個統一的語言與視覺模型成為一項巨大的挑戰。

在本工作中，我們致力于進一步推動自回歸模型在圖像生成中的發展：繼續其研究方法，并為開源社區做出貢獻。回顧2024年前的圖像生成文獻，我們識別出現有先進模型的三大關鍵[Peebles & Xie 2023; Podell et al 2023; Xue et al 2023; Chen et al 2023b,c; Betker et al 2023; Li et al 2024; Esser et al 2024]：1）設計良好的圖像壓縮器，2）可擴展的圖像生成模型，3）高質量的訓練數據。受到這一啟發，我們重新審視了圖像標記器（自回歸模型的圖像壓縮器）的設計，圖像生成模型的可擴展性特性，以及訓練數據的影響。

為了實現語言與視覺之間潛在的統一模型，我們的設計是減少對視覺信號的歸納偏置，并采用與LLM相同的架構。這屬于與近期工作[Chang et al 2022; Yu et al 2023b; Tian et al 2024]在視覺導向設計指導下修改架構的不同研究哲學。例如，MaskGIT [Chang et al 2022]、MAGVIT [Yu et al 2023a,b]采用了遮擋圖像建模策略，VAR [Tian et al 2024]使用了分層多尺度特性。盡管它們在實現領先的圖像生成性能方面取得了成功，甚至超過了擴散模型，但仍不清楚原始語言模型架構是否能夠實現這一點。相反，我們的工作表明，應用與語言模型完全相同的“下一個token預測”的原始自回歸模型同樣能夠實現最先進的圖像生成性能。作為額外的好處，我們可以利用LLM社區開發的技術[Dao et al 2022; Rasley et al 2020; Shoeybi et al 2019; Zhao et al 2023; Kwon et al 2023; Chen et al 2023a; Dettmers 2022]來優化訓練配方和模型推理速度。

總之，我們的貢獻包括：1. 圖像標記器：一個下采樣比率為16的圖像標記器，在ImageNet基準測試中實現了0.94 rFID的重建質量和97%的代碼本使用率。使用下采樣比率為8時，我們的標記器在擴散模型中使用的連續VAE [Rombach et al 2022; Podell et al 2023; OpenAI 2023a]方面具有競爭力，甚至表現得更好。這表明，在圖像標記器中的離散表示不再是圖像重建的瓶頸。

2. 可擴展的圖像生成模型：開發了一系列類條件圖像生成模型，參數范圍從111M到3.1B，基于Llama架構[Touvron et al 2023a,b]。最大模型在ImageNet 256×256基準測試中實現了2.18的FID，超越了流行的擴散模型如LDM [Rombach et al 2022]、DiT [Peebles & Xie 2023]。這表明，沒有視覺信號歸納偏置的純粹自回歸模型可以作為圖像生成系統的基礎。

3. 高質量的訓練數據：一個具有775M參數的文本條件圖像生成模型，首先在LAION-COCO的50M子集上進行訓練，然后在1000萬張內部高美學質量圖像上進行微調。它展示了在視覺質量和文本對齊方面的競爭性表現。

4. 優化的推理速度：我們采用了vLLM [Kwon et al 2023]，這是最流行的LLM服務框架之一，來優化我們圖像生成模型的推理速度，并實現了326% - 414%的加速效果。

我們發布了所有模型和代碼，以促進視覺生成和多模態基礎模型的開源社區發展。值得注意的是，我們發布的模型仍然落后于基于擴散模型的最先進視覺生成模型 [Alpha-VLLM 2024; Esser et al 2024; Brooks et al 2024]。當未來有更多的訓練數據和計算資源時，我們將探索基于自回歸的大規模視覺生成模型，例如超過7B參數的模型。

2 Autoregressive Models for Image Generation

2.1 Overview

---

首先，圖像像素$x\in {\mathbb{R}}^{H\times W\times 3}$通過圖像標記器被量化為$q\in {\mathbb{Q}}^{h\times w}$的離散標記[Van Den Oord et al. 2017; Esser et al. 2021; Yu et al. 2021]，其中$h=H/p$，$w=W/p$，$p$是圖像標記器的下采樣比率，$q^{(i,j)}$是圖像碼本中的索引。這些圖像標記被重新排列為$h?w$長度的序列，按光柵掃描順序排列，并用于訓練基于Transformer [Vaswani et al. 2017] 的自回歸模型。

在圖像生成過程中，圖像標記$(q_1,q_2,\dots ,q_{h?w})$由自回歸模型生成[Radford et al. 2018, 2019; Brown et al. 2020; Touvron et al. 2023a]，方式是通過下一個token預測：
$$\prod _{t=1}^{h?w}p(q_t\mid q_{<t},c),$$
其中$c$是類別標簽嵌入或文本嵌入。最后，這些圖像標記通過圖像標記器解碼器轉換為圖像像素。

2.2 Image Tokenizer

---

量化自編碼器架構。我們使用與VQGAN [Esser et al. 2021] 相同的架構，即編碼器-量化器-解碼器。編碼器和解碼器是帶有下采樣比率$p$的卷積網絡。量化器包含一個碼本$Z\in {\mathbb{R}}^{K\times C}$，其中有$K$個可學習的向量。編碼器將圖像像素$x$投影到特征圖$f$。量化過程將特征圖中每個向量$f^{(i,j)}$映射到其最近的碼本向量$z^{(i,j)}$的索引$q^{(i,j)}$。在解碼過程中，索引$q^{(i,j)}$會重新映射為特征向量$z^{(i,j)}$，解碼器將這些特征向量轉換回圖像像素$\hat{x}$。

碼本對圖像標記器的性能有關鍵影響。參考[Yu et al. 2021]，我們對碼本向量使用${?}_2$歸一化，采用較低的碼本維度$C$和較大的碼本大小$K$。這些設計顯著提升了重建質量和碼本使用率。更多細節將在實驗中討論。

訓練損失。由于量化是不可導操作，因此使用直通梯度估計器 [Bengio et al. 2013] 來保留從解碼器到編碼器的梯度。定義：
$$z=\text{sg}[z?f]+f,$$
其中$\text{sg}[?]$是stop-gradient操作。碼本學習的損失函數為：
$${\mathcal{L}}_{VQ}=\Vert \text{sg}[f]?z{\Vert }_2^2+\beta \Vert f?\text{sg}[z]{\Vert }_2^2,$$
第二項是承諾損失 [Van Den Oord et al. 2017]，鼓勵編碼器提取的特征向量接近碼本向量，$\beta$是承諾損失的權重。為簡化起見，我們沒有添加熵損失 [Yu et al. 2023a; Chang et al. 2022]。

對于圖像重建訓練，損失為：
$${\mathcal{L}}_{AE}={?}_2(x,\hat{x})+{\mathcal{L}}_P(x,\hat{x})+{\lambda }_G{\mathcal{L}}_G(\hat{x}),$$
其中${?}_2$是圖像像素的重建損失，${\mathcal{L}}_P(?)$ 是來自LPIPS [Zhang et al. 2018] 的感知損失，${\mathcal{L}}_G(?)$ 是使用與圖像標記器同時訓練的PatchGAN [Isola et al. 2017] 判別器得到的對抗損失，${\lambda }_G$ 是對抗損失的權重。

表1：LlamaGen的模型規模和架構配置。配置參考了前人的工作 [Radford et al 2019; Touvron et al 2023a; OpenLM-Research 2023]。

2.3 Image Generation by Autoregressive Models

---

Llama 架構。我們的模型架構主要基于 Llama [Touvron et al 2023a,b]，采用 RMSNorm [Zhang & Sennrich 2019] 進行預歸一化，使用 SwiGLU 激活函數 [Shazeer 2020]，以及旋轉位置嵌入（rotary positional embeddings）[Su et al 2024]。具體來說，我們在模型的每一層中使用二維 RoPE，參考了 [Lu et al 2023; Fang et al 2023] 的實現。我們未使用 AdaLN 技術 [Peebles & Xie 2023]，以保持結構與原始 LLM 相同。

類別條件圖像生成。類別嵌入從一組可學習的嵌入中索引得到 [Peebles & Xie 2023; Esser et al 2021]，并作為前填充 token 嵌入使用。

從該 token 嵌入開始，模型通過下一個 token 預測的方式生成圖像 token 序列，并在達到預設最大長度的位置停止。

文本條件圖像生成。為了將文本條件集成到自回歸模型中，我們使用 FLAN-T5 XL [Chung et al 2024] 作為文本編碼器，編碼后的文本特征通過一個額外的 MLP 投影 [Chen et al 2023b,c]，并作為自回歸模型中的前填充 token 嵌入使用。我們指出，這種設計并不是多模態基礎模型的最終方案，后者會在語言和視覺之間建立統一的詞匯表 [Lu et al 2023; Team et al 2023]。

我們將這一方向留待未來研究探索。

Classifier-free guidance（無分類器引導）。該方法起源于擴散模型社區，由 Ho & Salimans (2022) 提出，以其提升圖像質量和文本-圖像對齊效果而廣為人知。我們在模型中采用了這一策略。在訓練期間，條件輸入會被隨機丟棄，替換為一個“無條件嵌入” [Peebles & Xie 2023; Chen et al. 2023b]。在推理階段，對于每一個 token，其 logit ${?}_g$ 的計算方式為：
$${?}_g={?}_u+s({?}_c?{?}_u),$$
其中 ${?}_c$ 是有條件的 logit，${?}_u$ 是無條件的 logit，$s$ 是無分類器引導的放大系數。

值得注意的是，目前所討論的設計選擇在很大程度上受到已有工作的啟發，例如，圖像標記器借鑒自 [Rombach et al. 2022; Yu et al. 2021]，圖像生成方法源于 [Peebles & Xie 2023; Chen et al. 2023b; Esser et al. 2021]。這些技術在擴散模型中已被廣泛研究，而在自回歸模型中尚屬稀缺。我們的工作將這些先進設計集成應用于基于自回歸的視覺生成模型中。

2.4 Scale Up

---

我們的模型架構與 Llama 幾乎相同，這使我們能夠無縫地采用 LLM 社區中的優化技術 [Zhang & Sennrich 2019; Shazeer 2020; Su et al 2024] 和訓練策略 [Dao et al 2022; Rasley et al 2020; Shoeybi et al 2019]。如表1所示，我們在本工作中將模型規模擴展至 31 億參數。所有模型均基于 PyTorch 2 [Ansel et al 2024] 實現，并在 80GB 的 A100 GPU 上進行訓練。對于參數量小于 14 億的模型，我們直接使用 DDP 進行訓練；而對于更大的模型，我們采用 PyTorch FSDP [Zhao et al 2023] 以優化 GPU 內存使用。

2.5 Serving

---

自回歸模型一直存在推理速度慢的問題。隨著大語言模型的快速發展，LLM 社區提出了一系列先進的推理優化技術 [Kwon et al 2023; Chen et al 2023a; Dettmers 2022] 來提升推理效率。

與訓練過程類似，LLM 社區開發的推理加速方法同樣可以應用于我們的方法中以優化模型推理性能。我們驗證了當前最流行的 LLM 服務框架之一 vLLM [Kwon et al 2023] 在圖像生成任務中的有效性。如表7所示，與基線設置相比，我們的推理速度提升達到了 326% 至 414%。

表7：使用 vLLM 服務框架優化后的推理速度。推理時間針對一個批次共16張圖像（使用 classifier-free guidance 生成8張圖像）。所有模型的圖像分辨率均為 384×384。