OrbitQuant:面向图像和视频扩散变换器的数据无关量化方法 OrbitQuant: Data-Agnostic Quantization for Image and Video Diffusion Transformers
通过旋转基共享 Lloyd-Max 码本实现数据无关的 DiT 低比特量化,无需校准数据即可跨模态应用
前置知识
后训练量化(PTQ)
后训练量化是指在模型训练完成后,不重新训练模型,直接将模型权重和激活值从高精度(如 FP16)转换为低比特表示(如 INT4)的技术。核心挑战是如何在不显著降低模型性能的前提下,减少量化误差。传统方法通常需要在少量校准数据上收集激活统计信息来确定量化范围和缩放因子。
本文聚焦于 PTQ 技术,理解其基本原理和传统方法的局限性,才能体会 OrbitQuant 绕过范围校准的创新价值。
扩散变换器(DiT)
扩散变换器是用 Transformer 架构替代传统 U-Net 作为去噪骨干网络的扩散模型。与 LLM 不同,DiT 需要在多个顺序去噪时间步上重复评估,而且即使是单批次推理,计算也是 compute-bound 而非 memory-bound。DiT 的激活值会随着时间步、提示词和分类器自由引导(CFG)分支的不同而发生偏移。
DiT 的这些特性(多步采样、计算密集、激活动态范围变化)使其量化问题与 LLM 量化有本质区别,这正是本文需要解决的核心挑战。
Hadamard 旋转
Hadamard 旋转是一种正交变换,由 Walsh-Hadamard 矩阵实现。快速 Walsh-Hadamard 变换(FWHT)可以在 O(d log d) 时间内完成,比密集矩阵乘法的 O(d²) 快得多。Hadamard 变换在量化中的核心作用是重新分布激活值,使得异常值被分散到更多通道上,从而更容易均匀量化。
本文的 RPBH 旋转基于 Hadamard 变换,理解其效率优势(快速变换)和数学性质(正交性、重新分布异常值)对于理解 OrbitQuant 的高效性至关重要。
Lloyd-Max 算法
Lloyd-Max 算法是一种标量量化的最优码本设计方法,通过迭代优化量化码本和分割边界来最小化均方误差(MSE)。它从初始码本开始,交替执行两个步骤:第一步,给定当前码本,找到每个数据点最近的码字(量化);第二步,给定量化结果,更新每个码本为对应量化区域的均值。算法在码本收敛时停止。
OrbitQuant 使用 Lloyd-Max 算法在离线阶段针对固定边际分布预计算最优码本,这是整个方法能够在推理时完全数据无关的理论基础。
研究动机
现有的扩散变换器后训练量化方法面临严重的数据依赖问题。DiT 的激活值会在多个维度上发生偏移:不同去噪时间步(如 FLUX.1-dev 的 50 步)、不同提示词(无条件生成 vs 条件生成)、不同 CFG 分支(正向 vs 负向)都会导致激活值的统计特性剧烈变化。例如,在 FLUX.1-schnell 模型中,原始激活值分布与目标分布存在显著偏差(Kolmogorov-Smirnov 距离较大)。这意味着传统基于范围校准的方法(如 SmoothQuant、SVDQuant、AdaTSQ、ViDiT-Q)必须为每个新检查点、新分辨率或新模态重新收集校准数据并重新拟合,在实际部署中带来了巨大的工程负担。当量化到极低比特(如 W2A4)时,现有方法完全崩溃,生成质量接近零,如表 1 所示,SmoothQuant、QuaRot 和 ViDiT-Q 在 W2A4 设置下的 GenEval Overall 分数几乎全部为 0.001。
本文的目标是本文提出 OrbitQuant,目标是设计一种完全数据无关的权重-激活量化方法,能够在无需任何校准数据的情况下,对图像和视频扩散变换器进行有效的低比特量化。该方法应当在跨模态(图像到视频)、跨模型(FLUX.1、Z-Image-Turbo、Wan 2.1、CogVideoX)、跨时间步、跨提示词的所有场景下都能保持接近全精度的生成质量,特别是在极低比特设置(W2A4、W3A3)下仍能产生可用的生成结果。
与已有工作不同的是,本文的独特切入角度在于绕过范围估计这一量化传统步骤,转而将低比特 DiT 激活量化问题重新表述为分布码本问题。与 LLM 量化中旋转异常值或缩放到权重的方法不同,OrbitQuant 利用数学性质:经过随机旋转后,归一化激活的每个坐标都会收敛到一个固定的、已知的边际分布,该分布仅由维度决定,与输入无关。这使得可以在离线阶段针对这个固定分布构建单一 Lloyd-Max 码本,并在所有层、时间步、提示词和模态间共享该码本,完全消除对校准数据的需求。这是首个将 TurboQuant 的分布码本思想从 KV-cache 压缩迁移到 DiT 投影内部,并通过共享旋转设计让权重和激活在同一基中量化的工作。
核心方法
OrbitQuant 的整体思路是:用分布性量化器替代按输入范围校准的传统方法,在共享、旋转、归一化的基中进行量化。由于权重和激活在同一基中量化,旋转在矩阵乘积中抵消,推理时只需对激活进行前向旋转。方法分为离线和在线两个阶段:离线阶段,针对每个维度构建随机置换块 Hadamard(RPBH)旋转和 Lloyd-Max 码本;对每个权重矩阵,用旋转后得到新权重,拆分为范数和方向,量化方向并重新附加范数,原地替换原权重。在线阶段,对每个输入激活,用前向旋转得到新激活,同样拆分并量化方向。由于旋转是随机采样且独立,单位和方向的坐标都遵循密度,因此针对该密度设计的 Lloyd-Max 码本是 MSE 最优的。
核心创新点是 RPBH 旋转的共享设计,它与现有方法的本质区别在于:(1) 使用均匀随机置换而非基于数据的校准置换(如 DuQuant 按异常值幅度排序、PeRQ 按每块质量平衡),证明了随机置换在概率上足以保证旋转边际保持良好行为;(2) 旋转被折叠到权重行中,在权重和激活的同一共享基中量化,旋转在每层线性变换内部抵消,无需运行时反向旋转;(3) 单一 Lloyd-Max 码本由固定的分布生成,服务所有维度的所有层、所有时间步、所有提示词,实现真正的数据无关量化。这使得 OrbitQuant 成为首个支持 W2A4(2-bit 权重、4-bit 激活)仍能产生可用图像的 PTQ 方法。
方法步骤详情
OrbitQuant 的完整方法步骤包括:(1) 离线初始化:对于所有目标维度集合中的每个维度,构建 RPBH 旋转,其中是 Walsh-Hadamard 矩阵,是 Rademacher 符号对角阵,是均匀随机置换矩阵;使用 Lloyd-Max 算法针对分布生成比特宽度的码本和。(2) 离线权重打补丁:对于每个目标线性层(输入维度为),计算;对每一行,计算范数和单位方向;使用码本量化方向得到;原地替换为。(3) 在线激活量化:对于每个输入激活(为 batch size),前向旋转;计算范数和单位方向(防止零范数);使用码本量化方向得到。量化层输出为。
技术新颖性
OrbitQuant 的技术新颖性体现在多个方面:(1) 首次将 DiT 低比特激活量化表述为分布码本问题,用单一 Lloyd-Max 码本替代按时间步范围校准,该码本适应所有去噪步骤;(2) 提出均匀随机置换的 RPBH 旋转,证明了随机置换在保持激活可量化性方面的充分性,并通过概率给出了方差集中界;(3) 首个在图像和视频 DiT 上完全免校准的 PTQ 方法,无需任何模型评估即可构建量化器;(4) 在 W2A4 极低比特设置下成为唯一仍能产生可用图像的 PTQ 方法,而现有方法完全崩溃;(5) 结构化旋转支持快速 Hadamard 变换内核,比密集 Haar 旋转快 26 倍(0.451s vs 11.65s 每图像),同时保持相同的量化质量。
实验结果
OrbitQuant 在多个图像和视频扩散变换器上取得了 SOTA 的 PTQ 性能。在图像生成任务(GenEval)上,W4A4 设置下 OrbitQuant 基本上是无损的:在 FLUX.1-schnell 上超越 FP16(Overall 0.703 vs 0.664),在 Z-Image-Turbo 上超越 FP16(Overall 0.767 vs 0.754),在 FLUX.1-dev 上略低于 FP16(Overall 0.633 vs 0.667),同时全面超越所有 PTQ 基线。W2A4 设置下优势更加明显,QuaRot、SmoothQuant、ViDiT-Q 等基线在所有三个骨干网络上崩溃到接近零分,而 OrbitQuant 在 FLUX.1-schnell 上保持 0.604 Overall,在 FLUX.1-dev 上保持 0.475,在 Z-Image-Turbo 上保持 0.319,成为唯一能产生有意义结果的方法。在视频生成任务(VBench)上,W4A6 设置下 OrbitQuant 在 Wan 2.1-1.3B 上以 24.35 Overall 超越所有 PTQ 基线(ViDiT-Q 19.58、SVDQuant 23.26),在 CogVideoX-2B 上以 24.55 Overall 达到最佳 PTQ 性能(ViDiT-Q 20.41)。W4A4 设置下 OrbitQuant 在 Wan 2.1-1.3B 上以 23.86 Overall 超越 SVDQuant 的 21.91,在 CogVideoX-2B 上以 23.86 匹配 SVDQuant 的 22.89。在极低比特(W3A3、W2A3)实验中,OrbitQuant 在 FLUX.1-schnell 上 W3A3 达到 0.678 Overall,W2A3 达到 0.517;在 FLUX.1-dev 上 W3A3 达到 0.584,W2A3 达到 0.372;在 Z-Image-Turbo 上 W3A3 达到 0.740(接近 FP16 的 0.754),而基线方法在 W2A3 上全部崩溃。在最大规模的 Wan 14B 上 W4A4,OrbitQuant 以 26.15 Overall 成为最佳 PTQ 方法,在 8 个维度中领先 7 个。种子鲁棒性实验显示 OrbitQuant 在 W4A4 下的 Overall 标准差最多 0.005,表明单次随机采样具有代表性。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 图像生成(GenEval Overall) | Overall Score | 0.703 (FLUX.1-schnell W4A4) | 0.680 (AdaTSQ) | +3.4% |
| 图像生成(GenEval Overall) | Overall Score | 0.604 (FLUX.1-schnell W2A4) | 0.001 (ViDiT-Q) | +60300% |
| 视频生成(VBench Overall Consistency) | Overall Consistency (%) | 24.35 (Wan 2.1-1.3B W4A6) | 19.58 (ViDiT-Q) | +24.4% |
| 视频生成(VBench Overall Consistency) | Overall Consistency (%) | 23.86 (Wan 2.1-1.3B W4A4) | 21.91 (SVDQuant) | +8.9% |
局限与改进
作者在讨论中承认了 OrbitQuant 的几个局限性:(1) 运行时旋转开销:与其他基于旋转的量化方法一样,OrbitQuant 需要在每次前向传播时对激活应用 RPBH 旋转,这会带来额外的计算开销,虽然作者测量显示在 H100 上每图像仅需 0.451 秒,但这是实现限制而非方法限制。(2) 缺乏原生低比特 GEMM 内核:由于 Lloyd-Max 质心是非均匀的,当前实现对量化码本进行反量化后在 BF16 中运行矩阵乘法(与所有基线的 fake quantization 设置一致),无法利用整数张量核心在均匀网格上的原生计算优势。(3) AdaLN 模态投影的限制:由于 AdaLN 的时变尺度和平移不能折叠到相邻权重中,作者将其保持为 INT4 权重量化而激活保持 BF16,这限制了压缩比从 4× 降到 2.21×。我观察到的额外限制包括:(4) 在极低比特(W2A3)下,Z-Image-Turbo 上的性能急剧下降到 0.105 Overall,表明分布码本方法在 2-bit 激活上达到极限;(5) 随机置换虽然理论上充分,但在特定极端情况下(当单个坐标携带过大的范数份额时),可能不够理想。
独立分析的弱点
独立分析的弱点包括:(1) 随机置换的次优性:在极端异常值集中场景下(如某些通道值远大于其他通道),均匀随机置换可能不如基于异常值幅度的校准置换(如 DuQuant)有效地分散质量。改进方向是研究自适应随机置换策略,在保持数据无关性的同时更好地处理极端异常值。(2) 旋转开销:虽然 RPBH 使用快速 Hadamard 变换,但在高分辨率图像生成(如 1024×1024)和长视频(81 帧)场景下,旋转成本仍然显著。改进方向是研究稀疏化或低秩近似旋转,或探索替代性的分布匹配方法。(3) AdaLN 量化瓶颈:当前方法对 AdaLN 投影使用 INT4 RTN,未能充分利用 OrbitQuant 的旋转优势。改进方向是研究如何将时变调制纳入旋转框架,或设计专门针对调制参数的量化策略。(4) 缺乏端到端量化感知训练:作为 PTQ 方法,OrbitQuant 无法像 QAT 方法那样通过微调模型来适应量化误差。改进方向是探索将 OrbitQuant 与轻量级 QAT 结合,通过少量梯度步骤进一步优化码本。
未来方向
作者提出的未来工作方向包括:(1) 开发非均匀权重量化的查找表 GEMM 内核:借鉴 Lut-GEMM 和非均匀权重量化的工作,构建融合内核,在 GEMM 前言中收集质心并在原生低比特格式中计算,真正实现端到端低比特推理。(2) 扩展到更多模态和架构:将 OrbitQuant 的分布码本思想迁移到其他生成模型(如 GANs、VAEs)和 Transformer 架构(如 Vision Transformers、音频模型)。(3) 理论分析完善:进一步研究随机置换的界限,特别是针对不同分布特性的激活值,给出更紧的收敛保证和失败模式分析。基于成果可延伸的方向包括:(4) 自适应码本设计:当前码本针对特定分布,可研究针对特定模型架构或层的码本微调,在保持数据无关性的同时提升精度。(5) 混合精度策略:根据层的敏感性自动分配比特宽度(如注意力层用 4-bit,前馈层用 3-bit),结合 OrbitQuant 的码本框架实现更精细的精度-性能权衡。(6) 分布式推理优化:研究如何将 RPBH 旋转和码本查找操作高效地映射到分布式推理系统,如张量并行和流水线并行。
复现评估
复现评估:作者在论文中提供了详细的方法描述和算法伪代码,项目页面承诺会开源代码。实验使用公开可用的模型(FLUX.1、Z-Image-Turbo、Wan 2.1、CogVideoX)和标准评估基准(GenEval、VBench),所有实验在 NVIDIA H100 GPU 上进行。模型设置明确:FLUX.1-schnell 使用 4 步和 guidance 0.0,FLUX.1-dev 使用 50 步和 guidance 3.5,Z-Image-Turbo 使用 10 步和 guidance 0.0;Wan 2.1-1.3B 使用 81 帧、480×832、50 步、CFG 5.0;CogVideoX-2B 使用 49 帧、480×720、50 步、CFG 6.0。作者提供了种子鲁棒性分析和旋转消融实验,显示方法在不同随机初始化下的稳定性。然而,论文未提供具体的超参数细节(如 Lloyd-Max 算法的收敛准则、值的敏感性分析),也未明确报告量化过程中的内存占用。复现难度中等:核心思想清晰但需要实现 RPBH 旋转和 Lloyd-Max 码本计算,熟悉量化研究的工程师可以在 1-2 周内完成复现。算力需求较高:需要在 H100 等高端 GPU 上运行大模型(如 Wan 14B)以复现完整实验,但可以在较小模型上验证核心思想。
论文图表
该图包含两个面板:左面板报告 AdaLN 调制权重比特降低时的 GenEval Overall;右面板报告 FLUX 架构在 AdaLN 权重为 BF16(2.21 倍压缩)和 INT4(4 倍压缩)时的模型压缩比。结果显示将 AdaLN 权重量化为 INT4 几乎匹配所有三个模型上的 BF16 结果,进一步降低到 W3 仍保持,但在 W2 时 FLUX.1-dev 和 -schnell 崩溃,而 Z-Image-Turbo 保持鲁棒。
这张图解释了 OrbitQuant 的设计选择:为什么将 AdaLN 保持为 INT4 而非进一步量化,平衡了压缩比(4 倍 vs 2.21 倍)和生成质量。