填补领域空白！TerDiT：首次探索大规模DiT模型量化问题（MMLab出品）

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

项目链接：https://github.com/Lucky-Lance/TerDiT

最近在大规模预训练的文本到图像扩散模型方面的发展显著提高了高保真图像的生成能力，特别是基于transformer架构的扩散模型（DiTs）的出现。在这些扩散模型中，扩散transformer展示了卓越的图像生成能力，降低了FID分数并提高了可扩展性。然而，由于其庞大的参数数量，部署大规模的DiT模型可能会非常昂贵。尽管现有研究已经探索了扩散模型的高效部署技术，如模型量化，但关于基于DiT模型的研究仍然很少。为了解决这一研究空白，本文提出了TerDiT，一种面向量化感知训练（QAT）和高效部署的基于transformer的三值化扩散模型方案。本文专注于DiT网络的三值化化，并将模型规模从600M扩展到4.2B。本文的工作为大规模DiT模型的高效部署策略探索做出了贡献，证明了从头训练极低比特的扩散transformer模型的可行性，同时在图像生成能力上保持与全精度模型相当的竞争力。

介绍

大规模预训练的文本到图像扩散模型的进步已经成功生成了复杂且高度保真于输入条件的图像。值得注意的是，基于transformer架构的扩散模型（DiTs）的出现代表了该研究领域的一个重要进展。与其他扩散模型相比，扩散transformer展示了在更高计算量下实现更低FID分数的能力。最新的研究突出了扩散transformer架构在图像生成方面的显著能力，例如在Stable Diffusion 3方法中展示的成果，以及在视频生成方面的出色表现，如Sora所展示的工作。

鉴于扩散transformer模型的出色表现，研究人员现在越来越多地深入研究这些视觉模型的扩展规律，这与大语言模型相似。例如，Stable Diffusion 3提供了一系列参数规模的训练DiT模型，从8亿到80亿。此外，有研究人员推测Sora可能拥有大约30亿个参数。由于这些模型拥有庞大的参数数量，部署这些DiT模型往往成本高昂，尤其是在某些终端设备上。

为了应对部署难题，最近已经有一些关于高效部署扩散模型的研究工作，其中大多数集中在模型量化方面。然而，据本文所知，目前的研究仍然存在两个主要缺陷。首先，虽然量化基于 U-Net 的扩散模型已经受到了很多关注，但对于基于transformer的扩散模型的量化方法探索仍然非常有限。其次，目前文献中的大多数主流方法主要依赖于后训练量化 (PTQ) 技术来进行模型量化，这会导致不可接受的性能下降，特别是在极低比特宽度（例如2比特和1比特）下。然而，神经网络的极低比特量化非常重要，因为它可以显著减少部署所需的计算资源，尤其对于具有巨大参数规模的模型。在本文的研究过程中，本文发现目前还没有研究考虑 DiT 模型的极低比特量化。

为了解决这些缺陷，本文提出利用量化感知训练（QAT）技术对大规模 DiT 模型进行极低比特量化。在大规模模型领域，低比特 QAT 方法已在大语言模型（LLM）领域进行了讨论。最近的研究表明，从头开始训练具有极低比特参数（例如二进制和三进制）的大语言模型，也可以达到与全精度模型相当的竞争性能。这表明大规模模型中仍然存在显著的精度冗余，并且暗示了 QAT 方案对于大规模 DiT 模型的可行性。

在本文中，本文主要关注三值权重网络，并提供了 TerDiT，这是本文所知的首个用于 DiT 的量化方案。本文的方法实现了三值扩散transformer模型的量化感知训练（仅限权重）和高效部署。与 LLM 和 CNN 中线性层的简单量化不同，本文发现直接对 DiT 模块中的 adaLN 模块进行权重三值化，会导致归一化层中的大尺寸尺度和偏移值（由于权重量化和梯度近似），这与全精度模型相比，导致收敛速度较慢和模型性能较差。因此，本文提出了一种 adaLN 的变体，通过在 adaLN 模块的三值线性层之后应用 RMS Norm，有效地缓解了这一训练问题。

通过这种修改，本文将三值 DiT 模型的参数规模从 600M（DiT-XL/2的规模）扩展到 4.2B（Large-DiT-4.2B的规模），发现具有更多参数的模型能够收敛到更好的结果。本文进一步采用现有的2-bit CUDA 内核来部署训练后的三值 DiT 模型，使模型checkpoint 大小减少了十倍以上，推理内存消耗减少了约六倍，同时实现了具有竞争力的生成质量。主要贡献总结如下：

受低位 LLMs 量化感知训练方案的启发，本文研究了针对三值 DiT 模型的 QAT 方法，并引入了 DiT 特定的改进以获得更好的训练效果，这在 DiT 文献中尚未被探索。
本文将三值DiT模型的参数规模从600M扩展到4.2B，并基于现有的2-bit CUDA内核在GPU上部署了训练后的三值DiT模型，使得4.2B DiT模型的推理内存消耗小于3GB。
与全精度模型在ImageNet基准测试（图像生成）中的对比评估结果展示了本文提出的TerDiT方案的有效性。

本文的研究是首次尝试探索DiT模型的量化问题。本文专注于量化感知训练和大规模三值DiT模型的高效部署，为未来研究在极低比特精度下部署DiT模型提供了宝贵的见解。

TerDiT

TerDiT，这是一个用于进行仅权重量化感知训练和高效部署大规模三值DiT模型的框架。本文首先简要回顾扩散transformer（DiT）模型。然后，基于之前开源的Large-DiT，阐述了量化函数和量化感知训练方案，并进行特定于QAT的模型结构改进以优化网络训练，并介绍了三值部署方案。

扩散transformer模型

扩散transformer（Diffusion Transformer）。扩散transformer（DiT）是一种架构，它用操作潜在patches的transformer替代了扩散模型中常用的U-Net骨干结构。类似于下图2（C）中展示的视觉transformer（ViT）架构，DiT首先将空间输入划分为一系列tokens，然后通过一系列transformer块（下图2（B））进行去噪处理。为了处理额外的条件信息（例如噪声时间步t、类别标签l、自然语言输入），DiT利用自适应归一化模块（adaLNZero）将这些额外的条件输入插入到transformer块中。在最后一个transformer块之后，应用标准线性解码器来预测最终的噪声和协方差。DiT模型的训练方式与基于U-Net的扩散模型相同。

DiT中的AdaLN模块。DiT与传统ViT的主要区别在于需要注入条件信息以进行图像生成。DiT在每个transformer块中使用零初始化的自适应层归一化（adaLN-Zero）模块，如上图2（B）红色部分所示，该模块根据输入条件c计算维度级的缩放和偏移值。

AdaLN 是 DiT 模型中的一个重要组件，其效果已被证明优于交叉注意力和上下文条件方法。在 DiT 架构中，AdaLN 模块集成了一个包含大量参数的 MLP 层，占模型总参数的约 10% 到 20%。在 TerDiT 的训练过程中，本文观察到直接对该模块进行权重三值化会导致不理想的训练结果。

模型量化

如上文所示，理解DiT模型的扩展规律越来越受到关注，这对于开发和优化大语言模型（LLM）至关重要。在最近的探索中，Large-DiT成功地将模型参数从600M扩展到7B，结合了LLaMA和DiT的方法。结果表明，参数扩展可以潜在地提升模型性能，并加快标签条件的ImageNet生成任务的收敛速度。受此启发，本文提出进一步研究DiT模型的三值化，这可以缓解部署大规模DiT模型相关的挑战。在本小节中，本文介绍量化函数和量化感知训练方案。

量化函数。 为了构建三值权重的 DiT 网络，本文将原始 Large-DiT 模块中的所有自注意力、前馈和 MLP 中的线性层替换为三值线性层，得到一组三值 DiT 模块（上图 2（A））。对于三值线性层，本文采用类似于 BitNet b1.58的 absmean 量化函数。首先，通过将权重矩阵中的每个元素除以矩阵中所有元素的平均绝对值来归一化权重矩阵。归一化后，权重矩阵中的每个值都会被四舍五入到最接近的整数，并限制在 ${-1, 0, +1\}$ 集合中。

将权重矩阵中的每个元素除以矩阵中所有元素的平均绝对值。归一化后，权重矩阵中的每个值都会被四舍五入到最接近的整数，并限制在集合 ${-1, 0, +1\}$ 中。

参考当前流行的大语言模型（LLMs）的量化方法，本文在量化后也将一个可学习的缩放参数 $\alpha$ 乘以每个三值线性矩阵，从而得到最终的值集合为 $\{-\alpha, 0, +\alpha\}$ 。量化函数公式如下：

其中 $\epsilon$ 被设置为一个小值（例如， $10^{-6}$ ）以避免除以 0，并且

TerDiT 是一种仅对权重进行量化的方案，本文不对激活值进行量化。

量化感知训练方案。 基于上述设计的量化函数，本文从头开始训练一个 DiT 模型，利用直接传递估计器（STE），允许梯度通过不可微分的网络组件传播。在整个训练过程中，本文保留网络的全精度参数。对于每一步训练，通过前向传播中的三值量化函数从全精度参数计算出三值权重，并在反向传播中将三值权重的梯度直接应用于全精度参数进行参数更新。

然而，本文发现收敛速度非常慢。即使经过多次训练迭代，损失值也无法降低到合理范围。本文认为这个问题可能源于三值线性层通常会导致较大的激活值，并提出在接下来的小节中通过针对 QAT（量化感知训练）特定的模型结构改进来解决这个问题。

QAT特定模型结构改进

三值线性层的激活分析。 在三值线性层中，所有参数都从集合 ${-α, 0, +α\}$ 中取一个值。通过该层的值会变成较大的激活值，这可能会阻碍神经网络的稳定训练。本文进行了一项初步研究，以定性地展示三值线性权重对激活值的影响。

本文随机初始化一个三值线性层，输入特征维度设置为1024，输出特征维度设置为9216（对应LargeDiT中adaLN模块的线性层）。权重参数通过量化函数，并接收一个大小为 $512 \times 1024$ 的矩阵输入（填充为1）。激活值分布的箱线图显示在图3的中间部分。本文还计算了通过使用相同随机种子生成的全精度线性层后的激活值分布，显示在下图3的右侧部分。可以看出，与全精度线性层相比，三值线性层导致了非常大的激活值。

通过对三值线性层的输出应用层归一化，可以缓解由三值线性权重带来的大激活值问题。本文在三值线性层之后添加了一个RMS归一化层（类似于LLaMA），并获得了激活值分布（如下图3左侧所示）。在这种情况下，激活值在通过归一化层后被缩放到一个合理范围，从而导致更稳定的训练行为。这一观察结果也与[17]中的结论一致，其中在每个量化线性层的激活量化之前应用了层归一化函数。

RMS归一化的AdaLN模块。 基于上述见解，本文分析了DiT模型以改进QAT特定的模型结构。在标准的ViT Transformer块中，层归一化应用于每个自注意力层和前馈层。DiT块中的自注意力层和前馈层也是如此，这有助于适当地缩放激活值范围。然而，由于上文中介绍的AdaLN模块的存在，DiT块与传统的Transformer块有所不同。值得注意的是，这个模块没有应用层归一化。在全精度训练的情况下，缺乏层归一化并不会产生显著影响。然而，对于三值DiT网络来说，其缺失可能会导致adaLN（归一化）模块中的维度尺度和偏移值过大，从而对模型训练产生不良影响。为了解决这个问题，本文在每个三值DiT块的AdaLN模块的MLP层之后引入了RMS归一化：

最终的TerDiT模型结构如上图2(A)所示。这个小的修改可以带来更快的收敛速度和更低的训练损失，从而在定量和定性评估中取得更好的结果。为了更好地展示这一效果，在原文附录中分析了模型训练后引入或不引入RMS Norm的实际激活分布。

部署方案

在训练了DiT模型之后，本文发现目前没有有效的开源三值网络部署解决方案。在这种情况下，本文使用2位实现来部署训练好的网络。具体来说，本文使用文献[44]提供的pack_2bit_u8()函数，将三值线性权重打包成int8值（4个三值数打包成一个int8数）。在DiT模型的推断过程中，本文即时调用相应的unpack_2bit_u8()函数，将打包的2位数字恢复为浮点数值，然后进行后续计算。添加解包操作会减慢推断过程，但本文相信，随着对模型三值化研究的深入，将会有更多硬件支持来加速推断过程。

实验

在本节中，本文进行了一系列实验来评估本文提出的TerDiT。本文展示了主要的评估结果，进行了部署效率比较，并说明了RMS Normalized adaLN模块的有效性。本文的DiT实现基于开源代码Large-DiT-ImageNet4。本文分别对具有600M（DiT-XL/2的大小）和4.2B（Large-DiT-4.2B的大小）参数的三值DiT模型进行了实验。

主要评价结果

本文在本小节中提供了TerDiT模型的定量和定性评估结果。据本文所知，目前尚无关于扩散transformer模型量化的已发表工作，因此本文主要在本小节中将其与具有代表性的全精度扩散模型进行比较。

关于TerDiT基线的备注。 据本文所知，目前仍没有研究DiT模型量化的工作。除了在本小节中与全精度模型进行比较外，本文还在其他小节中建立了一些基线进行比较。对于QAT基线，本文直接训练了在Sec. 4.3中的adaLN模块中没有RMS Norm的三值DiT模型。为了与现有的PTQ方法进行比较，本文对预训练模型进行了4位权重量化，使用与TerDiT相同的一组参数，结果发现它们无法生成可视的图像。

实验设置。 按照原始DiT论文的评估设置，本文在ImageNet数据集上训练了600M和4.2B的三值DiT模型。由于计算资源的限制，本文在256×256分辨率下训练和评估模型，但本文认为评估结果已经具有很强的代表性。本文将TerDiT与一系列全精度扩散模型进行比较，并报告FID、sFID、Inception Score、Precision和Recall（50k生成图像），参考[48]。本文还提供了训练阶段的总图像数量（百万），如[23]所示，以进一步了解不同生成模型的收敛速度。

训练细节。 本文在 8 个 A100-80G GPU 上训练 600M TerDiT 模型，共进行了 1750k 步，批量大小设置为 256；在 16 个 A100-80G GPU 上训练 4.2B 模型，共进行了 1180k 步，批量大小设置为 512。初始学习率设置为 $\times 10^{-4}$ 。在 600M 模型训练 1550k 步后和 4.2B 模型训练 550k 步后，本文将学习率降低到 $\times 10^{-4}$ 以进行更精细的参数更新。

定量结果分析。 评估结果列在下表1中。TerDiT是针对DiT模型的QAT方案，因此在所有全精度模型中，本文特别关注DiT-XL/2（675M）和Large-DiT-4.2B。在没有分类器自由指导的情况下，TerDiT-4.2B在测试结果上与DiT-XL/2非常相似（使用的训练图像数量要少得多）。在有分类器自由指导（cfg=1.5）的情况下，TerDiT-4.2B-G的表现优于LDM-G，同时与两个全精度DiT结构模型相比仅带来了非常轻微的性能下降。此外，TerDiT-4.2B-G的评估结果优于TerDiT-600M-G，这表明参数更多的模型在量化后会带来更小的性能下降。

为了直观地展示TerDiT的有效性，本文在下图4中展示了一些定性比较结果，涉及TerDiT-4.2B、DiT-XL/2和Large-DiT4.2B。从视觉感知的角度来看，TerDiT生成的图像与全精度模型生成的图像之间没有显著差异。

部署效率对比

部署效率的提升是本文提出TerDiT方案的动机。在本小节中，本文对TerDiT-600M/4.2B、DiT-XL/2和Large-DiT-4.2B进行了比较，以讨论TerDiT在实际部署中所能带来的效率提升。下表2展示了四种DiT模型的checkpoint 大小。本文还记录了在单个A100-80G GPU上，总扩散采样循环（步数=250）的内存使用情况和推理时间。

从表格中可以看出，TerDiT大大减少了checkpoint 大小和内存使用。4.2B三值化DiT模型的checkpoint 大小和内存使用显著小于Large-DiT-4.2B，甚至比DiT-XL/2还要小。这为在终端设备（如手机）上部署模型带来了显著优势。尽管由于需要解包操作，本文观察到推理速度较慢，但本文相信，随着未来硬件支持的提升，三值化权重网络的计算优势将会得到充分展示。

RMS归一化AdaLN模块的讨论

TerDiT对DiT模型结构的主要修改是在adaLN模块中的MLP之后增加了RMS Norm。在这一部分，本文与基线三值化模型进行比较，以展示RMS Norm对训练过程和训练结果的影响。

实验设置。 本文在ImageNet数据集上以256×256分辨率训练具有600M和4.2B参数的三值化DiT模型。对于每种参数规模，本文训练了两个模型，一个在adaLN模块中使用了RMS Norm，另一个则没有（本文的基线模型）。本文记录了训练过程中的损失曲线，并每100k训练步测量一次FID-50k分数（cfg=1.5）。

训练细节： 为了进行公平的比较，本文在8个A100-80G GPU上训练所有的三值化DiT模型，批量大小设置为256。整个训练过程中，学习率设置为 $5e^{-4}$ 。

结果分析： 训练损失和评估得分分别显示在下图5和下图6中。如图所示，使用RMS Normalized adaLN模块进行训练将导致更快的收敛速度和更低的FID分数。另一个观察结果是，参数更多的模型相比参数较少的模型能实现更快且更好的训练效果。这在一定程度上也反映了三值化DiT模型的扩展规律。

讨论和未来展望

本文在成功的大语言模型低比特训练方法的基础上，提出了针对大规模三值化DiT模型的量化感知训练（QAT）和高效部署方法。在ImageNet数据集（256×256）上的竞争性评估结果证明了从头开始训练大型三值化DiT模型的可行性，同时实现了与全精度模型相当的结果。据本文所知，这是首个关于DiT模型量化的研究。

虽然本文认为这项工作为DiT模型的低比特量化提供了有价值的见解，但它仍然存在一些局限性。首先，训练三值化DiT比全精度网络更不稳定且耗时。在本文中，尽管本文讨论了通过添加归一化方法来使训练更稳定，但相较于训练全精度网络（如Large-DiT-4.2B），它仍然更耗时，这将在更广泛的背景下导致模型训练期间二氧化碳排放量的增加。

其次，由于计算资源的限制，本文没有进行ImageNet 512×512实验，也没有进行文本到图像生成任务的实验。然而，本文相信ImageNet 256×256基准上的评估结果已经相当具有代表性。剩余的任务将留待本文未来的工作中进行。本文希望本文的工作可以减少图像生成模型的部署需求，并能激励社区在未来加入本文，共同促进这一研究领域的更广泛发展。