上交大、上海 AI 实验室、ViVO 推出 TerDiT ，极低比特量化感知训练和和高效部署方案！

近期大规模预训练的文本到图像扩散模型的发展显著提高了高保真图像的生成，特别是基于 Transformer 架构的扩散模型（DiTs）的出现。在这些扩散模型中，扩散 Transformer 展示了卓越的图像生成能力，降低了FID分数并提高了可扩展性。然而，部署大规模的DiT模型可能由于参数数量庞大而变得昂贵。
尽管现有研究已经探索了扩散模型的效率部署技术，如模型量化，但关于基于DiT的模型的研究仍然很少。为了填补这一研究空白，在本文中，作者提出了TerDiT，一个针对三元扩散 Transformer 模型的量化感知训练（QAT）和高效部署方案。
作者专注于DiT网络的三元化，并将模型尺寸从6亿扩展到42亿。
作者的工作为大规模DiT模型的效率部署策略的探索做出了贡献，证明了从零开始训练极低比特扩散 Transformer 模型的可行性，同时与全精度模型相比保持了有竞争力的图像生成能力。
代码将在https://github.com/Lucky-Lance/TerDiT。

1 引言

大规模预训练文本到图像扩散模型的进展导致了成功生成具有复杂性和对输入条件高保真的图像。特别是基于 Transformer 架构的扩散模型的出现，在这一研究领域中代表了重要的进步。与其他扩散模型相比，扩散 Transformer 已经展示了以更高的计算Gflops实现更低FID分数的能力[6]。近期的研究突显了扩散 Transformer 架构在图像生成能力方面的卓越表现，如Stable Diffusion 3[7]等方法，以及在视频生成方面，如Sora2所展示的出色性能。鉴于扩散 Transformer 模型的卓越性能，研究行人现在越来越多地研究这些视觉模型的扩展规律[8]，这与大型语言模型（LLMs）相似。

例如，Stable Diffusion 3提供了跨越从80亿到800亿参数范围的训练好的DiT模型。此外，研究行人推测Sora可能拥有大约300亿个参数。鉴于巨大的参数数量，在特定终端设备上部署这些DiT模型往往是昂贵的。

为了解决部署难题，最近已经有了一些关于扩散模型高效部署的工作[9; 10; 11; 12]，其中大部分聚焦于模型量化。然而，在作者看来，当前研究仍存在两个主要不足。

首先，虽然对基于U-Net的扩散模型的量化给予了大量关注，但对基于 Transformer（transformer）的扩散模型量化方法的探索仍然相当有限。其次，当前文献中大多数流行的方法严重依赖于后训练量化（PTQ）技术进行模型量化，这导致在极低比特宽度（例如，2比特和1比特）时性能下降到不可接受的程度。

然而，神经网络的极低比特量化非常重要，因为它可以显著减少部署所需的计算资源[16]，特别是对于参数规模巨大的模型。在作者的研究中，作者发现仍没有考虑对DiT模型进行极低比特量化。

为了克服这些不足，作者提出利用量化感知训练（QAT）技术对大规模DiT模型进行极低比特量化。在LLM领域，已经讨论了针对大规模模型的低比特QAT方法。

最近的研究发现，从零开始训练具有极低比特参数（例如，二进制和三元）的大型语言模型也能与全精度模型相比具有竞争性的性能。这表明大规模模型中仍然存在显著的精度冗余，并暗示了QAT方案对于大规模DiT模型的可行性。

在本文中，作者主要关注三元权重网络[19]，并提供TerDiT，据作者所知，这是第一个针对DiTs的量化方案。作者的方法实现了量化感知训练（仅权重）并高效部署三元扩散 Transformer 模型。与LLMs和CNNs中线性层的简单量化不同，作者发现直接对DiT块[6; 21]中的adaLN模块[20]进行权重三元化会导致归一化层中维度上的尺度和偏移值与全精度模型相比变大（由于权重量化、梯度近似），这会导致收敛速度减慢和模型性能下降。

因此，作者提出了adaLN的一个变体，通过在adaLN模块的三元线性层后应用RMS Norm[22]来有效缓解这一训练问题。

通过这一修改，作者将三元DiT模型的参数规模从6亿（DiT-XL/2[6]的大小）扩展到42亿（Large-DiT-4.2B[23]的大小），发现参数更多的模型可以收敛到更好的结果。

作者进一步使用现有的2比特CUDA Kernel 来部署训练的三元DiT模型，使得模型预训练权重大小减少了十倍以上，推理内存消耗减少了大约六倍，同时实现了具有竞争力的生成质量。

贡献总结如下：

受到针对低比特LLM模型的量化感知训练方案的启发，作者研究了针对三元DiT模型的QAT方法，并引入了针对DiT的特定改进以获得更好的训练效果，这在DiT文献中尚未被探索。
作者将三元DiT模型的规模从600M扩展到42B个参数，并进一步在基于现有2比特CUDA Kernel 的GPU上部署训练好的三元DiT，使得4.2B DiT模型的推理可使用不到3GB的GPU内存。
与全精度模型在ImageNet [24] 基准测试（图像生成）上的竞争性评估结果展示了作者提出的三元DiT方案的有效性。

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

2 Related Works

扩散模型。近年来，由于扩散模型能够生成高质量的图像以及其在各种应用中的潜力，已经引起了广泛关注。[25]首次提出了扩散模型的概念，提出了一个学习逆转扩散过程的生成模型。这项工作为该领域后续的研究奠定了基础。[1]进一步扩展了这一想法，引入了去噪扩散概率模型（DDPMs），这已成为图像生成任务的流行选择。DDPMs已经被应用于广泛的领域，包括无条件图像生成[1]，图像修复[26]和图像超分辨率[27]。此外，扩散模型也被用于文本到图像的合成，如DALL-E模型[28]和Imagen模型[5]所示。这些模型展示了扩散模型根据文本描述生成高度逼真和多样化图像的能力。此外，扩散模型已被扩展到其他模态，如音频合成[29]和视频生成[30]，展示了它们的多样性和多模态应用的潜力。

扩散模型的量化。近年来，对扩散模型的量化研究旨在提高扩散模型的效率。如中提出的后训练量化（PTQ）方法，在量化时间和数据使用方面具有优势。然而，当这些方法应用于低比特设置时，往往导致性能不佳。[31]提出将量化感知低秩 Adapter （QALoA）与PTQ方法结合使用，从而提高了评估结果。作为PTQ的替代方法，专门为低比特扩散模型量化引入了量化感知训练（QAT）方法。尽管这些QAT方法有效，但目前它们仅限于基于小型U-Net的扩散模型，这表明在大型DiT模型中应用QAT的研究还有空白。进一步探索针对具有极低比特宽度的大型DiT模型的QAT技术，可能会带来更大的效率提升，并使扩散模型在资源受限环境中得到有效的部署。

三值权重网络。三值权重网络[19]作为一种内存高效和计算高效的网络结构出现，具有显著减少推理内存使用的潜力。在专用硬件的支持下，三值权重网络还可以实现计算加速。在量化方法中，三值权重网络受到了广泛关注，目前主要探索两种方法：仅权重量化和权重激活量化。在[35]中讨论的仅权重量化中，只有权重被量化为三值。另一方面，如[36; 37]所示，权重激活量化涉及将权重和激活都量化为三值。最近的研究已经证明了三值权重网络在训练大型语言模型[18]中的适用性，取得了与全精度模型相当的结果。基于这些进展，作者首次引入了针对三值DiT模型的量化感知训练和高效部署方案。通过利用三值量化在DiT模型中的优势，作者旨在推动效率的边界，使强大的扩散模型能够在资源受限的环境中部署，为实际应用开辟新的可能性。

3 TerDiT

在本节中，作者介绍了TerDiT框架，该框架旨在进行仅权重量化感知训练以及大规模三元DiT模型的效率部署。

首先在3.1节中对扩散 Transformer （DiT）模型进行了简要回顾。

然后在3.2节中，基于之前开源的Large-DiT [23]，作者说明了量化函数和量化感知训练方案，并在3.3节中进行了针对更好网络训练的QAT特定模型结构改进，最后在3.4节中介绍了三元部署方案。

Diffusion Transformer Models

扩散 Transformer 。 扩散 Transformer [6]（DiT）是一种架构，它用操作潜在块的 Transformer 替代了扩散模型中常用的U-Net主干网络。与图2（C）中显示的视觉 Transformer （ViT）架构类似，DiT首先将空间输入分割成一系列 Token ，然后通过一系列 Transformer 块（图2（B））进行去噪处理。为了处理额外的条件信息（例如，噪声时间步长，类别标签，自然语言输入），DiT利用自适应规范化模块 [38]（adaLN-Zero）将这些额外的条件输入插入到 Transformer 块中。在最后的 Transformer 块之后，应用标准的线性解码器来预测最终的噪声和协方差。DiT模型可以像基于U-Net的扩散模型一样进行训练。

DiT中的AdaLN模块。 DiT与传统ViT的主要区别在于图像生成时需要注入条件信息。DiT在每个 Transformer 块中采用零初始化的自适应层规范化（adaLN-Zero）模块，如图2（B）中的红色部分所示，它从输入条件计算逐维的比例和偏移值：

AdaLN是DiT模型 [6] 中的一个重要组成部分，并且已经证明比交叉注意力方法和上下文条件方法更有效。在DiT架构中，adaLN模块集成了一个具有大量参数的MLP层，大约占模型总参数的10%到20%。在TerDiT的训练过程中，作者观察到直接对这一模块进行权重三元化会导致不理想的训练结果（在3.3节中进行分析）。

Model Quantization

如图1节所示，对DiT模型的缩放法则的理解越来越受欢迎，这对于开发和优化LLM已被证明是至关重要的。在最近的探索中，Large-DiT [23] 通过结合LLaMA [39; 40]的方法和DiT，成功地将模型参数从6亿扩展到70亿。结果表明，参数扩展有可能提高模型性能，并加快标签条件下的 ImageNet 生成任务的收敛速度。受到这一结果的启发，作者 Proposal 进一步研究DiT模型的三元化，这可以减轻部署大规模DiT模型所面临的挑战。在本小节中，作者将介绍量化函数和量化感知训练方案。

图2：（A）三元DiT块，（B）Large-DiT块，以及（C）原始ViT块的模型结构比较。Large-DiT（DiT）块在原始ViT块中加入了adaLN模块进行条件注入。三元DiT块进一步在adaLN模块中添加了RMS规范，以进行更好的三元化感知训练。

量化函数。为了构建一个三元权重DiT网络，作者将原始Large-DiT块中的自注意力、前馈和MLP中的所有线性层替换为三元线性层，得到一组三元DiT块（图2（A））。对于三元线性层，作者采用了类似于BitNet b1.58 [18]的_absmean_量化函数。首先，将权重矩阵归一化，即每个元素除以矩阵中所有元素的平均绝对值。归一化后，权重矩阵中的每个值四舍五入到最近的整数，并钳制在集合中。

量化感知训练方案。基于上述设计的量化函数，作者使用直接通过估计器（STE）[43]从头开始训练一个DiT模型，允许梯度通过不可微网络组件传播。作者在整个训练过程中保持网络的全程精度参数。对于每个训练步骤，通过前向传递的三元量化函数从全程精度参数计算三元权重，而在反向传递中进行参数更新时，将三元权重的梯度直接应用于全程精度参数。然而，作者发现收敛速度非常慢。即使在许多训练迭代之后，损失也不能降低到一个合理的范围内。作者发现这可能是由于三元线性层通常导致大的激活值，作者将在下面的小节中通过特定的QAT模型结构改进来解决这个问题。

QAT-specific Model Structure Improvement

三元线性层激活分析。 在三元线性层中，所有参数取自集合{-，0，+}中的一个值。通过这个层的数据会产生较大的激活值，这可能会阻碍神经网络的稳定训练。作者进行了一个初步研究，定性地展示了三元线性权重对激活值的影响。

图3：激活值分析。作者比较了通过三元权重线性层（带与不带RMS归一化）的激活值，并以全精度线性层作为参考。不带RMS归一化的三元线性层导致激活值极大，给神经网络训练带来不稳定。然而，当应用归一化层时，激活值被缩放到一个合理的范围内，与全精度层观察到的相似。

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

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

RMS归一化的AdaLN模块。 基于上述洞察，作者分析了针对QAT特定模型结构改进的DiT模型。在标准的ViT Transformer 块中，层归一化应用于每个自注意力层和前馈层。DiT块中的自注意力层和前馈层也是如此，这有助于适当缩放激活的范围。然而，由于Sec. 3.1中引入的AdaLN模块的存在，DiT块与传统的 Transformer 块不同。值得注意的是，这个模块没有应用层归一化。在全精度训练的背景下，层归一化的缺失不会产生重大影响。然而，对于三元DiT网络，其缺失可能导致adaLN（归一化）模块中维度方向的大尺度和平移值，对模型训练产生不良影响。为了缓解这个问题，作者在每个三元DiT块中的adaLN模块的MLP层后引入了RMS归一化：

最终的TerDiT模型结构如图2（A）所示。这个小的修改可以使模型训练更快地收敛，并降低训练损失，导致更好的定性和定量评估结果。为了更好地展示效果，实际训练后带与不带RMS归一化的激活分布在Sec. A.1中进行了分析。

Deployment Scheme

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

4 Experiments

在本节中，作者进行了一系列实验来评估作者提出的TerDiT。

作者在4.1节展示了作者的主要评估结果，在4.2节进行了部署效率的比较，并在4.3节说明了RMS归一化adaLN模块的有效性。

作者的DiT实现基于开源的Large-DiT-ImageNet4代码。

作者分别对具有600M（DiT-XL/2的大小）和4.2B5（Large-DiT-4.2B的大小）参数的三元DiT模型进行了实验。

主要评估结果

本小节作者提供了TerDiT的定性和定量评估结果。据作者所知，到目前为止还没有关于扩散变换模型量化的已发表工作，因此本小节作者主要关注与代表性全精度扩散模型的比较。

关于TerDiT Baseline 的说明。据作者所知，还没有研究DiT模型量化的工作。除了在本小节与全精度模型进行比较外，在其他小节作者还建立了一些进行比较的 Baseline 。对于QAT Baseline ，作者在4.3节的adaLN模块中直接训练不带RMS归一化的三元DiT模型。为了与现有的PTQ [45]方法进行比较，作者对与TerDiT相同参数集上的预训练模型进行了4位权重量化，发现它们无法生成可观看的图像（详细内容见附录A.2）。

实验设置。遵循原始DiT论文[6]的评估设置，作者在ImageNet数据集上训练了600M和4.2B的三元DiT模型。由于计算资源限制，作者在256×256分辨率下训练和评估模型，但作者认为这些评估结果已经相当有代表性。作者将TerDiT与一系列全精度扩散模型进行了比较，根据[48]报告了FID [46]，sFID [47]，Inception Score，Precision和Recall（5万张生成图像）。作者还提供了训练阶段的总图像数量（百万张），如[23]所示，以便进一步了解不同生成模型的收敛速度。

训练细节。作者在8个A100-80G GPU上用1750k步训练了600M的TerDiT模型，批处理大小设置为256；在16个A100-80G GPU上用1180k步训练了4.2B模型，批处理大小设置为512。作者设置初始学习率为5e-4。在600M模型训练了1550k步，4.2B模型训练了550k步之后，作者将学习率降低到1e-4以进行更细致的参数更新（关于此学习率降低的消融研究在附录A.3中提供）。

定量结果分析。评估结果列于表1中。TerDiT是针对DiT模型的QAT方案，因此在所有全精度模型中，作者特别关注DiT-XL/2（675M）和Large-DiT-4.2B。在没有分类器自由引导的情况下，TerDiT-4.2B与DiT-XL/2的测试结果非常相似（训练图像数量少得多）。在使用分类器自由引导的情况下