首次超越扩散模型和非自回归Transformer模型！字节开源RAR：自回归生成最新SOTA！

文章链接：https://arxiv.org/pdf/2411.00776
项目链接：https://yucornetto.github.io/projects/rar.html
代码&模型链接：https://github.com/bytedance/1d-tokenizer

亮点直击

• RAR（随机排列自回归训练策略），这是一种改进的训练策略，使得标准的自回归图像生成器能够实现SOTA性能。
• 引入双向上下文学习：RAR通过最大化所有可能的分解顺序的期望似然值，打破了自回归模型在视觉任务中的单向上下文限制，使模型能够在图像生成中更有效地利用双向上下文信息。
• 保持与语言建模框架的兼容性：RAR在提升图像生成性能的同时，保留了自回归建模的核心结构，它与大语言模型（LLM）的优化技术（如KV-cache）完全兼容，相比于MAR-H或MaskBit，采样速度显著更快，同时保持了更好的性能，便于在多模态统一模型中应用。
• 创新的退火训练策略：训练初期将输入序列随机排列，随着训练逐步回归至标准光栅顺序。这一过程使模型在各类上下文排列中均能获得优化，提升生成质量。
• 显著的性能提升：在 ImageNet-256 基准测试中，RAR实现了1.48的 FID 分数，显著超越了先前的自回归图像生成器，显示了其在图像生成任务中的突破性改进。

总结速览

解决的问题：
RAR（随机自回归建模）旨在提升图像生成任务的表现，同时保持与语言建模框架的完全兼容性。

提出的方案：
RAR采用了一种简单的方法，通过标准的自回归训练过程并结合下一个 token 预测目标，将输入序列随机打乱到不同的分解顺序。在训练过程中，使用概率 r 将输入序列的排列顺序随机化，其中 r 从 1 开始并逐渐线性衰减至 0，以让模型学习所有分解顺序的期望似然值。

应用的技术：
RAR在训练中应用了一种退火策略，使模型能够更好地利用双向上下文信息来建模，同时仍然保持自回归建模框架的完整性，从而实现了语言建模的完全兼容性。

达到的效果：
在 ImageNet-256 基准测试中，RAR 获得了 1.48 的 FID 分数，超越了之前最先进的自回归图像生成器，并优于领先的基于扩散和掩码Transformer的方法。

方法

背景

RAR: 随机自回归建模

视觉信号天然存在双向关联，因此全局上下文建模在视觉任务中至关重要。然而，传统的自回归模型依赖因果注意力掩码，仅允许 token 序列呈单向依赖，这与视觉数据的双向关联性不符。已有研究表明，对于视觉模态，双向注意力显著优于因果注意力。

此外，图像 token 排列成因果序列时没有统一的“正确”方式，常用的光栅扫描顺序虽有效果，但在训练过程中引入了偏差，因为每个 token 仅依赖于扫描顺序中的前序 token，限制了模型学习来自其他方向的依赖关系。

为了解决这些问题，本文提出了一种随机自回归建模方法，将双向上下文的优化目标纳入自回归建模。

与 BERT 风格 或 MaskGIT 风格方法不同，本文的方法采用了置换目标的方法，在所有可能的分解顺序上以自回归方式训练模型。这使得模型在保持自回归框架的同时，能够在期望上收集双向上下文信息。公式表示如下：

虽然方法简单，但这种修改显著提升了图像生成性能，突显了双向上下文在提升自回归图像生成器能力方面的重要性。与自然语言处理 (NLP) 中的自回归训练观察结果一致。

讨论：尽管置换目标允许在自回归框架中实现双向上下文学习，但在生成过程中完全捕获“全局上下文”仍具挑战性。这是因为在生成过程中，总会有一些 tokens 在其他 tokens 之前生成，无法完全利用全局上下文。这种限制不仅存在于自回归方法中，也存在于非自回归模型中。重新采样或精炼等技术可能有助于确保每个 token 都能在充分的上下文下生成。然而，这些设计可能会增加系统的复杂性，因此探索这些解决方案超出了本文的范围，留待未来研究。

随机退火：虽然通过置换的随机自回归训练使模型能够在单向框架中捕获双向上下文，但它可能会因两个主要因素而引入次优的视觉生成行为：(1) 可能的置换数量极多，可能导致模型过于关注如何处理不同的置换顺序，而不是提高生成质量。例如，对于长度为 256 的 token 序列，可能的置换数为 (256! > 10^{506})，这会让模型不堪重负，降低训练效率。(2) 尽管图像可以按任意顺序处理，某些扫描顺序往往优于其他顺序。例如，[22] 评估了六种不同的扫描顺序（行优先、向内螺旋、向外螺旋、Z 曲线、子采样和交替顺序），并发现行优先（即栅格顺序）始终表现最佳，这使其成为视觉生成中最广泛使用的顺序。

为了解决这些问题，本文提出了随机退火策略，以平衡置换的随机性与行优先顺序的已知效果。此方法引入了一个控制随机置换和行优先顺序使用概率的单一参数 r。在训练开始时，r = 1，意味着模型完全使用随机置换。随着训练的进行，r 线性衰减至 0，逐步将模型转换为行优先顺序。具体而言，定义了r 的训练调度，由两个超参数 ​​start​​​ 和 ​​end​​ 控制，分别表示 r 开始退火和退火结束的训练轮次。公式如下：

其中，​​epoch​​​ 表示当前的训练轮次。在实验中，我们将对超参数 ​​start​​​ 和 ​​end​​ 进行消融实验。该调度策略允许模型在初期探索多种随机置换，以更好地学习双向表示，最终收敛到更有效的行优先扫描顺序，从而提升视觉生成质量，类似于其他典型的自回归方法。值得注意的是，该策略不仅提升了生成性能，还保持了与先前工作中使用的标准扫描顺序的兼容性。

实验结果

本节首先介绍方法的实现细节。接这展示了关于关键设计选择的消融研究。再讨论了主要结果，最后还包括了缩放研究和可视化内容。

实现细节

本文在语言建模自回归框架的基础上实现了RAR方法，做了最小的改动。

VQ分词器：采用了与先前工作 [10, 22] 相似的VQ分词器，将输入图像转换为离散的token 序列。我们使用的是基于CNN的MaskGIT-VQGAN [10] 分词器，使用在ImageNet上训练的官方权重。该分词器将256 × 256的图像转化为256个离散的token （下采样因子为16），并且字典大小（即词汇表大小）为1024。

自回归Transformer：本文使用不同配置的视觉Transformer（ViT），包括RAR-S（133M）、RAR-B（261M）、RAR-L（461M）、RAR-XL（955M）和RAR-XXL（1499M）。对于所有这些模型变体，我们在自注意力模块中应用了因果注意力掩码，并使用QK层归一化来稳定大规模模型的训练。为了加速实验，在消融研究中使用了普通的ViT，而在最终模型中增强了AdaLN 。架构配置和模型大小可以在下表1中找到。

位置嵌入：本文为ViT中的原始位置嵌入和目标感知位置嵌入使用了可学习的嵌入。值得注意的是，由于我们的模型在训练结束后会退火到基于栅格顺序的自回归图像生成，最终这两种位置嵌入可以合并为一个，使得最终模型与传统的自回归图像生成器相同。

数据集：研究者们在ImageNet-1K训练集上训练我们的模型，该数据集包含128,1167张图像，涵盖1000个物体类别。我们使用MaskGIT-VQGAN分词器对整个训练集进行预分词，以加速训练。对于消融研究，仅使用中心裁剪和水平翻转数据增强进行预分词，而对于最终模型，使用了十裁剪变换来增强数据集的多样性。

训练协议：研究者们所有模型变体使用相同的训练超参数。模型使用批量大小2048训练400个epoch（250k步）。在前100个epoch（热身阶段）内，学习率从0线性增加到4 × 10⁻⁴，然后按照余弦衰减计划逐渐衰减至1 × 10⁻⁵。使用AdamW优化器，其中beta1为0.9，beta2为0.96，权重衰减为0.03。我们对梯度进行了裁剪，最大梯度范数为1.0。在训练过程中，类别条件会以0.1的概率被丢弃。消融研究和主要结果中所有RAR模型变体的训练设置保持一致。

采样协议：本文使用[18]的评估代码对50,000张图像进行FID计算。我们不使用任何基于top-k或top-p的过滤技术。还遵循先前的工作使用无分类器指导。在消融研究中，使用更简单的线性指导调度，而在最终模型中使用改进的幂余弦指导调度。

消融研究

本文研究了RAR的不同配置，包括随机退火策略和RAR最终收敛的扫描顺序。

随机退火策略：在下表2中，比较了不同的随机退火策略。采用了线性衰减的调度，并通过改变超参数 ​​start​​​ 和 ​​end​​​ 来研究何时应该开始和结束随机化退火，具体定义见公式(5)。对于持续400个epoch的训练，我们枚举了每100个epoch的所有可能组合。例如，当 ​​start = 200​​​ 和 ​​end = 300​​​ 时，模型在前200个epoch采用随机排列，在后100个epoch采用栅格顺序。在第200到300个epoch之间，模型通过以概率 ​​r​​​ 进行随机排列，或者以概率 ​​1−r​​​ 进行栅格顺序训练，其中 ​​r​​​ 按照公式(5)计算。值得注意的是，当 ​​start = end = 0​​​ 时，模型仅使用栅格顺序进行训练，即标准的自回归训练；当 ​​start = end = 400​​​ 时，模型始终使用随机排列的输入序列进行训练。两种情况都是提出的随机退火方法的重要基准，分别达到了FID得分3.08和3.01。令人感兴趣的是，我们观察到所有其他变体都比这两个基准取得了显著的改进。例如，简单地将前100个epoch的栅格顺序替换为随机排列（即，​​start = 100​​​ 和 ​​end = 100​​），就将FID得分提高到了2.48，提升了0.6。此外，模型倾向于保留一些初期的epoch进行纯随机排列训练，并且在最后一些epoch更好地适应栅格顺序，这通常比其他变体表现更好。所有结果表明，通过引入带有排列目标的随机化自回归训练，有助于自回归视觉生成器的性能，并提升FID得分，这得益于改进的双向表示学习过程。

此外，在所有变体中，发现当 ​​start = 200​​​ 和 ​​end = 300​​ 时表现最佳，将基准（纯栅格顺序）的FID从3.08提高到2.18。该策略将稍多的计算分配给随机排列顺序的训练，并将最后100个epoch专注于纯栅格顺序。因此，我们默认采用这种退火策略用于所有RAR模型。

不同的扫描顺序（除了栅格顺序）：尽管行主序（即栅格扫描）一直是视觉生成中事实上的扫描顺序，但缺乏系统的研究来比较它与其他扫描顺序的优劣。我们注意到，四年前的工作 [22] 进行了类似的研究。然而，考虑到近年来生成模型取得的显著进展，值得重新审视这一结论。具体来说，我们考虑了6种不同的扫描顺序（行主序、螺旋内、螺旋外、Z曲线、子采样和替代扫描顺序），这些扫描顺序是RAR可能最终收敛的目标。本文没有像那样报告训练损失和验证损失作为对比指标，而是直接评估它们的生成性能。结果总结在下表3中。有趣的是，我们观察到所有变体的得分都相当不错，这表明RAR能够处理不同的扫描顺序。考虑到行主序（栅格扫描）仍然在其他扫描顺序中表现出优势，我们因此为所有最终RAR模型使用栅格扫描顺序。

主要结果

本文报告了RAR与最先进的图像生成器在ImageNet-1K 256×256基准测试上的结果。

如下表4所示，RAR相较于之前的AR图像生成器表现出了显著更好的性能。具体来说，最紧凑的RAR-B（仅有261M参数）就达到了FID得分1.95，已经显著超越了当前最先进的AR图像生成器LlamaGen-3B-384（3.1B，FID 2.18，裁剪尺寸384）和 Open-MAGVIT2-XL（1.5B，FID 2.33），并且分别减少了91%和81%的模型参数。它还超越了广泛使用的扩散模型，例如DiT-XL/2（FID 1.95 vs. 2.27）和SiT-XL（FID 1.95 vs. 2.06），并且仅使用了相对于这些模型的39%的参数。

在表4中，进一步探讨了不同模型尺寸下的RAR（从261M到1.5B），我们观察到RAR在不同尺寸下具有强大的可扩展性，并且随着模型尺寸的增大，性能不断提升。特别地，最大的变体RAR-XXL在ImageNet基准测试上创下了新的最先进结果，FID得分为1.48。与其他两种近期方法VAR和MAR相比，这两种方法都尝试通过改进AR公式来提升视觉生成质量，RAR不仅在性能上表现更优（RAR的FID为1.48，而VAR为1.73，MAR为1.55），而且保持了整个框架与语言建模的兼容性，因此更适合将成熟的优化和加速技术应用于大型语言模型，从而推动视觉生成的发展。

此外，RAR在不同框架中的表现超越了最先进的视觉生成器。它在对比领先的自回归模型、扩散模型和掩蔽Transformer模型时，表现得更好，分别超越了LlamaGen-3B-384、MDTv2-XL/2和 MaskBit（RAR的FID为1.48，相比之下LlamaGen为2.18，MDTv2为1.58，MaskBit为1.52）。据所知，这是首次语言建模风格的自回归视觉生成器超越最先进的扩散模型和掩蔽Transformer模型。

采样速度：自回归方法的一个关键优势是能够利用LLM中已建立的优化技术，如KV缓存。在表5中，我们将RAR与其他类型的生成模型进行采样速度比较（以每秒生成图像数为衡量标准），包括扩散模型、masked transformer器、VAR和 MAR。其中，自回归模型（RAR）和VAR模型（VAR-d30）与KV缓存优化兼容，因此在生成速度上显著优于其他方法。如表5所示，RAR不仅在FID得分上达到了最先进的水平，同时在生成速度上也大大超越了其他方法。例如，在FID得分约为1.5时，MaskBit和 MAR-H的生成速度分别为每秒0.7和0.3张图像。相比之下，RAR-XL不仅获得了更好的FID得分，还能每秒生成8.3张高质量视觉样本——比MaskBit快11.9倍，比MAR-H快27.7倍。最大的RAR变体RAR-XXL进一步提高了FID得分，同时保持了显著的速度优势，速度是MaskBit的9.1倍，是MAR-H的21.3倍。此外，RAR可能进一步受益于LLM优化技术，例如vLLM，这一点与其他AR方法一致。

扩展性行为

本文研究了RAR的扩展性行为。具体来说，我们绘制了训练损失曲线和FID得分曲线（有无无分类器引导的情况）如下图4所示。如图所示，我们观察到RAR在不同模型尺寸下均表现出良好的扩展性，较大的模型尺寸在训练损失和FID得分上持续表现出较好的性能，无论是否使用无分类器引导增强。我们注意到，由于RAR保持了AR公式和框架的完整性，它也继承了AR方法的可扩展性。

可视化

在下图5中可视化了不同RAR变体生成的样本，结果表明RAR能够生成高质量、具有高度保真度和多样性的样本。更多可视化结果见附录。

结论

本文提出了一种简单而有效的策略，以增强与语言建模兼容的自回归图像生成器的视觉生成质量。通过采用随机化排列目标，本文的方法在保持自回归结构的同时，改善了双向上下文学习。因此，所提出的RAR模型不仅超越了以前的最先进自回归图像生成模型，还超过了领先的非自回归Transformer和扩散模型。希望本研究有助于推动自回归Transformer朝着视觉理解与生成统一框架的方向发展。

本文转自AI生成未来 ，作者：AI生成未来

原文链接:​​https://mp.weixin.qq.com/s/pIG5FJc1maOtz96S4CiJIA​​