Autoregressive Model Beats Diffusion: Llama for Scalable Image Generation

目录

0. 摘要

1. 引言

2. 自回归模型在图像生成中的应用

2.1 概述

2.2 图像 tokenizer

2.3 自回归模型生成图像

2.4 规模扩展

2.5 服务

3. 实验

5. 结论

0. 摘要

我们介绍 LlamaGen，一个新的图像生成模型家族，该模型将大型语言模型（LLM）的原始 “下一 token 预测” 范式应用于视觉生成领域。这是对纯粹的自回归模型（例如 Llama）在没有视觉信号上的归纳偏差的情况下，通过适当扩展，是否能实现最先进的图像生成性能的肯定回答。我们重新审视了图像 tokenizer 的设计空间、图像生成模型的可扩展性属性及其训练数据质量。这次探索的成果包括：

• 一个图像 tokenizer，具有 16 倍的下采样比，在 ImageNet 基准测试中实现了 0.94 的 rFID 重建质量和 97% 的码本使用率。

• 一系列类别条件的图像生成模型，参数从 111M 到 3.1B 不等，在 ImageNet 256×256 基准测试中实现了 2.18 的 FID，超过了流行的扩散模型如 LDM 和 DiT。

• 一个具有 775M 参数的文本条件图像生成模型，通过在 LAION-COCO 和高美学质量图像上的两阶段训练，展示了在视觉质量和文本对齐方面的竞争性能。

• 我们验证了 LLM 服务框架在优化图像生成模型推理速度方面的有效性，实现了 326% - 414% 的加速。

我们发布了所有模型和代码，以促进开源社区的视觉生成和多模态基础模型的发展。

1. 引言

基于自回归模型，大型语言模型（LLMs）通过预测序列中的下一个 token 来生成文本。这种 “下一 token 预测” 范式在以类似人类对话的方式解决语言任务方面表现出了前所未有的能力，并展现了惊人的可扩展性，显示了通用人工智能模型的潜在路径。

自回归模型在大型语言模型上的可扩展性被验证后，先驱性工作尝试探索自回归模型在图像生成中的应用，例如，VQVAE ，DALL-E，Parti。他们引入图像 tokenizer，将连续图像转换为离散 token，并应用自回归模型以 “下一 token 预测” 的方式生成图像 token。他们在 2022 年前的同类中表现出了强大的性能。然而，他们的开源社区发展不完善，这在很大程度上限制了其进一步的改进。

同期，另一种图像生成方法，扩散模型快速发展。伴随着其开源社区，他们主导了直到今天的视觉生成领域。然而，扩散模型与自回归语言模型有着截然不同的范式，这对构建统一的语言和视觉模型提出了巨大的挑战。

在这项工作中，我们致力于进一步推进自回归模型在图像生成方面的应用：继续其研究方法并为开源社区做出贡献。回顾 2024 年前的图像生成文献，我们确定了现有高级模型的三个关键点：1）精心设计的图像压缩器，2）可扩展的图像生成模型，和 3）高质量的训练数据。受此启发，我们重新审视了图像 tokenizer（自回归模型的图像压缩器）的设计、图像生成模型的可扩展性属性及其训练数据的影响。

为了实现语言和视觉之间潜在的统一模型，我们的设计是减少对视觉信号的归纳偏差，并采用与 LLM 相同的架构。这与最近的一些工作所采用的以视觉为导向的设计不同。例如，MaskGIT 和 MAGVIT 采用了掩码图像建模策略，VAR 使用了分层多尺度属性。尽管它们在实现领先的图像生成性能方面取得了成功，甚至超过了扩散模型，但仍不清楚原始语言模型架构是否具备这种能力。相反，我们的工作揭示了应用与语言模型完全相同的 “下一 token 预测” 的纯粹自回归模型也能够实现最先进的图像生成性能。额外的好处是，我们可以利用 LLM 社区开发的技术来优化我们模型的训练配方和推理速度。

我们发布了所有模型和代码，以促进视觉生成和多模态基础模型的开源社区的发展。值得注意的是，我们发布的模型在性能上仍落后于基于扩散模型的最先进视觉生成模型 [Alpha-VLLM 2024; Esser等人 2024; Brooks等人 2024]。当未来有更多的训练数据和计算资源可用时，将探索更大规模的基于 AR 的视觉生成模型，例如超过 7B 参数的模型。

2. 自回归模型在图像生成中的应用

2.1 概述

首先，图像像素 x∈R^(H×W×3) 被图像 tokenizer 量化为 q∈Q^(h×w) 的离散 token，其中 h=H/p，w=W/p，p 是图像 tokenizer 的下采样率，q^(i,j) 是图像码本的索引。然后，这些图像 token 被重新排列为按照栅格扫描顺序的 h⋅w 个 token 序列，并用于训练基于 Transformer 的自回归模型。

在图像生成过程中，图像 token (q1,q2,...,q_(h⋅w)) 由自回归模型通过 “下一 token 预测” 的方式生成：

其中 c 是类别标签嵌入或文本嵌入。最后，这些图像 token 通过图像 tokenizer 解码器转换为图像像素。

2.2 图像 tokenizer

量化自编码器架构。我们使用与 VQGAN 相同的架构，即编码器-量化器-解码器。编码器和解码器是具有下采样比率 p 的卷积网络。量化器包含一个码本 Z∈R^(K×C)，其中有 K 个可学习的向量。编码器将图像像素 x 投射到特征图 f。量化过程将特征图中的每个向量 f^(i,j) 映射到码本中最接近的向量 z^(i,j) 的编码索引 q^(i,j)。在解码过程中，编码索引 q^(i,j) 被重新映射到特征向量 z^(i,j)，解码器将这些特征向量转换回图像像素 ^x。

码本对图像 tokenization 性能有重大影响。根据 [Yu等人 2021]，我们对码本向量使用 ℓ_2 正则化、低码本向量维度 C 和大码本大小 K。这些设计显著提高了重建质量和码本的使用率。更多细节将在实验中讨论。

训练损失。由于量化是不可微操作，使用直通梯度估计器（straight-through gradient
estimator）来保持从解码器到编码器的梯度：z = sg[z−f] + f，其中 sg[⋅] 是停止梯度操作。对于码本学习：L_VQ = || sg[f] − z ||^2_2 + β·|| f − sg[z] ||^2_2，其中第二项是 commitment loss，迫使从编码器提取的特征向量接近码本向量，β 是 权重。为了简化，我们在码本学习中不添加熵损失 [Yu等人 2023a; Chang等人 2022]。

对于图像重建训练：

其中 ℓ_2 是图像像素的重建损失，L_P(⋅) 是来自 LPIPS 的感知损失，LG(⋅) 是与图像 tokenizer 同时训练的 PatchGAN [Isola等人 2017] 鉴别器的对抗损失，λ_G 是对抗损失权重。

2.3 自回归模型生成图像

Llama 架构。我们的模型架构主要基于 Llama，应用了 RMSNorm 进行预归一化（pre-normalization），SwiGLU 激活函数，以及旋转位置嵌入。具体来说，我们在每一层模型中使用 2D RoPE，按照 [Lu et al. 2023; Fang et al. 2023] 的实现方法。我们不使用 AdaLN 技术，以保持我们的结构与 LLM 一致。

类别条件图像生成。类别嵌入从一组可学习的嵌入中索引，并用作预填充 token 嵌入。从这个 token 嵌入开始，模型通过下一 token 预测的方式生成图像 token 序列，并在预定义的最大长度位置停止。

文本条件图像生成。为了将文本条件整合到自回归模型中，我们使用 FLAN-T5 XL 作为文本编码器，编码后的文本特征通过一个额外的 MLP 投影，并用作自回归模型中的预填充 token 嵌入。我们注意到这种设计并不是多模态基础模型的最终设计，其中语言和视觉之间建立了统一的词汇表。我们将其留待未来研究。

无分类器指导。在扩散模型社区中开发的无分类器指导以其提高视觉质量和文本图像对齐而闻名。我们在我们的模型中采用了这一方法。在训练过程中，条件随机丢弃，并被一个空的无条件嵌入取代。在推理时，对于每个 token，其 logit ℓ_g 由 ℓ_g = ℓ_u + s(ℓ_c − ℓ+u) 形成，其中 ℓ_c 是条件 logit，ℓ_u 是无条件 logit，s 是无分类器指导的缩放比例。

值得注意的是，迄今为止讨论的所有设计选择主要受到先前工作的启发，例如，图像 tokenizer 借鉴自 [Rombach et al. 2022; Yu et al. 2021]，图像生成来自 [Peebles & Xie 2023; Chen et al. 2023b; Esser et al. 2021]。这些技术的大部分在扩散模型中研究得很好，但在自回归模型中研究得较少。我们的工作将这些先进的设计集体适应于基于自回归的视觉生成模型。

2.4 规模扩展

我们的模型架构与 Llama 几乎相同，这使我们能够无缝采用 LLM 社区中的优化技术和训练方案。如表 1 所示，我们在这项工作中将模型规模扩展到 3.1B 参数。所有模型都使用 PyTorch 2 实现，并在 80GB A100 GPU 上进行训练。对于参数少于 1.4B 的模型，我们直接使用 DDP，否则，我们采用 PyTorch FSDP 来优化 GPU 内存使用。

2.5 服务

自回归模型一直受困于其低推理速度。随着大规模语言模型的快速发展，LLM 社区提出了高级推理技术来优化推理速度。与训练类似，LLM 社区开发的推理技术也可以用来优化我们的模型。我们验证了 vLLM 这一最流行的 LLM 服务框架之一在我们的图像生成方法上的有效性。如表 7 所示，与基线设置相比，实现了 326% - 414% 的加速。

3. 实验

5. 结论

在这项工作中，我们深入研究了用于可扩展图像生成的普通自回归模型。通过重新审视其图像 tokenizer、图像生成模型和训练数据，我们的类别条件模型优于流行的扩散模型，而我们的文本条件模型在视觉质量和文本对齐方面表现出竞争力。 

论文地址：https://arxiv.org/abs/2406.06525

项目页面：https://github.com/FoundationVision/LlamaGen 

进 Q 交流群：922230617 或加 VX：CV_EDPJ 进 V 交流群
加 VX 群请备注学校 / 单位 + 研究方向