自回归模型优于扩散：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\cdot w$长度的序列，按光栅扫描顺序排列，并用于训练基于Transformer [Vaswani et al. 2017] 的自回归模型。

在图像生成过程中，图像标记$(q_1,q_2,\dots ,q_{h\cdot w})$由自回归模型生成[Radford et al. 2018, 2019; Brown et al. 2020; Touvron et al. 2023a]，方式是通过下一个token预测：
$$\prod _{t=1}^{h\cdot 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]，我们对码本向量使用${\ell }_2$归一化，采用较低的码本维度$C$和较大的码本大小$K$。这些设计显著提升了重建质量和码本使用率。更多细节将在实验中讨论。

训练损失。由于量化是不可导操作，因此使用直通梯度估计器 [Bengio et al. 2013] 来保留从解码器到编码器的梯度。定义：
$$z=\text{sg}[z-f]+f,$$
其中$\text{sg}[\cdot ]$是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}={\ell }_2(x,\hat{x})+{\mathcal{L}}_P(x,\hat{x})+{\lambda }_G{\mathcal{L}}_G(\hat{x}),$$
其中${\ell }_2$是图像像素的重建损失，${\mathcal{L}}_P(\cdot )$ 是来自LPIPS [Zhang et al. 2018] 的感知损失，${\mathcal{L}}_G(\cdot )$ 是使用与图像标记器同时训练的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 ${\ell }_g$ 的计算方式为：
$${\ell }_g={\ell }_u+s({\ell }_c-{\ell }_u),$$
其中 ${\ell }_c$ 是有条件的 logit，${\ell }_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。