Autoregressive Model Beats Diffusion: Llama for Scalable Image Generation

1 研究背景

基于自回归模型的 LLMs 取得了显著进展，部分研究开始探讨自回归模型在图像生成领域的应用，并引入图像标记化技术将连续图像转换为离散标记，从而实现图像标记的生成。

首先，图像被量化为离散标记，并通过光栅扫描顺序将这些标记序列化。然后，基于Transformer的自回归模型逐步预测下一个标记，直到生成完整的图像标记序列。最后，这些标记通过解码器还原为原始图像。

使用与VQGAN相同的三部分结构：编码器、量化器和解码器。

编码器
- 使用下采样比为 p 的 ConvNet
- 将图像像素 x 投影到特征图 f
量化器
- 包含一个codebook ，codebook中有 K 个可学习的向量。在量化过程中，编码器输出的每个特征向量会被映射到 codebook 中最接近的特征向量的代码索引。
解码器
- 使用下采样比为 p 的 ConvNet
- 在解码阶段，代码索引被重新映射回其对应的特征向量，解码器将这些特征向量转换回原始图像像素 x 。

codebook 的设计

量化是将连续值（图像像素）转换为离散值（ codebook 中的索引）的过程。这一过程是非可微分的，因此标准的梯度下降方法无法直接应用于量化步骤。
为了解决这个问题，使用了直通梯度估计器（Straight-Through Gradient Estimator），可以在反向传播时允许梯度流动。
- 公式为：是停止梯度操作，它确保在反向传播时，量化的影响不会被直接传播，而是保持输入特征 f 的影响。

codebook 学习是使编码器生成的特征向量 f 与 codebook 中的向量 z 之间的距离尽量接近。 codebook 学习损失

图像重建损失

Llama architecture.
- 模型架构主要基于 Llama
- 使用 RMSNorm
- SwiGLU 激活函数
- 旋转位置嵌入
  - 在模型的每一层都使用 2D RoPE。
- 不使用 AdaLN 来保持我们的结构与 LLM 相同。
Class-conditional image generation.
- 类嵌入是指将每个图像类别（如狗、猫、车等）映射为一个可学习的嵌入向量。这个嵌入向量是从一个嵌入矩阵中索引出来的。并用作预填充令牌嵌入。
- 从该令牌嵌入开始，模型通过下一个令牌预测方式生成图像令牌序列，并在预定义的最大长度的位置停止。
Text-conditional image generation.
- 使用 FLAN-T5 XL 作为文本编码器
- 编码的文本特征由额外的 MLP 投影
- MLP 处理后的文本特征会被用作自回归模型中的预填充 token 嵌入。
- 未来的研究可能会探索一种多模态基础模型,在语言和视觉之间建立一个统一的词汇表。
Classifier-free guidance.
- 无分类器引导是在扩散模型社区中开发的因其提高视觉质量和文本图像对齐而闻名。
- 在训练期间，条件被随机丢弃并被空条件嵌入取代。在推理中，对于每个标记，其 logit 由形成，其中是条件 logit，是无条件 logit，s 是无分类器指导的尺度。

本文的模型架构几乎与 Llama 相同，因此可以无缝地采用 LLM 社区中的优化技术和训练方法。

在这项工作中将模型大小扩展到 3.1B 参数。所有模型均使用 PyTorch 2 实现并在 80GB A100 GPU 上进行训练。对于训练参数低于1.4B的模型，我们直接使用DDP，否则，我们采用 PyTorch FSDP 优化 GPU 内存使用。

与训练类似，LLM社区开发的推理技术也可以用来优化本文的模型。作者在本文的图像生成方法上验证了 vLLM（最流行的 LLM 服务框架之一）的有效性。与基线设置相比，实现了 326% 和 414% 的加速。

增加图像令牌可以提升图像重建质量，但对图像生成质量的影响不显著。
当模型参数小于 1B 时，使用 256 个令牌（16×16 格式）能带来比 576 个令牌（24×24 格式）更好的图像生成效果。
然而，当令牌数量较少时，生成质量会受到限制。例如对于3B的model，256 令牌的 FID 为 3.06，而 576 令牌则能将 FID 降低到更低的值2.61。

不同规模的模型（B, L, XL, XXL, 3B）展示了在模型参数和训练轮次增加时，FID 值逐步降低，模型性能随之提升。但当模型扩展至 3B 后，训练轮次增加带来的进一步提升的效果变得有限。
这可能是由于数据集大小的限制（ImageNet 仅有约 100 万张图像），扩展数据集或增强数据增广可能进一步提升性能。

比较了该模型与其他模型如GAN、Diffusion Models、Masked Prediction Models等主流模型的性能。在多个评估指标上，模型表现出极具竞争力的性能，尤其是 3B 模型超越了一些著名的扩散模型，如 LDM 和 DiT。
该模型的优势来自于图像标记器的改进设计以及图像生成模型的良好可扩展性。