如何量化文图/视频生成模型? 清华&Infinigence提出高效且准确的DiT量化方法ViDiT-Q

我们提出了ViDiT-Q, 一种面向Diffusion Transformer (DiTs)的后训练量化(PTQ)方法，对基于DiT的文生图/文生视频模型的量化做出了先行探索。

• 我们分析了现有Diffusion量化方案在处理DiT模型量化时的独特问题，概括出DiT模型量化的若干关键挑战，并针对这些挑战提出对应的解决方案，设计了针对DiT的量化方案ViDiT-Q。

• 针对更低比特量化，我们定位了“量化被“瓶颈”在若干敏感层”的关键问题。并基于文生图/视频的任务特性，提出了一种指标解耦的混合位宽分配方案。

• 我们在主流实验场景与模型上进行了广泛的评估。针对文图生成模型（PixArt-alpha & sigma），ViDiT-Q可实现数值指标与视觉效果完全无损的W8A8与W4A8量化。针对文视频生成模型 (Open-SORA) ，ViDiT-Q在W8A8时实现数值指标无损，在W4A8时无明显视觉损失。

论文标题： ViDiT-Q: Efficient and Accurate Quantization of Diffusion Transformers for Image and Video Generation 论文链接： https://arxiv.org/abs/2406.02540 代码链接： https://github.com/A-suozhang/ViDiT-Q 项目主页: https://a-suozhang.xyz/viditq.github.io/

一、前言

近年来，扩散模型(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的延迟优化。

二、方案概述

在设计DiT的量化算法时，我们发现了现有扩散模型量化方案面临着以下挑战：

• 挑战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）差异：扩散模型的推理过程涉及在多次的神经网络迭代推理，我们发现在不同时间步中，模型中的激活值存在着显著性的差异。

（4）通道维度（Channel）差异：对模型的权重与激活值，我们都发现了不同通道的激活值存在着显著的差异。

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

（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维度与时间步维度的差异问题。

（3）时间步感知的通道均衡（Timestep-aware Channel Balancing）：现有量化工作（如SmoothQuant）通常通过采用通道均衡的卷积技巧，来解决通道之间数据分布差异大的问题。通过引入一个逐通道的Mask：S，在权重上除以S并在激活值上乘以S，再进行量化。该过程将权重量化的难度转移到了激活值量化上以平衡二者的难度。当将该技巧应用在DiT量化时，我们发现其仍然不能弥补性能损失。我们进一步分析得到，激活值的逐通道数据分布，随着扩散模型的时间步变化明显。基于此观察，我们对通道均衡进行了“时间步感知”的改进，在不同的时间步阶段采用不同的通道均衡Mask（如下图所示）。

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

虽然ViDiT-Q能够有效的实现W8A8的无损量化，但面向更低位宽（W6A6，W4A8）的量化，量化后模型仍然存在着明显的生成质量损失。经过进一步分析，我们发现了更低比特下，模型量化失效的主要原因是：量化被少部分极端敏感层所“瓶颈”。如下图所示，W4A16量化会产生纯黑的图片，而仅仅保留某一层为全精度，就能够产生有内容的视频。

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

因此，考虑到视频生成任务的独特特点，我们针对视频评估的多方面：文本吻合度，视觉质量，与时间一致性，选取了对应的指标。经过分析得到，不同类型的层对不同方面的影响有着较强的相关性（如下热力图）。因此，我们提出了一种“指标解耦”的量化方法，对每种类型层，采用单独量化该层的特定指标劣化程度，来作为量化敏感性的衡量标准。

五、实验结果与分析

我们在多种评测场景下，对多种模型进行了全面的评估。

5.1 文生视频模型的全面评测Benchmark：VBench

我们在VBench上评测了我们量化后的OpenSORA（STDiT）模型，如下表与下图所示，ViDiT-Q量化后的模型，在多方面都与全精度浮点模型获得了类似的性能。ViDiT-Q-MP的混合精度模型，在W4A8时，取得了比基线量化方案W8A8显著更优的性能。

我们通过下述具体案例进一步分析ViDiT-Q量化模型的性能保持：

（1）“物体一致性（Subject Consistency）”：基线量化方法所生成的视频中，熊的耳朵在视频中途突然出现了，未能完全保持物体的一致。

（2）“时间一致性”：基线量化方法所生成的视频中，镜头未能保持恒定，视频中的楼房快速变化且跳动，未能保持视频的时间一致性。此外，楼房本身也

（3）“视觉效果”：基线量化方法所产生的视频相比全精度模型，出现了明显的色差，且生成图像存在着明显的抖动。

5.2 视频生成模型的多方面指标：UCF-101数据集

我们在UCF-101数据集上测试了STDiT与Latte模型的量化结果，如下表展示，ViDiT-Q在各数据位宽上都展示了显著更优的生成质量。

5.3 文生视频与基线量化方法的对比

我们基于Open-SORA的STDiT模型，在多种位宽下，进行了多种量化方案的对比。如下表所示，各种基线量化方法都难以单独解决视频模型的量化问题，ViDiT-Q在各种位宽下都取得了更优且与全精度模型类似的性能。

我们给出以下具体的视频样例来展示ViDiT-Q与现有量化方案的视觉效果对比：

5.4 文生图模型的对比

我们将ViDiT-Q方案应用于主流文生图DiT模型（Pixart-alpha/sigma），如下表与下图所示，基线量化方法只能产生模糊的图片，几乎难以辨识物体。而ViDiT-Q在W8A8与W4A8量化都能产生和全精度模型几乎一样的图片。

5.5 硬件效率的提升

我们在Nvidia A100 GPU上进行了实际硬件Profiling评估。由于目前缺乏开源的支持动态量化的INT GPU Kernel，我们参考了静态INT GPU Kernel的执行速度，并考虑了动态量化的额外开销以进行估计。ViDiT-Q量化后的模型能够取得2-3x的内存优化，与约1.47x的延迟优化。我们正在实现兼容ViDiT-Q方案的高效硬件GPU算子，并计划开源（欢迎持续关注我们的项目主页： https://a-suozhang.xyz/viditq.github.io/ ）。

5.6 消融实验

我们进行了消融实验以展示ViDiT-Q的各技巧的作用，我们选取了W4A8这一较为具有挑战性的场景。如下表与下图所示，引入各技巧后，量化后模型产生的视频获得了不同程度的性能提升。

六、总结与未来指引

本文提出了ViDiT-Q，对视觉生成Diffusion transformer模型的低比特量化进行了先行探索。我们系统分析了DiT模型在量化上的独特挑战，并提出对应的解决方案。并面向更低比特量化，结合视觉生成任务的任务特性，提出了指标解耦的混合比特量化方案。ViDiT-Q在广泛的评估场景（全面Benchmark，多方面指标），任务（文生图，文生视频），与主流开源模型上取得了优异的生成质量。

llustration From IconScout By Delesign Graphics

-The End-

扫码观看！

本周上新！

“AI技术流”原创投稿计划

TechBeat是由将门创投建立的AI学习社区（

www.techbeat.net

投稿内容

// 最新技术解读/系统性知识分享 //

// 前沿资讯解说/心得经历讲述 //

投稿须知

稿件需要为原创文章，并标明作者信息。

我们会选择部分在深度技术解析及科研心得方向，对用户启发更大的文章，做原创性内容奖励

投稿方式

发送邮件到

chenhongyuan@thejiangmen.com

或添加工作人员微信（chemn493）投稿，沟通投稿详情；还可以关注“将门创投”公众号，后台回复“投稿”二字，获得投稿说明。

关于我“门”

将门是一家以专注于数智核心科技领域的新型创投机构，也是北京市标杆型孵化器。 公司致力于通过连接技术与商业，发掘和培育具有全球影响力的科技创新企业，推动企业创新发展与产业升级。

将门成立于2015年底，创始团队由微软创投在中国的创始团队原班人马构建而成，曾为微软优选和深度孵化了126家创新的技术型创业公司。

如果您是技术领域的初创企业，不仅想获得投资，还希望获得一系列持续性、有价值的投后服务，欢迎发送或者推荐项目给我“门”:

bp@thejiangmen.com

点击右上角，把文章分享到朋友圈