如何量化文图/视频生成模型? 清华\u0026Infinigence提出高效且准确的DiT量化方法ViDiT-Q(量化模型生成编辑器视频)

近年来，扩散模型(Diffusion Model)在视觉生成领域取得了显著的进展，Stable Diffusion模型能够依据文本信息生成高度拟真且美观的图像，OpenAI-SORA更是在视频生成任务上取得了惊人的视觉效果。
随着SORA的提出，扩散变换器（Diffusion Transformer，DiT) 模型受到了越来越多的关注，一系列基于DiT的模型在文生图（PixArt-alpha, PixArt-Sigma, Hunyuan-DiT），与文生视频 (OpenSORA，Vidu) 任务上展现了出色的生成能力。

然而，由于文生图大模型具有巨大的参数规模（Stable Diffusion XL: 3.5B, 35亿参数，PixArt-alpha：0.9B，9亿参数）与扩散模型循环迭代式的推理特点（单次生成图片/视频需要进行数十次的大模型推理），其运行的硬件资源消耗十分巨大，而视频生成需要同时生成多帧图像，进一步增加了模型的硬件开销，对其实际应用带来了巨大挑战。
例如，Open-SORA模型生成2s16帧的视频，大概需要消耗10余GB的GPU显存，在Nvidia A100 GPU上需要花费约1分钟。
这难以满足实际应用场景的效率要求。

低比特量化是一种被广泛使用的减少模型计算存储开销的方法，通过将原本高精度浮点（FP32/FP16）的模型全权重与激活值 (Weight and Activation, 简称W&A)，转化为低比特定点数（INT8/INT4），可以显著减少模型显存开销与计算复杂度。

为缓解基于DiT视觉生成模型的效率问题，来自清华大学电子工程系、无问芯穹、微软、和上海交通大学研究团队，对文图/文视频生成Diffusion Transformer的量化做出了先行探索，提出了一种新颖的扩散模型低比特量化方法：《ViDiT-Q: Efficient and Accurate Quantization of Diffusion Transformers for Image and Video Generation》这项工作中，研究人员分析了DiT量化的独特挑战，并针对性设计了解决方案。
现有量化方案在W8A8量化时劣化明显，难以生成符合文本的内容，在W4A8下只能生成模糊的色块。
而ViDiT-Q量化方案能够生成与全精度模型几乎相同的图片/视频。
在生成质量基本无损的前提下，ViDiT-Q能够获得2-3x的显存优化，与约1.5x的延迟优化。

二、方案概述

挑战1：在W8A8时会造成明显的视觉效果损失，在W4A8时只能产生单色的图片。
为解决这一挑战，我们针对DiT的模型与算法特性，设计了改进量化方案ViDiT-Q，能够实现无损的文生图模型的W8A8/W4A8量化，与文生视频的W8A8量化。

挑战2：但在文生视频的更低比特W6A6, W4A8时，仍然存在着图像质量的损失。
针对更低比特量化的新挑战，我们分析并定位了关键问题：更低比特量化时某些极端敏感层“瓶颈”住了量化性能。
考虑到文生视频任务的特性，我们提出了指标解耦的低比特量化方法。

三、考虑DiT模型特性：ViDiT-Q量化

首先，针对现有Diffusion Quantization方案在DiT量化时遭遇的挑战，我们通过分析数据分布探索其性能损失的原因。
我们将DiT量化的独特关键问题概括为：“在多个不同维度上存在显著的数据动态差异”，而现有的扩散模型量化方法大多采取了固定且粗粒度的量化参数，难以应对高度动态的数据变化。
具体的，我们将DiT模型中的数据差异概括为以下四个维度（如下图所示）：

（1）令牌维度（Token）的差异：在DiT中，激活值被表达为一系列视觉令牌（Visual Tokens，对于视频模型中还包含着时间令牌，Temporal Tokens），我们发现不同令牌的特征存在着显著的差异。

（2）控制信号维度（Classifier-free guidance, CFG）的差异：无分类器的控制信号引入（Classifier-free Guidance）是可控生成的一种主流范式，通过加权组合两次模型推理（一次带监督信号，一次不带监督信号）来实现可控的生成。
我们发现在有无监督信号时，模型中的激活值存在着显著的差异。

（3）时间步维度（Timestep）差异：扩散模型的推理过程涉及在多次的神经网络迭代推理，我们发现在不同时间步中，模型中的激活值存在着显著性的差异。

针对上述挑战，我们提出了以下针对性的量化方法改进：

（1）逐Token的量化参数确定（Token-wise Quantization）：DiT所采用的Transformer架构与CNN模型的关键区别在于，卷积涉及对局部像素的特征聚合，这些参与聚合的像素需要采用相同的量化参数。
因此，面向CNN的量化方法通常对整个激活值张量采用统一的量化参数（Tensor-wise quantization parameter）。
与此不同，DiT的特征聚合主要由Attention算子完成，而网络的主要计算开销为大量的线性（Linear）层，对Linear层来说，每个Token的计算是独立的，因此，可以采用逐Token的量化参数，来应对不同Token之间特征分布差异大的问题。
采用逐Token的量化参数，引入的额外存储开销仅为激活值张量的约千分之一，却能显著提升量化后模型的性能。

（2）动态量化参数（Dynamic Quantization）：在上述挑战中，CFG维度的差异与时间步维度的差异是扩散模型的两个特有问题。
受此前语言模型量化工作的启发，我们通过采用动态量化，即在线进行量化参数的统计，可以以少量的额外开销（小于LInear层计算过程约1%的延迟开销），自然的解决CFG维度与时间步维度的差异问题。

四、考虑视觉生成任务特性：指标解耦的混合位宽设计

为应对某些对量化极端敏感的层，一个直观的解决方案是采用混合位宽量化，对这些敏感层采用更高位宽。
然而，我们发现了与全精度模型输出的均方误差（Mean Squared Error, MSE）更大，并不一定意味着更差的生成质量。
如下图所示，左侧的视频具有着更大的MSE误差，然而，视觉效果比右侧的视频更好（右侧视频不能准确的生成文本描述中的“海鸥”，而产生了很多白色的碎片状物体）。
由于视频生成的效果需要从多方面评估，量化对多方面都会产生影响，直接采用MSE Erorr这样的数据层指标难以准确的评估量化的敏感性。

五、实验结果与分析