文生图新架构！清华提出MADFormer！混合自回归与扩散的Transformer模型！
仅用于站内搜索，没有排版格式，具体信息请跳转上方微信公众号内链接
最新论文解读系列
论文名：MADFormer:MixedAutoregressiveandDiffusionTransformersforContinuousImageGeneration
论文链接：https ://arxiv. org/pdf/2506. 07999
多模态生成模型在诸如文本到图像生成、图像描述和视频合成等视觉-语言任务中取得了显著成功。这些系统通常采用自回归(AR)模型处理语言，一次生成一个离散标记，具有很强的上下文连贯性；采用扩散模型进行图像生成，通过逐步对潜在表示进行去噪，在连续域中运行。这种分工充分发挥了两者的优势:自回归模型擅长结构化序列建模，而扩散模型通过迭代细化实现高图像保真度。
多模态生成领域的最新进展越来越多地将自回归(AR)方法与基于扩散的方法相结合，充分发挥它们的互补优势:自回归模型能够捕捉长距离依赖关系，生成连贯且具有上下文感知的输出；而扩散模型则在连续潜在空间中运行，以优化高保真的视觉细节。然而，现有的混合模型往往缺乏关于如何以及为何在这两种范式之间分配模型容量的系统指导。在这项工作中，我们提出了MADFormer，一种混合自回归与扩散的Transformer模型，它可作为分析自回归-扩散权衡的试验平台。MADFormer将图像生成划分为空间块，使用自回归层对各块进行一次性全局条件约束，并使用扩散层对每个块进行迭代局部细化。通过在FFHQ-1024和ImageNet数据集上进行的可控实验，我们得到了两个关键见解:(1)按块划分显著提高了高分辨率图像的生成性能；(2)垂直混合自回归层和扩散层可实现更好的质量-效率平衡，在推理计算受限的情况下，可将FID(弗雷歇inception距离)提高多达75%。我们的研究结果为未来的混合生成模型提供了实用的设计原则。
我们介绍了MADFormer，这是一种统一的Transformer架构，不仅作为一种独立方法，还作为一个试验台，用于分析连续图像生成中AR和扩散建模之间的架构权衡。基于混合生成的最新进展，MADFormer将先前工作的关键组件集成到一个连贯且可扩展的框架中。其模块化设计能够系统地探索如何通过AR条件进行全局序列建模，以及如何通过扩散在标记序列和网络深度上进行细粒度的细节细化。
数据表示。我们根据每个模态的特征对其进行预处理。文本输入由Llama3标记器标记为离散标记。图像通过稳定扩散变分自编码器(VAE)映射到连续潜在空间。我们按从左到右、从上到下的顺序将图像潜在变量线性化，用图像开始(BOI)和图像结束(EOI)标记对其进行界定，并根据原始序列将它们与文本标记交错。这产生了离散文本标记和连续图像潜在变量的统一序列。图像潜在变量块被分组为连续块，内部具有双向注意力，并且这些块在自回归序列中被视为标记。
模型流程。我们的架构主要由一个基于标准Llama解码器块构建的单个Transformer组成，能够在潜在空间中进行统一的AR建模和基于扩散的生成。为了将不同模态投影到这个空间，我们采用特定于模态的组件:文本标记通过学习的嵌入矩阵进行嵌入，而图像潜在变量块通过U-Net上采样和下采样块进行处理。与DiT范式不同，我们不会将显式条件(例如时间步)注入到单个Transformer层中。相反，时间步信息通过U-Net下采样器直接编码到图像潜在变量中。
我们从概念上将Transformer的解码器堆栈分为两个阶段:早期层作为AR条件模块，而后期层执行递归扩散去噪。在每个扩散步骤中，早期层生成的条件在去噪循环之前添加到有噪声的潜在变量中。模型根据当前有噪声的潜在变量预测去噪后的潜在变量，遵循时间步调度。
此外，我们采用了ACDiT的策略，将干净的图像块前置到有噪声的图像块之前。这增强了自回归(AR)建模的上下文信息，提高了去噪保真度，尽管这会带来额外的计算成本。所有模态共享Transformer主干，但使用单独的参数集进行处理:文本、干净的图像块和有噪声的图像块——分别称为文本塔、干净塔和噪声塔。跨模态交互通过交错因果注意力实现，在图像块内应用双向注意力。为了支持混合的干净和有噪声输入，使用FlexAttention动态构建注意力掩码。我们在4. 4节通过消融实验证明了这种架构设计的有效性。
我们在模型的设计空间中探索了三个主要维度:扩散深度、自回归长度和损失函数。这些组件从根本上影响训练收敛性、生成速度和输出质量，且与模型大小无关。
扩散深度。我们将扩散深度定义为专用于去噪过程的Transformer层数，对应于图1图像塔中所示的解码器块。虽然先前的工作通常使用整个模型进行扩散，但我们认为这会导致不必要的计算开销。我们认为，单遍自回归条件阶段可以有效地捕捉模态间和块间的依赖关系，只允许后面的一小部分层通过递归去噪来细化块内细节。在4. 2节中，我们比较了使用整个模型、、和层进行扩散的情况，在4. 1节中，我们提出了在推理预算受限的情况下的最优自回归/扩散层比率。
我们通过公式将这种逐层结构形式化:图像塔由若干层组成，前层用于自回归条件，最后层用于基于扩散的去噪。自回归阶段计算先前生成的图像块的条件表示:
这里，PosEnc是ACDiT[Hu等人，2024]中提出的多维旋转位置编码(RoPE-ND)的简写。为了开始扩散阶段，我们注入加噪的真实潜在变量，其中是标准高斯噪声，以及条件。扩散层的目标是将去噪为干净的潜在变量
最终输出作为模型对干净潜在变量的预测。
自回归长度。我们研究了图像块自回归处理的最优粒度。自回归长度指的是输入图像被划分的图像块数量(例如，图1的示例中自回归长度为2)。不同的划分策略在带来最小计算开销的同时，会显著影响自回归条件建模的质量。在4. 2节中，我们评估了在不同图像分辨率下划分为4、16和64个块的情况，并根据实验结果提出了推荐的划分策略。
损失函数。先前的工作强调了损失函数设计对于稳定训练和高质量生成的重要性。按照惯例，我们对文本标记应用负对数似然损失，对图像潜在预测应用均方误差损失。此外，我们的架构启用了两个新的辅助损失项，每个都由一个可调超参数调节。我们在4. 6节报告它们的效果。
图3:MADFormer图像塔的架构细节。隐藏损失和干净塔损失是使用干净塔输出和条件输入计算的，并与真实图像块进行比较。
隐藏损失。为了鼓励自回归模块生成对去噪有信息价值的条件，我们引入了自回归生成的条件与下一个图像块的真实潜在变量之间的损失。理想情况下，该条件应该完全编码它之前的干净潜在变量。
干净塔损失。我们在每个干净图像块的输出与其对应的下一个干净潜在块之间应用辅助损失，类似于自回归(AR)模型中的下一个标记预测。结合通过涉及噪声块的注意力反向传播的梯度，这一目标促使干净塔编码对去噪有益的预测信号。
总损失计算为各个目标的加权和:
其中和如图3所示。和是固定权重，遵循Transfusion；和是第4. 6节中分析的可调超参数
通过一系列实验和评估，我们展示了模型设计空间的广度，并呈现了从全面的消融研究中得出的关键见解。
数据集我们在两个广泛使用的数据集上进行消融研究。FFHQ数据集包含70,000张分辨率为的高质量人脸图像。ImageNet数据集包含约128万张图像，分为1000个类别，分辨率为。
训练对于FFHQ-1024的实验，我们使用一个具有13亿参数、28个解码器层的模型，该模型由用于处理干净和有噪声图像块的两组参数组成。对于ImageNet的消融实验，我们采用一个具有21亿参数的模型，该模型包括用于文本处理的额外组件:一个文本塔、标记嵌入和一个语言建模头。在这两种配置中，U-Net上采样器和下采样器贡献了大约个参数。变分自编码器(VAE)在训练期间保持冻结；仅更新模型的其余部分。
我们使用AdamW优化器和热身稳定衰减(WSD)学习率调度进行训练，峰值学习率为，权重衰减为。应用衰减因子为0. 9999的指数移动平均(EMA)来稳定训练。在训练期间进行图像去噪时，我们使用1000个扩散步骤，遵循DDPM时间步调度器。
在FFHQ-1024上，我们以64的批量大小训练210,000步。在ImageNet上，我们以256的批量大小训练250,000步。附录A中提供了额外的训练配置和架构细节。请注意，与MAR和ACDiT相比，我们在ImageNet数据集消融实验中相对较高的FID(弗雷歇inception距离)分数，主要是由于训练轮数显著减少(50轮，而MAR为400轮，ACDiT为800轮)以及省略了无分类器引导(CFG)，因为我们的重点是进行可控的设计空间分析，而不是最终性能。
评估我们使用弗雷歇inception距离(FID)作为图像质量的主要指标。对于FFHQ-1024，FID是在8000个样本上计算的；对于ImageNet，我们遵循标准的FID-50K协议，使用50,000张生成的图像。除非在消融实验中另有说明，样本使用DDIM采样器生成，FFHQ-1024使用250步，ImageNet使用100步。最终的FID分数是对最后五个检查点(每10,000步)进行平均，以减少方差。我们还报告函数评估次数(NFE)，以比较不同配置下的推理速度和计算效率。
利用我们的MADFormer测试平台，我们沿着两个主要维度系统地探索了混合自回归(AR)-扩散模型的设计空间:(1)层维度，它控制着自回归和扩散之间Transformer层的划分；(2)标记维度，它决定了图像潜在变量如何进行自回归分割。这两个维度决定了生成质量和计算效率之间的权衡。我们还对架构变体和训练策略进行了辅助消融实验，为构建高效、高保真的自回归-扩散生成器提供了实用指导。
当推理预算有限时，最优的自回归-扩散层分配应优先考虑自回归建模而非扩散建模(图4)。
这是我们消融实验得出的一个关键见解。我们在不同的计算预算下评估模型性能，计算预算通过推理期间的函数评估次数(NFE)来衡量——具体而言，对于FFHQ数据集，NFE分别为280、700、1400和2800；对于ImageNet数据集，NFE分别为84、140、224、280和700。对于每个设置，我们在固定的28层预算内比较不同自回归-扩散层划分的模型。如图4所示，在低NFE设置下，为自回归分配更多层始终能提高弗雷歇inception距离(FID)分数。在两个数据集上，自回归占比高的配置都大幅优于扩散占比高的配置——在计算资源紧张的情况下，FID提升幅度高达60-75%。然而，随着计算资源的增加，趋势发生逆转:扩散占比高的配置开始占据主导地位，尤其是在高分辨率数据集上。这些结果表明，自回归层在建模全局结构方面计算效率更高，而在预算允许的情况下，扩散在细粒度细化方面表现出色。
图4:不同自回归-扩散层比例下的FID与NFE关系。在低NFE(即计算资源受限)情况下，自回归占比高的配置(例如，自回归:扩散为3:1，)始终优于扩散占比高的配置(例如，全扩散层，)，FID最多降低。随着计算资源的增加，扩散占比高的设置表现更好，这表明在结构建模和扩散在细化方面的互补作用。
增加扩散深度可提高生成保真度。我们进一步在总深度固定为28且扩散步数相同的情况下，对分配给扩散的层数进行消融实验。如表1所示，增加扩散深度始终能提高FID分数。这表明自回归和扩散组件都受益于足够的容量，并强调了根据生成设置明智平衡这些组件的重要性。
表1:扩散深度消融实验。我们实验中的所有模型在FFHQ数据集上训练步，在ImageNet数据集上训练250k步(50个epoch)。
最优自回归(AR)长度取决于图像分辨率。我们将长度定义为按顺序处理的图像块数量，这由图像为自回归建模进行分区的方式决定。例如，分辨率为且的FFHQ图像被划分为四个补丁。如表2所示，最优自回归长度因数据集而异:FFHQ在使用16个块时表现最佳，而ImageNet则更倾向于相同大小的单个块。这些发现与ACDiT在ImageNet上的研究结果一致，即较长的自回归序列会降低质量。然而，我们在FFHQ-1024上的实验揭示了一个更微妙的趋势，表明高分辨率图像受益于更细粒度的自回归分解。我们认为，最优自回归长度取决于图像分辨率、架构和数据集特征，这凸显了未来在这些方向开展研究的必要性。
干净块和自回归条件都能提高扩散质量。我们引入了两个辅助模块来引导扩散过程:(1)干净块——未损坏的图像表示，添加到输入序列的前面；(2)自回归条件——从先前块自回归生成的上下文。这些模块将结构注入去噪轨迹，并作为互补先验。如表3所示，在FFHQ和ImageNet上移除任何一个模块都会导致FID(弗雷歇inception距离)持续下降，这证实了它们各自的有效性以及将它们结合起来的好处。
在我们的实验中，为文本、干净图像块和含噪图像块使用单独的参数集效果甚微。我们探索了通过为文本、干净图像块和含噪图像块分配不同的参数集(例如，前馈网络(FFNs)、查询-键-值-输出投影(QKVOprojections)和层归一化(layernorms))将稀疏性引入模型。这种对模型设计的消融实验的灵感来源于这样一种直觉，即分离参数空间可以促进跨模态学习不同的分布，这一想法得到了LMFusion[Shi等人，2024]的支持。然而，表4中的消融实验结果显示性能提升并不明显。换句话说，在所有模态间共享参数的密集模型仍然有效。
扩散过程中的序列级因果注意力至关重要。受先前工作(如)的启发，该工作用辅助多层感知机(MLP)替代注意力机制，以独立地对每个图像块进行去噪。我们通过修改注意力掩码，将因果注意力限制在块内范围，从而对这一设置进行消融实验，有效地切断了扩散阶段各块之间的信息流。这模拟了独立的逐块MLP去噪过程。如表5所示，这导致生成质量显著下降，在FFHQ数据集上，弗雷歇inception距离(FID)分数从17. 8恶化到21.2，在ImageNet数据集上从30. 0恶化到96.5。这些结果强调了序列级因果注意力在促进跨块连贯细化方面的关键作用，我们假设这对于在生成过程中保持空间一致性和细粒度细节至关重要。
辅助损失可以改善训练动态。为了促进结构化的潜在表示并促进信息流，我们引入了两种辅助损失:对自回归(AR)条件潜在变量的隐藏损失和对干净块潜在变量的干净塔损失，这两种损失均由移位后的干净块目标监督(见第3. 2节)。如表6所示，隐藏损失显著改善了FID，当系数为0. 1时，FID从19. 4降至17.8。干净塔损失对最终生成质量的影响相对较小。然而，我们通过实验观察到，这两种损失在训练开始时都有助于加速收敛。
我们提出了多模态自适应扩散变压器(MADFormer)，这是一种统一的变压器(Transformer)架构，它在图像生成流水线中灵活地结合了自回归(AR)和扩散建模。通过将图像划分为空间块，并在模型深度上分配自回归(AR)层和扩散层，多模态自适应扩散变压器(MADFormer)实现了结构化的全局条件控制和高保真的局部细化。通过一系列可控实验，我们证明了逐块和逐层的自回归(AR)-扩散混合在不同的生成设置下都有助于提高质量和效率。我们的分析强调了自回归(AR)和扩散组件之间的权衡，为在不同计算预算下平衡模型容量提供了实用指南——特别是，增加自回归(AR)层对推理计算量较低的设置有益。我们希望多模态自适应扩散变压器(MADFormer)能为未来对混合生成架构的研究迈出坚实的一步，这种架构能自适应地利用自回归(AR)和扩散范式的优势。
感谢你看到这里，也欢迎点击关注下方公众号并添加公众号小助手加入官方读者交流群，一个有趣有AI的AIGC公众号:关注AI、深度学习、计算机视觉、AIGC、StableDiffusion、Sora等相关技术，欢迎一起交流学习💗～
