Command Palette
Search for a command to run...
OrbitQuant:面向图像与视频扩散Transformer的数据无关量化方法
OrbitQuant:面向图像与视频扩散Transformer的数据无关量化方法
Donghyun Lee Jitesh Chavan Duy Nguyen Sam Huang Liming Jiang Priyadarshini Panda Timo Mertens Saurabh Shukla
摘要
扩散Transformer(DiTs)在图像和视频生成领域达到了最优水平,但其多步采样和不断增长的参数量使得推理代价高昂。训练后量化(PTQ)是解决问题的自然手段,然而DiT的激活值会随时间步、提示词和引导分支发生变化,迫使先前的方法需要为每个新的模型检查点或模态重新拟合校准数据。我们提出了OrbitQuant,一种数据无关的权重-激活量化器,通过在归一化、旋转后的基中进行量化来绕过范围估计问题。在该基中,随机置换分块哈达玛(RPBH)旋转将每个坐标集中到一个固定的已知边缘分布附近,而与输入无关,因此对于给定输入维度,单个Lloyd--Max码本即可服务于所有时间步、提示词和层。我们将同一量化器离线扩展到权重行中,将旋转吸收到权重内,使其在每个线性层内部抵消,运行时仅保留对激活值的前向旋转。同一套方案从图像转移到视频时无需任何按模态调整。在FLUX.1、Z-Image-Turbo、Wan 2.1和CogVideoX上,该方法在多种低位宽设置下均达到了训练后量化的最优水平,并将图像扩散Transformer的训练后量化推进至W2A4。
一句话总结
Cantina Labs、USC 和 UIUC 的研究者提出 OrbitQuant,一种数据无关的权值-激活量化器,通过随机置换块哈达玛(RPBH)旋转将激活归一化到固定分布,使用跨时间步和模态的单一 Lloyd–Max 码本,将旋转离线吸收到权值中,并在图像和视频扩散 transformer(包括 FLUX.1、Z-Image-Turbo、Wan 2.1 和 CogVideoX)上实现了最先进的低比特训练后量化,同时将图像扩散 transformer 推至 W2A4 并保持可用的生成质量。
核心贡献
- OrbitQuant 引入了一种数据无关的权值-激活量化器,在归一化的旋转基下工作,其中单一 Lloyd–Max 码本服务于给定输入维度的所有时间步、提示和层。
- 旋转被离线吸收到权值中,在每个线性层内部抵消,运行时仅对激活进行前向 RPBH 旋转。RPBH 旋转的随机排列确保在低比特宽度下无需校准即可获得行为良好的旋转边缘分布。
- 在 FLUX.1、Z-Image-Turbo、Wan 2.1 和 CogVideoX 上,相同的配方从图像到视频无需按模态调整即可迁移。在多个低比特设置下,它在 GenEval 和 VBench 上设定了 PTQ 的最先进水平,并且在 W2A4 下,它是唯一能产生可用图像的方法,而其他 PTQ 基线则退化到噪声。
引言
扩散 Transformer(DiTs)目前驱动着最先进的图像和视频生成,但其计算受限的推理使得低比特权值-激活量化变得至关重要。现有的训练后量化(PTQ)方法继承自大型语言模型,依赖于校准数据来吸收异常值偏移;然而,DiT 激活会随去噪步骤、提示和分类器无关引导分支漂移,打破了静态统计假设,并需要为每个检查点或模态重新校准。本文提出 OrbitQuant,一种免校准的 PTQ 框架,将激活旋转到一个公共的归一化基中,其中它们的坐标遵循固定的、已知的边缘分布,与输入无关。然后,一个离线 Lloyd–Max 码本对所有去噪步骤的激活进行统一量化。旋转被折叠到权值行中,使其在每个线性层内抵消,仅在运行时留下轻量级的随机置换块哈达玛(RPBH)变换。OrbitQuant 在图像和视频 DiTs 上均取得了最先进的结果,无需任何校准数据,并且在 W2A4 比特宽度下仍能使用,而此前的方法已经崩溃。
方法
作者提出 OrbitQuant,一种免校准的量化框架,用基于分布的量化器替代逐输入范围校准,该量化器应用于共享的、旋转和归一化的基中。由于权值和激活在同一基中量化,旋转在矩阵乘积中抵消,仅在运行时对激活进行前向旋转。
该框架通过两个不同的阶段运行:离线权值量化和在线激活量化。对于一个输入维度为 d 的线性层,作者应用一个共享旋转 Πd。在离线阶段,权值矩阵被旋转到该基中,即 W′=WΠd⊤。然后,W′ 的每一行被分解为幅度 ri′=∥wi′∥2 和单位方向 w~i′=wi′/ri′。方向使用 Lloyd-Max 码本进行量化,幅度重新附加,得到 W^′=diag(r′)⋅Q^bw(d)(W~′)。行范数向量以 BF16 存储,开销可忽略不计。
在推理时,每个输入激活 x 在进入层之前先通过 Πd 旋转,产生 x′=Πdx。这被分解为幅度 s=∥x′∥2 和单位方向 x~′=x′/(s+ε)。方向使用 Lloyd-Max 量化器 Q^ba(d) 量化,并由 s 重新缩放,得到 x^′=s⋅Q^ba(d)(x~′)。由于权值吸收了 Πd⊤ 而激活应用了 Πd,两者在乘积中抵消,允许量化层计算 W^′x^′≈Wx,而无需在运行时进行逆旋转。
为了使旋转高效,作者将 Πd 实现为随机置换块哈达玛(RPBH)变换。其定义为 Πd=blkdiag(HhD1,…,HhDd/h)⋅Pπ,其中 Hh 是 Walsh-Hadamard 矩阵,Di 是 Rademacher 符号对角矩阵,Pπ 是均匀随机置换矩阵。前导置换 Pπ 至关重要,因为它将坐标分散到各个块,确保 Πdx~ 每个坐标的方差接近于 1/d。这使边缘分布接近 N(0,1/d),使得 Lloyd-Max 码本接近最优。
由于旋转后的坐标始终遵循固定的边缘分布 fd,作者离线对 fd 运行 Lloyd-Max,以获得每个维度的单一码本 Cd。该数据无关的码本服务于维度 d 的每个时间步、提示、层和权值行,完全消除了对输入统计收集或校准的需求。如图所示,原始激活明显偏离 fd,但经过 RPBH 旋转后,权值和激活都像密集 Haar 旋转一样紧密匹配目标分布,使得单一码本能够有效适应所有操作数。
实验
在图像和视频扩散 transformer 上的评估显示,OrbitQuant 实现了最先进的训练后量化质量,具有近乎无损的 W4A4 性能,在 GenEval 和 VBench 上可超过全精度,并且具有稳健的 W2A4 行为,而基线则崩溃。OrbitQuant 在延迟比较中也表现出最低的量化开销,定性样本证实了即使在极低比特宽度下其保真度。消融研究支持使用随机置换块哈达玛旋转和 INT4 AdaLN 投影,以获得最优的精度-效率权衡。
OrbitQuant 在三个图像扩散 transformer 中的两个上,在 W4A4 下实现了无损或优于 FP16 的 GenEval 性能,在第三个上仅以微小差距落后,超越了所有训练后量化基线,设定了新的最先进水平。在 W2A4 下,优势更加明显:旋转和平滑方法在每个模型上都崩溃到接近零的分数,而 OrbitQuant 仍能正常运作,在 FLUX 模型上保留了大部分质量,并且是唯一在 Z-Image-Turbo 上产生有意义结果的方法。对于 FLUX.1-schnell 和 Z-Image-Turbo,OrbitQuant 在 W4A4 下的 GenEval 总体分数超过 FP16,优于所有 PTQ 基线。在 W2A4 下,竞争方法退化到接近零的分数,而 OrbitQuant 仍可使用,并且在 Z-Image-Turbo 上唯一提供了有意义的分数。
在视频生成基准测试中,OrbitQuant 在 W4A6 和 W4A4 精度下均优于其他训练后量化方法。它在 Wan 2.1-1.3B 和 CogVideoX-2B 上获得了最高的 Overall Consistency,在大多数单项质量维度上领先,而竞争方法在较低比特宽度下退化得更严重。在 Wan 2.1-1.3B 的 W4A6 下,OrbitQuant 获得了最佳的 Overall Consistency(24.35),领先于次优方法 SVDQuant(23.26)。在 Wan 2.1-1.3B 的 W4A6 下,OrbitQuant 在 Imaging Quality、Aesthetic Quality、Dynamic Degree、Background Consistency、Subject Consistency 和 Scene 上领先。Motion Smoothness 是 OrbitQuant 在 W4A6 下唯一不是最高指标的维度;它排名第二(97.76),仅次于 SmoothQuant(98.01)。在 W4A4 下,OrbitQuant 在大多数维度上最接近全精度,并在 Wan 2.1-1.3B 和 CogVideoX-2B 上再次在 Overall Consistency 上排名第一,而基线则失去优势。
在 W4A4 下,所有旋转方法的性能在噪声范围内,但在较低比特宽度下,RPBH 获得了最佳的 GenEval 分数,同时与其他快速 Hadamard 变换的速度相当。密集 Haar 旋转速度慢一个数量级,因为它无法利用快速内核。RPBH 在 W3A3 和 W2A4 下获得了最高的 GenEval Overall,优于 Block-RHT、Full RHT 和 Haar。移除随机置换会降低低比特性能(Block-RHT),表明置换在分散异常值以实现码本兼容性方面的作用。结构化 Hadamard 变换(Full RHT、Block-RHT、RPBH)的运行速度比密集 Haar 旋转快 25 倍以上,而 RPBH 仅比 Block-RHT 增加少量开销。
OrbitQuant 为扩散模型的训练后量化设定了新的最先进水平,在 4 比特权值和激活下实现了无损或优于 FP16 的图像质量,并且在 2 比特权值下仍是唯一能正常运作的方法,而所有基线都崩溃了。在视频生成方面,它在 W4A6 和 W4A4 精度下在总体一致性和大多数质量维度上领先。对旋转变换的消融表明,随机置换哈达玛(RPBH)通过有效分散异常值,获得了最佳的低比特性能,且与其它快速结构化变换相比仅增加了少量速度开销。