DiRotQ: Rotation-Aware Quantization for 4-bit Diffusion Transformers¶

会议: NeurIPS 2026
arXiv: 2605.16732
OpenReview: https://openreview.net/forum?id=2vvKmRgMuy
领域: 模型压缩
关键词: 扩散Transformer (DiT), 4比特量化 (W4A4 PTQ), 主成分分析 (PCA), 正交旋转, 硬件加速

一句话总结¶

DiRotQ 通过时间平均 PCA 将 DiT 激活解耦为高方差高精度子空间与残差子空间，并对残差子空间引入正交旋转均衡通道方差、离线融合逆变换至权重中，辅以 GPTQ 权重补偿与自研 Triton 算子，在 PixArt-Σ、FLUX.1-dev 和 SANA-1.6B 上实现了无损级 W4A4 极低比特推理与最高 9.2 倍端到端加速。

研究背景与动机¶

扩散 Transformer（Diffusion Transformers, DiT）凭借卓越的可扩展性与生成质量，已成为前沿图像、音频及多模态生成的基础骨干架构，但其多步去噪推断与自注意力机制带来的计算与显存开销极其高昂。后训练量化（PTQ）是无需重训练即可大幅压缩模型并降低推理延迟的关键途径。然而，扩散模型在生成过程中处于计算受限（compute-bound）状态，单纯的纯权重（weight-only）量化无法在硬件上直接调用低比特张量计算核心，权重在线反量化甚至会抵消提速效果；要真正释放硬件潜能，必须对权重和激活值同时实施极低比特量化（如 W4A4）。

在将精度压缩至 4 比特的过程中，DiT 面临着严重的激活异常值（outliers）与分布失衡问题。与大语言模型（LLM）中分布相对静态的激活不同，DiT 的激活展现出显著的跨去噪时间步漂移、通道间极值差异（最大值与中位数之比可达 80 倍以上）以及由条件引导（CFG）和 AdaLN 带来的动态波动。以 SVDQuant 为代表的既有工作尝试通过通道平滑将激活异常值转移至权重中，再通过 SVD 低秩分支补偿；然而此类方法本质上在权重空间操作，无法根本抑制激活本身的动态方差与通道偏斜。与此同时，经典 LLM 旋转量化方法因 DiT 中非线性算子阻断变换融合以及跨时间步剧烈波动的特性，长期被学界认为难以直接落地于扩散模型。

针对这一困境，本文发现激活异常值在统计上高度集中于少数方向，只要对激活空间进行正交基对齐，就能以极低的高精度开销捕获绝大部分异常方差。核心 idea：利用时间平均激活的主成分分析（PCA）将激活特征空间正交解耦为保留高精度的主子空间与残差子空间，结合随机正交旋转平抑残差通道方差并离线融合逆变换至静态权重，实现高精度、零显存膨胀的 W4A4 DiT 纯整型与浮点高效量化。

方法详解¶

整体框架¶

DiRotQ 的整体处理流水线涵盖离线标定基变换构建、静态权重吸收重构、以及在线激活低开销正交投影与双分支低比特推理。在离线阶段，算法利用少量标定样本统计跨时间步的激活二阶矩矩阵，经特征分解得到全局时间平均 PCA 基矩阵 $\mathbf{U} = [\mathbf{U}_h \;\; \mathbf{U}_l]$，其中 $\mathbf{U}_h$ 对应方差最大的前 $k$ 个主方向，$\mathbf{U}_l$ 对应低方差残差子空间；同时在残差子空间构建独立于数据的随机正交旋转矩阵 $\mathbf{R}$。在推理阶段，激活向量仅需在线经历正投影变换 $\mathbf{V} = [\mathbf{U}_h \;\; \mathbf{U}_l\mathbf{R}]$，将激活拆分为保留 16 比特的高方差分支与经过随机旋转均化方差后量化为 4 比特的残差分支；逆变换则在离线阶段完全吸收进静态线性层权重中，辅以 GPTQ 补偿量化误差。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入激活 X 与静态权重 W"] --> B["时间平均 PCA 激活旋转与子空间解耦<br/>二阶矩分解分解为高方差 Uh 与残差 Ul"]
    B --> C["残差子空间随机正交旋转与静态权重融合<br/>正交矩阵 R 均衡残差方差，逆变换 (Ul R)T 离线吸收入权重"]
    C --> D["结合 GPTQ 的残差权重补偿与 Triton 高效算子实现<br/>双分支矩阵乘相加 Y = X Uh Wh + Q(Xl R) Q(Wl)"]
    D --> E["4-bit (W4A4) 低比特生成结果"]

关键设计¶

1. 基于时间平均的 PCA 激活旋转与子空间解耦：从激活源头捕获并隔离主方差异常值 针对 DiT 激活值通道间差异极大（部分通道方差远超中位数数十倍）且跨去噪时间步剧烈漂移的痛点，DiRotQ 直接在激活空间寻找最佳投影基底。在离线阶段，算法利用 128 条 COCO 文本在多去噪时间步下聚合各线性层的输入激活 $\mathbf{X} \in \mathbb{R}^{b \times m}$，计算非中心化经验二阶矩矩阵： $$\mathbf{\Sigma} = \frac{1}{N} \mathbf{X}^\top \mathbf{X} = \mathbf{U} \mathbf{\Lambda} \mathbf{U}^\top, \quad \mathbf{\Lambda} = \text{diag}(\lambda_1, \dots, \lambda_m), \quad \lambda_1 \ge \dots \ge \lambda_m$$ 由于 DiT 在 LayerNorm 后的激活均值接近于 0，$\mathbf{\Sigma}$ 良好近似于样本协方差矩阵。利用正交基 $\mathbf{U} = [\mathbf{U}_h \;\; \mathbf{U}_l]$ 旋转激活后，通道间完全去相关，各旋转分量的方差直接对应特征值 $\lambda_j$。由于均匀量化的期望均方误差正比于通道方差 $\lambda_j / 4^b$，DiRotQ 仅需提取方差占比绝对主导的前 $k = r \cdot m$（如 $r = 10\%$）个通道作为高方差子空间 $\hat{\mathbf{X}}_h = \mathbf{X}\mathbf{U}_h$，并以 FP16/BF16 原生精度保留；剩余 $(1-r)m$ 个残差通道 $\hat{\mathbf{X}}_l = \mathbf{X}\mathbf{U}_l$ 的最大/中位方差比从原本的 83 倍陡降至 16 倍，大幅减轻了后续 4 比特量化的动态范围压力。

2. 残差子空间随机正交旋转与静态权重融合：零推理开销平抑残差并吸收逆变换 虽然 PCA 解耦剔除了极端异常值，但在残差子空间 $\hat{\mathbf{X}}_l$ 中直接采用 4 比特量化时，有限的量化台阶（16 级）依然会对轻微的不均匀分布极其敏感。为解决该问题，DiRotQ 为残差空间引入无数据依赖的随机正交旋转矩阵 $\mathbf{R} \in \mathbb{R}^{(m-k) \times (m-k)}$（满足 $\mathbf{R}^\top \mathbf{R} = \mathbf{I}$），生成均化残差 $\tilde{\mathbf{X}}_l = \hat{\mathbf{X}}_l \mathbf{R}$。数学上，$\mathbf{U}_l \mathbf{R}$ 的各列依旧严格正交，使得合成投影基 $\mathbf{V} = [\mathbf{U}_h \;\; \mathbf{U}_l\mathbf{R}]$ 保持正交性质（$\mathbf{V}\mathbf{V}^\top = \mathbf{I}$）。旋转变换将残差通道的最大/中位数比进一步平抑至 2 倍，使得通道共享量化比例因子的效率达到最优化。更为关键的是，由于 $\mathbf{U}$ 与 $\mathbf{R}$ 在各时间步全局共享且线性层静态权重固定，逆旋转矩阵可在线下被数学等价地直接吸收入权重矩阵： $$\tilde{\mathbf{W}}_h = \mathbf{U}_h^\top \mathbf{W}, \quad \tilde{\mathbf{W}}_l = (\mathbf{U}_l \mathbf{R})^\top \mathbf{W}$$ 在推理阶段，线性层计算转化为双分支并行结构： $$\mathbf{Y} \approx [\mathbf{X}\mathbf{U}_h \;\; Q(\tilde{\mathbf{X}}_l)] \begin{bmatrix} \tilde{\mathbf{W}}_h \\ Q(\tilde{\mathbf{W}}_l) \end{bmatrix} = \mathbf{X}\mathbf{U}_h \tilde{\mathbf{W}}_h + Q(\mathbf{X}\mathbf{U}_l\mathbf{R}) Q(\tilde{\mathbf{W}}_l)$$ 逆变换在运行时开销彻底归零，完美规避了 DiT 中 AdaLN 等非线性算子对权重反向吸收的阻隔。

3. 结合 GPTQ 的残差权重补偿与 Triton 高效算子实现：全流程误差闭环与端到端硬件加速 在完成投影基融合后，残差权重 $\tilde{\mathbf{W}}_l$ 的量化精度直接决定了重构质量。DiRotQ 并没有简单采用就近舍入（RTN），而是利用激活二阶矩矩阵在残差权重空间执行 GPTQ（阻尼系数 $\lambda = 0.01$，分块大小 128），以二阶 Hessian 逆矩阵贪心更新剩余列权重，使得权重误差与残差激活误差产生协同补偿效应。在系统工程层面，由于正向旋转 $\mathbf{X}\mathbf{V}$ 需要在每次前向传播中在线执行，团队设计了自研 Triton 高性能 GPU 算子，实现高效的切片正交投影计算；低比特矩阵乘法复用了 Nunchaku 引擎的高吞吐 INT4/NVFP4 张量核核心，使高精度分支（占总权重仅 2% 显存开销）与低比特残差分支能够紧凑并行，彻底释放显存带宽与算力瓶颈。

损失函数 / 训练策略¶

DiRotQ 属于免训练的后训练量化（PTQ）方法，无需反向传播更新模型参数。标定过程严格基于 128 条来自 COCO Captions 的提示词进行，对 PixArt-Σ、FLUX.1-dev 和 SANA-1.6B 运行完整的去噪轨迹以收集各层激活二阶矩。算法在计算时间平均二阶矩矩阵 $\mathbf{\Sigma}$ 时加入 $\lambda = 0.01$ 的岭正则化阻尼项以确保特征分解的数值稳定性。对于 4 比特整型量化（INT W4A4），采用组大小为 64（group size 64）的分组非对称激活量化与对称权重量化；对于 4 比特浮点量化（NVFP W4A4），支持 Blackwell / Ada 架构下的 NVFP4 标准格式（组大小 16，FP8 缩放因子）。整个标定与量化流程在单张 NVIDIA A100 GPU 上即可快速完成。

实验关键数据¶

主实验¶

在代表高多样性艺术风格的 MJHQ-30K 和复杂细粒度描述的 sDCI 数据集上，DiRotQ 针对 3 款不同规模的 DiT 架构（PixArt-Σ 600M、FLUX.1-dev 12B、SANA-1.6B）与多种主流基线进行了全面评测。在极端 INT W4A4 与 NVFP W4A4 条件下，DiRotQ 的表现显著优于先前最先进的 SVDQuant 与 ViDiT-Q。

模型	量化精度	方法	MJHQ-30K FID (↓)	MJHQ-30K IR (↑)	MJHQ-30K PSNR (↑)	sDCI FID (↓)	sDCI PSNR (↑)
PixArt-Σ	FP16 (基线)	-	16.6	0.952	-	24.8	-
PixArt-Σ	INT W4A8	PTQ4DiT	39.8	0.572	12.8	47.1	11.0
PixArt-Σ	INT W4A8	Q-DiT	19.6	0.874	13.8	26.1	13.0
PixArt-Σ	INT W4A4	ViDiT-Q	21.1	0.896	10.0	26.5	9.3
PixArt-Σ	INT W4A4	SVDQuant	18.9	0.871	17.6	25.8	16.4
PixArt-Σ	INT W4A4	DiRotQ (本文)	15.9	0.941	19.1	23.1	18.0
PixArt-Σ	NVFP W4A4	SVDQuant	16.8	0.933	18.6	23.2	17.4
PixArt-Σ	NVFP W4A4	DiRotQ (本文)	16.0	0.946	19.6	23.0	18.4
FLUX.1-dev	BF16 (基线)	-	19.9	0.939	-	25.0	-
FLUX.1-dev	INT W4A4	SVDQuant	19.7	0.920	21.0	24.7	19.2
FLUX.1-dev	INT W4A4	DiRotQ (本文)	19.8	0.934	24.6	24.9	22.8
FLUX.1-dev	NVFP W4A4	DiRotQ (本文)	19.9	0.932	24.4	25.0	22.7
SANA-1.6B	BF16 (基线)	-	16.0	1.100	-	22.8	-
SANA-1.6B	INT W4A4	SVDQuant	16.1	1.080	18.1	21.5	16.8
SANA-1.6B	INT W4A4	DiRotQ (本文)	16.0	1.100	23.5	22.9	21.4
SANA-1.6B	NVFP W4A4	DiRotQ (本文)	16.0	1.090	23.2	22.4	21.2

消融实验¶

在 PixArt-Σ 与 MJHQ-30K 基准下，对 PCA 子空间解耦、残差正交旋转及残差权重 GPTQ 补偿各组件的贡献进行了逐步消融实验：

精度设置	PCA 分解	残差旋转	GPTQ 权重优化	FID (↓)	Image Reward (↑)	LPIPS (↓)	PSNR (↑)	说明
FP16 基线	–	–	–	16.6	0.952	–	–	全精度浮点参考
INT W4A4	✗	✗	✗	40.3	0.512	0.603	13.0	朴素 RTN，激活异常值致性能崩溃
INT W4A4	✓	✗	✗	19.7	0.892	0.452	14.6	引入 PCA 隔离异常方差，FID 骤降
INT W4A4	✓	✓	✗	19.2	0.904	0.438	14.8	残差正交旋转均化方差，指标稳步提升
INT W4A4	✓	✓	✓	15.9	0.941	0.236	19.1	结合 GPTQ 形成完整闭环，超越 FP16
NVFP W4A4	✗	✗	✗	27.7	0.710	0.515	14.7	朴素 NVFP4 量化
NVFP W4A4	✓	✗	✗	16.7	0.899	0.335	17.1	PCA 解耦大幅降低 LPIPS
NVFP W4A4	✓	✓	✓	16.0	0.946	0.224	19.6	完整版本达到最佳像素与感知保真度

关键发现¶

PCA 旋转是精度恢复的核心推力：单纯引入 PCA 隔离前 10% 异常分量，即令 INT W4A4 的 FID 从 40.3 锐减至 19.7，各层激活 QSNR 普遍提升 10–20 dB，证明了对激活空间做正交基对齐能直击异常值的本质分布。
高精度尾部分数 $r$ 在 10% 处呈现显著饱和：扫参实验表明，当保留比例 $r$ 从 5% 增至 10% 时，激活 QSNR 获得 2–3 dB 的快速增长；而当 $r > 10\%$ 进一步增加至 25% 时，每增加 5% 带来的 QSNR 增益不足 0.5 dB。选取 $r=10\%$ 完美兼顾了生成品质与高精度计算/显存开销。
硬件显存与端到端延迟跨越式突破：在消费级 NVIDIA RTX 4090（24GB 显存）单卡上，针对 12B 的 FLUX.1-dev，DiRotQ 实现了 2.1 倍的显存压缩。在 Batch Size = 1 时实现每步 2.3 倍加速；在 Batch Size = 2 时，由于原生 BF16 显存超限必须依赖极慢的逐层 CPU Offloading，而 DiRotQ 依然能完全常驻显存，最终实现了惊人的 9.2 倍端到端加速。
VLM-as-a-Judge 全维一致占优：在首次引入的 InternVL2-8B 与 Qwen2-VL-7B 评测体系中，DiRotQ 在语义一致性（SC）与感知质量（PQ）两维度全面领先 SVDQuant，在所有 10 个图像子类别的两两盲测中均取得显著胜率优势。

亮点与洞察¶

激活空间直接对齐取代间接平滑：相比于传统 SmoothQuant 或 SVDQuant 将激活极值“推”入权重空间的折中做法，DiRotQ 直接在激活空间执行 PCA 正交坐标变换，从源头消解了通道与时间步异常，将异常值集中至极小的高精度子空间。
正交复合与逆变换离线融合的数学优雅性：利用 $\mathbf{U}_l$ 与随机正交矩阵 $\mathbf{R}$ 乘积依然严格正交的性质，将复杂的逆旋转变换直接乘入静态全连接权重，使运行时无任何逆投影开销，成功化解了 DiT 结构中 AdaLN 等算子对传统 LLM 旋转融合的阻碍。
首创 VLM 裁判协议评估扩散量化：率先探索了以多模态大模型作为裁判（VLM-as-a-Judge）对扩散生成量化损失进行量化打分，克服了 FID、LPIPS 对细微伪影和文本指令对齐漂移不敏感的弊端。

局限与展望¶

时间步共享投影基的表达折中：为使权重能够离线融合，当前方案在所有时间步共享同一套时间平均 PCA 变换矩阵；若未来探索更细粒度的按时间步分组投影，需要在额外的权重显存与更优的动态适应性之间做进一步权衡。
应用范围尚未拓展至视频扩散模型：当前方法主要验证于文本生成图像 DiT 模型；视频扩散 Transformer（如 Sora 类架构）涉及极其巨大的时空三维注意力激活漂移，将 DiRotQ 扩展至长时序视频模型仍需解决跨帧激活协同正交化的问题。

评分¶

新颖性: ⭐⭐⭐⭐☆ [巧用时间平均 PCA 与正交残差复合，破除扩散模型激活旋转融合的结构障碍]
实验充分度: ⭐⭐⭐⭐⭐ [覆盖 3 个主流 DiT、双数据集、传统像素/感知指标及全新 VLM-as-a-Judge 评测与真实单卡系统基准]
写作质量: ⭐⭐⭐⭐⭐ [数学推导严密、问题针对性极强、图表分析与消融论证完整详实]
价值: ⭐⭐⭐⭐⭐ [成功打通 W4A4 在 12B DiT 上的实用化推理，在单张 4090 实现 2.1 倍省存与最高 9.2 倍加速]