江大白 | 何凯明入职 MIT，首次带队提出Diffusion Loss，扩散模型思想提升生成速度和效果！

本文来源公众号“江大白”，仅用于学术分享，侵权删，干货满满。

原文链接：何凯明入职 MIT，首次带队提出Diffusion Loss，扩散模型思想提升生成速度和效果！

导读

在图像生成领域中，作者观察到向量量化标记并不是自回归建模的必要条件，并提出通过在连续值域上，操作的扩散过程来对每个标记的概率分布，进行建模。改进后的图像生成器在生成速度与效果上，都取得了巨大的提升。

传统观点认为，用于图像生成的自回归模型通常伴随着向量量化标记。作者观察到，尽管离散值空间可以促进表示分类分布，但这并非自回归建模的必要条件。

在这项工作中，作者提出使用扩散过程来建模每个标记的概率分布，这使得作者能够在连续值空间中应用自回归模型。作者不是使用分类交叉熵损失，而是定义了一个扩散损失函数来建模每个标记的概率。

这种方法消除了对离散值标记器的需求。作者在广泛的情况下评估了其有效性，包括标准自回归模型和广义 Mask 自回归（MAR）变体。

通过移除向量量化，作者的图像生成器在享受序列建模速度优势的同时，取得了强大的成果。

作者希望这项工作将激发在其他连续值领域和应用中使用自回归生成的兴趣。

论文链接：https://arxiv.org/pdf/2406.11838

1 Introduction

自回归模型目前是自然语言处理中生成模型的实际解决方案。这些模型基于前面的词作为输入来预测序列中的下一个词或标记。由于语言的离散性质，这些模型的输入和输出处于分类的、离散值的空间。这种普遍的方法导致人们普遍认为自回归模型与离散表示固有地联系在一起。

因此，将自回归模型推广到连续值域（尤其是图像生成）的研究一直集中在数据离散化上[6; 13; 40]。通常采用的策略是训练一个针对图像的离散值分词器，这涉及到通过向量量化（VQ）获得的有限词汇[51; 41]。

然后自回归模型在离散值标记空间上操作，类似于它们的语言对应物。

在这项工作中，作者旨在回答以下问题：“自回归模型与向量量化表示结合是必要的吗？”作者注意到自回归的特性，即“基于前面的标记预测下一个标记”，与值是离散的还是连续的无关。需要的是对每个标记的概率分布进行建模，这可以通过损失函数来衡量，并从中抽取样本。离散值表示可以通过分类分布方便地建模，但这在概念上并非必要。如果提出每个标记概率分布的其他模型，自回归模型可以在没有向量量化的情况下处理。

基于这一观察，作者提出通过在连续值域上操作的扩散过程来对每个标记的概率分布进行建模。作者的方法论利用了扩散模型[45; 24; 33; 10]的原理来表示任意的概率分布。具体来说，作者的方法自回归地为每个标记预测一个向量z，该向量作为去噪网络（如一个小型MLP）的条件。去噪扩散过程使作者能够表示输出x的潜在分布p(x|z)（图1）。这个小型的去噪网络与自回归模型一起训练，以连续值标记作为输入和目标。从概念上讲，这个应用于每个标记的小型预测头就像一个用于衡量z质量的损失函数。作者将这个损失函数称为“扩散损失”。

作者的方法消除了对离散值分词器的需求。向量量化分词器难以训练，且对梯度近似策略敏感。它们的重建质量通常比连续值对应物差[42]。作者的方法允许自回归模型享受高质量、非量化分词器的优点。

为了扩大范围，作者进一步将标准的自回归（AR）模型[13]和 Mask 生成模型[4; 29]统一到一个广义的自回归框架中（图3）。从概念上讲， Mask 生成模型以随机顺序同时预测多个输出标记，同时仍保持基于已知标记预测下一个标记的自回归性质。这导致了一个可以与扩散损失无缝使用的“ Mask 自回归”（MAR）模型。

作者通过实验证明了扩散损失在各种情况下的有效性，包括AR和MAR模型。它消除了对向量量化分词器的需求，并一致提高了生成质量。作者的损失函数可以灵活地应用于不同类型的分词器。此外，作者的方法享有序列模型快速速度的优势。作者的带扩散损失的MAR模型在ImageNet 256x256上的生成速度小于0.3秒/张图像，同时达到强大的FID小于2.0。作者最好的模型可以达到1.55 FID。

作者的方法的有效性揭示了一个在很大程度上尚未探索的图像生成领域：通过自回归建模标记的“相互依赖”，同时通过扩散对每个标记的分布进行建模。这与典型的潜在扩散模型[42; 37]形成对比，在后者中，扩散过程建模了所有标记的联合分布。鉴于作者方法的效率、速度和灵活性，作者希望扩散损失能推进自回归图像生成，并在未来的研究中推广到其他领域。

2 Related Work

图像生成的序列模型。在自回归图像模型方面的开创性工作是在像素序列上进行的。自回归可以通过RNNs [50]、CNNs [49; 7] 来实现，最近最流行的是Transformers [36; 6]。受到语言模型的启发，另一系列工作 [51; 41; 13; 40] 将图像建模为离散值标记。自回归 [13; 40] 和 Mask 生成模型 [4; 29] 可以在离散值标记空间上操作。但是离散标记器难以训练，这最近引起了特别的关注 [27; 54; 32]。

与作者的工作相关的是，最近关于GIVT [48] 的工作也专注于序列模型中的连续值标记。GIVT 和作者的工作都揭示了这一方向的重要性和潜力。在GIVT中，标记分布由高斯混合模型表示。它使用预定义的混合数量，这可能限制了它可以表示的分布类型。相比之下，作者的方法利用扩散过程在建模任意分布方面的有效性。

表示学习的扩散过程。去噪扩散过程已经被探索作为视觉自监督学习的准则。例如，DiffMAE [53] 用去噪扩散解码器替换了原始 MAE [21] 中的 L2 损失；DARL [30] 使用去噪扩散块解码器训练自回归模型。这些努力主要集中在表示学习上，而不是图像生成。在它们的场景中，生成 多样化 的图像并非目标；这些方法尚未展示从零开始生成新图像的能力。策略学习的扩散。作者的工作与机器人领域的扩散策略 [8] 在概念上是相关的。在这些场景中，_采取行动_ 的分布被制定为机器人观察到的去噪过程，可以是像素或潜在表示 [8; 34]。在图像生成中，作者可以将生成一个标记视为“行动”。尽管存在这种概念上的联系，但在机器人领域生成的样本多样性不如图像生成核心考量。

注意：公式部分按照您的要求保持原始输出，未翻译。

3 Method

总之，作者的图像生成方法是在标记化的潜在空间上操作的序列模型[6; 13; 40]。但与以往基于向量量化标记器的方法（例如，VQ-VAE的变体[51; 13]）不同，作者旨在使用连续值标记器（例如，[42]）。作者提出了扩散损失，这使得序列模型能够与连续值标记兼容。

3.1 Rethinking Discrete-Valued Tokens

这种分析表明，离散值标记对自回归模型来说_并非_必要。相反，建模一个分布是本质上的要求。离散值标记空间意味着一个分类分布，其损失函数和采样器容易定义。作者实际上需要的是用于分布建模的损失函数及其相应的采样器。

3.2 Diffusion Loss

去噪扩散模型[24]为建模任意分布提供了一个有效的框架。但与将扩散模型用于表示所有像素或所有 Token 的联合分布的常见用法不同，在作者的案例中，扩散模型用于表示每个 Token 的分布。

3.3 Diffusion Loss for Autoregressive Models

3.4 Unifying Autoregressive and Masked Generative Models

作者展示了诸如MaskGIT [4]和MAGE [29]之类的 Mask 生成模型可以在自回归的广泛概念下进行泛化，即下一个标记的预测。

双向注意力可以执行自回归。自回归的概念与网络架构正交：自回归可以通过RNNs [50]、CNNs [49; 7]和Transformers [38; 36; 6]来完成。当使用Transformers时，尽管自回归模型通常通过_因果_注意力来实现，但作者展示了它们也可以通过_双向_注意力来完成。请见图2。需要注意的是，自回归的目标是给定前一个标记后_预测下一个标记_；它并不限制前一个标记如何与下一个标记进行通信。

作者可以采用如Masked Autoencoder (MAE) [21]中所示的双向注意力实现。见图2(b)。具体来说，作者首先在已知标记上应用MAE风格的编码器1（带有位置嵌入[52]）。然后作者将编码后的序列与 Mask 标记（再次加上位置嵌入）连接起来，并用MAE风格的解码器映射这个序列。Mask 标记上的位置嵌入可以让解码器知道需要预测哪些位置。与因果注意力不同，这里的损失只在未知标记上计算[21]。利用MAE风格的技巧，作者允许_所有_已知标记相互看见，也允许所有未知标记看见所有已知标记。这种_全注意力_比因果注意力在标记之间的通信引入了更好的效果。在推理时，作者可以使用这种双向公式生成标记（每步一个或多个），这是一种自回归形式。作为一种妥协，作者无法使用因果注意力的键值（kv）缓存[44]来加速推理。但正如作者可以同时生成多个标记，作者可以减少生成步骤以加速推理。标记之间的全注意力可以显著提高质量，并提供更好的速度/精度权衡。

随机顺序的自回归模型。为了与 Mask 生成模型[4, 29]相联系，作者考虑了一种随机顺序下的自回归变体。模型接收到一个随机排列的序列。这种随机排列对于每个样本都是不同的。见图3(b)。在这种情况下，下一个待预测标记的位置需要模型能够访问。作者采用了与MAE[21]类似的策略：作者在解码层中添加位置嵌入（与未Shuffle的位置相对应），这可以告诉模型预测哪些位置。这种策略适用于因果和双向版本。

如图3(b)(c)所示，随机顺序自回归的行为像是一种特殊的 Mask 生成形式，一次生成一个标记。作者以下进行详细阐述。

3.5 Diffusion Loss

3.6 Autoregressive and Masked Autoregressive Image Generation

分词器。作者使用了LDM [42]提供的公开分词器。作者的实验将涉及它们的VQ-16和KL-16版本[42]。VQ-16是一个VQ-GAN[13]，即带有GAN损失[15]和感知损失[56]的VQ-VAE[51]；KL-16是它的对应版本，通过Kullback-Leibler（KL）散度进行正则化，不进行向量量化。16表示分词器的步长。

Transformer。作者的架构遵循ViT[11]中的Transformer[52]实现。给定一个来自分词器的标记序列，作者添加位置嵌入[52]并在序列开头附加类别标记[cls]；然后通过Transformer处理该序列。默认情况下，作者的Transformer有32个块和1024的宽度，作者称之为大型（-L）（约400M个参数）。

**自回归 Baseline **。因果注意力遵循GPT[38]的常见实践（图2(a)）。输入序列通过一个标记（此处为[cls]）进行移位。三角 Mask [52]应用于注意力矩阵。在推理时，应用温度（τ）采样。作者使用kv-cache[44]进行高效推理。

** Mask 自回归模型**。使用双向注意力（图2(b)），作者可以根据任意数量的已知标记预测任意数量的未知标记。在训练时，作者在[0.7, 1.0]范围内随机采样 Mask 比例[21, 4, 29]：例如，0.7意味着70%的标记是未知的。由于采样的序列可能非常短，作者在编码器序列的开始处始终填充64个[cls]标记，这提高了作者编码的稳定性和容量。如图2所示，在解码器中引入 Mask 标记[m]，并添加位置嵌入。为了简单起见，与[21]不同，作者让编码器和解码器具有相同的尺寸：每个都有所有块的一半（例如，在MAR-L中为16）。

在推理时，MAR执行“下一组标记预测”。它按照余弦计划逐步将 Mask 比例从1.0减少到0[4, 29]。默认情况下，作者在该计划中使用64步。应用温度（τ）采样。与[4, 29]不同，MAR始终使用完全随机顺序。

4 Experiments

作者在ImageNet [9] 数据集上进行了实验，分辨率为256×256。作者评估了FID [22] 和IS [43]，并根据常见做法[10]提供了精确度和召回率作为参考。作者遵循了[10]提供的评估套件。

4.1 Properties of Diffusion Loss

扩散损失与交叉熵损失的比较。作者首先比较了使用扩散损失的连续值标记与使用交叉熵损失的标准离散值标记（表1）。为了公平比较，两个分词器（"VQ-16" 和 "KL-16"）均从LDM代码库[42]下载。这些分词器被广泛使用（例如[13, 42, 37]）。

比较在AR/MAR的四种变体中进行。如表1所示，在所有情况下，扩散损失一致优于交叉熵损失。特别是，在MAR中（例如，默认设置），使用扩散损失可以相对减少约50%-60%的FID。这是因为连续值的KL-16比VQ-16（下面表2中讨论）具有更小的压缩损失，也因为扩散过程比分类过程更有效地建模分布。

在以下消融研究中，除非另有指定，作者遵循表1中“默认”MAR设置。

扩散损失的灵活性。扩散损失的一个显著优势是它适用于各种分词器。作者在表2中比较了几个公开可用的分词器。

即使给定VQ分词器，扩散损失也可以轻松使用。作者简单地将VQ层之前的连续值潜在标记视为标记。这个变体使作者得到了7.82 FID（无CFG），与使用相同VQ分词器的交叉熵损失的8.79 FID（表1）相比，表现良好。这表明扩散在建模分布方面具有更好的能力。

这个变体还使作者能够在使用相同损失的情况下比较VQ-16和KL-16分词器。如表2所示，VQ-16的重建FID（rFID）比KL-16差得多，这也导致了生成FID（例如，表2中的7.82 vs 3.50）差得多。

有趣的是，扩散损失还使作者能够使用步长不匹配的分词器。在表2中，作者研究了一个步长为8，输出序列长度为32×32的KL-8分词器。在没有增加生成器序列长度的情况下，作者将2×2标记组合成一个新的标记。尽管步长不匹配，作者仍然能够获得不错的结果，例如，KL-8给作者带来了2.05 FID，而KL-16是1.98 FID。此外，这个特性允许作者研究其他分词器，例如一致性解码器[35]，这是一种不同架构/步长的非VQ分词器，专为不同目标设计。

为了全面性，作者还使用[42]的代码在ImageNet上训练了一个KL-16分词器，注意到[42]中的原始KL-16是在OpenImages[28]上训练的。比较在表2的最后一行。作者将在以下探索中使用这个分词器。

扩散损失中的去噪MLP。作者在表3中研究了去噪MLP。即使是非常小的MLP（例如，2M）也能带来有竞争力的结果。如预期的那样，增加MLP的宽度有助于提高生成质量；作者还探索了增加深度并观察到类似情况。请注意，作者的默认MLP大小（1024宽度，21M）仅为MAR-L模型增加了约5%的额外参数。在推理期间，扩散采样器的整体运行时间成本适中，约为10%。在作者的实现中，增加MLP宽度几乎没有额外成本（表3），部分原因是因为主要开销不是关于计算而是内存通信。

扩散损失的采样步骤。作者的扩散过程遵循DDPM[24, 10]的常见做法：作者使用1000步噪声计划进行训练，但使用更少的步骤进行推理。图4显示，在推理时使用100个扩散步骤就足以实现强烈的生成质量。

扩散损失的温度。在交叉熵损失的情况下，温度至关重要。扩散损失也为控制多样性和保真度提供了一个温度对应项。图5显示了推理时扩散采样器中温度τ的影响（见第3.2节）。温度τ在作者的模型中发挥着重要作用，与基于交叉熵的对应项的观察类似（注意表1中的交叉熵结果是它们最优温度下的结果）。

4.2 Properties of Generalized Autoregressive Models

从AR到MAR。表1也是对AR/MAR变体的比较，作者接下来讨论。首先，将AR中的光栅顺序替换为_随机_顺序可以带来显著的增益，例如，在没有CFG的情况下将FID从19.23降低到13.07。接下来，用双向注意力替换因果注意力会导致另一个巨大的增益，例如，在没有CFG的情况下将FID从13.07降低到3.43。

随机顺序、双向的AR本质上是一种MAR形式，一次预测一个标记。在每一步预测_多个_标记（'>1'）可以有效地减少自回归步骤的数量。在表1中，作者展示了进行64步的MAR变体略微牺牲了生成质量。接下来将讨论更全面的权衡比较。

速度/精度权衡。遵循MaskGIT [4]，作者的MAR具有一次预测多个标记的灵活性。这由推理时的自回归步骤数量控制。图6绘制了速度/精度权衡。与具有高效kv-cache的AR相比，MAR具有更好的权衡。

在扩散损失下，与近期流行的Diffusion Transformer (DiT) [37]相比，MAR也显示出有利的权衡。作为一种潜在扩散模型，DiT通过扩散过程建模_所有_标记之间的相互依赖。DiT的速度/精度权衡主要由其扩散步骤控制。与作者对小型MLP的扩散过程不同，DiT的扩散过程涉及_整个_Transformer架构。作者的方法更准确且更快。值得注意的是，作者的方法可以在强FID（<2.0）下以每张图像小于0.3秒的速度生成。

4.3 Benchmarking with Previous Systems

作者在表4中与最先进的系统进行了比较。作者探索了各种模型大小（见附录B），并训练了800个周期。类似于自回归语言模型[3]，作者观察到了鼓舞人心的缩放行为。进一步研究缩放可能是有希望的。在指标方面，作者在没有CFG的情况下报告了2.35的FID，大大优于其他基于标记的方法。作者最佳的成绩是1.55的FID，并且与最先进的系统相比表现良好。图7展示了定性结果。