DiRotQ: 扩散 Transformer 的旋转感知 4 比特量化¶
会议: NeurIPS 2026
arXiv: 2605.16732
领域: 模型压缩
关键词: 扩散 Transformer、训练后量化、W4A4 量化、PCA 旋转、VLM-as-a-Judge
一句话总结¶
针对扩散 Transformer（DiT）在激进 4 比特量化下因激活离群值与时间步剧烈漂移导致的严重画质恶化，DiRotQ 提出以 PCA 主成分旋转隔离高方差子空间保留高精度、残差子空间正交旋转均衡后 4 比特量化，并离线融合逆变换至静态权重，首次在 FLUX.1-dev 等模型上实现无损级别的 W4A4 极致加速与内存节省。
研究背景与动机¶
扩散 Transformer（Diffusion Transformers, DiT）凭借卓越的可扩展性与生成质量，在文生图、图像合成及多模态生成任务中占据了主导地位。然而，DiT 的自回归或多步迭代去噪过程伴随着巨大的计算量与显存占用，每一去噪时间步的自注意力计算更具平方复杂度瓶颈。训练后量化（Post-Training Quantization, PTQ）作为无需昂贵重训的压缩方案备受关注，且因为扩散模型即使在小批量下依然处于计算受限（compute-bound）状态，单纯的仅权重量化（weight-only）在推理时仍需将权重反量化为浮点，无法真正降低延迟，必须对权重与激活进行协同的 4 比特低位量化（W4A4）。
然而，在激进的 W4A4 设置下，现有量化方案遭遇了严重的质量瓶颈。与大语言模型（LLM）相对平稳的激活分布不同，DiT 存在两类极其特殊的破坏性分布特性：一是跨去噪时间步的动态剧烈漂移，二是通道间与空间 token 间存在极端离群值（最大值与中位数比值可高达 83 倍）。以往方案如 SVDQuant 试图通过平滑技术将激活离群值转移至权重并进行低秩分解，但这仅是在权重空间进行间接补偿，未能直接对齐激活分布的内在几何结构；而大模型中经典的旋转量化（如 QuaRot、SpinQuant）因 DiT 内部存在自适应归一化（AdaLN）、Softmax 及 GELU 等中间非线性层，无法将旋转完全跨层吸收进前序权重，且时间步动态变化导致朴素全局旋转失效。
面对这一矛盾，本文从激活空间的几何本质切入：与其被动平滑或跨层转移，不如在激活空间直接借助主成分分析（PCA）构建旋转基底，将主导方差能量紧凑汇聚于少量极高方差的低秩子空间（保留高精度），而让庞大的残差子空间经正交旋转均衡后安全降阶至 4 比特。核心 idea：利用基于校准样本跨时间步聚合构建的每层 PCA 正交基底，将激活投影分解为高方差保持分支（保留 16 比特）与正交均匀化的残差分支（量化为 4 比特），逆旋转离线吸收入层权重，实现端到端高保真 W4A4 推理。
方法详解¶
整体框架¶
DiRotQ 的整体处理流水线涵盖离线校准分解、静态权重融合吸收、以及在线高效双分支推理计算三个核心环节。在离线校准阶段，算法基于少量校准提示词生成多时间步去噪激活，通过计算未中心化二阶矩矩阵提取各线性层的 PCA 正交基底 \(\mathbf{U} = [\mathbf{U}_h, \mathbf{U}_l]\)，其中排名前 \(k\) 的主成分 \(\mathbf{U}_h\) 对应高方差子空间，其余分量 \(\mathbf{U}_l\) 对应残差低方差子空间。为消除残差维度内部的方差不均衡，进一步引入数据无关的随机正交矩阵 \(\mathbf{R}\) 对残差基底进行旋转均衡。在权重处理上，由于各层逆旋转 \((\mathbf{U}_l \mathbf{R})^\top\) 与 \(\mathbf{U}_h^\top\) 均完全静态，算法将它们离线融合进当前层的静态权重矩阵中，无需在线额外反向变换。在线推理阶段，输入激活经单次正交前向旋转映射至新基底，高精度分支与 4 比特量化残差分支分别与对应权重切片相乘并求和，整个流程如下图所示：
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入激活 X 与静态权重 W"] --> B["时间聚合 PCA 正交分解<br/>提取高方差 Uh 与残差基底 Ul"]
B --> C["残差正交旋转均衡<br/>构建复合正交基底 Ul·R"]
C --> D["权重离线吸收与 GPTQ 量化<br/>融合逆变换生成 Wh 与 Q(Wl)"]
D --> E["双分支在线前向推理<br/>高精度分支 XUh 与量化分支 Q(Xl)"]
E --> F["输出张量 Y = XUh Wh + Q(Xl) Q(Wl)"]关键设计¶
1. 时间聚合 PCA 正交分解：在激活空间分离极值离群能量 针对 DiT 激活张量在不同去噪时间步及通道间存在巨大动态方差的痛点，DiRotQ 放弃在权重侧间接修补的思路，直接在激活空间进行正交变换。利用来自校准集（如 128 条 COCO 提示词）在 20 个完整去噪时间步收集的各线性层输入激活 \(\mathbf{X}\)，构建经验未中心化二阶矩矩阵： $\(\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 \geq \dots \geq \lambda_m\)$ 由于 DiT 中层归一化后的激活近乎零均值，\(\mathbf{\Sigma}\) 良好近似协方差。通过特征分解得到按方差降序排列的基底 \(\mathbf{U} = [\mathbf{U}_h, \mathbf{U}_l]\)。前 \(k = r \cdot m\)（其中 \(r \approx 10\%\)）个高方差分量构成高精度子空间 \(\hat{\mathbf{X}}_h = \mathbf{X} \mathbf{U}_h\)，其激活方差为 \(\lambda_1, \dots, \lambda_k\)。根据均匀量化误差理论，单通道均方误差与通道方差正比并反比于 \(4^b\)。将最陡峭的极端离群能量（原空间中 max/median 达 83 倍）收拢进高精度分支（16 比特存储与计算），能从根源上截断绝大部分绝对量化误差，使剩余通道的方差比骤降至 16 倍左右。
2. 残差正交旋转均衡：消除低位量化中的通道非均匀失真 分离出高方差成分后，剩余的 \((m-k)\) 维残差子空间 \(\hat{\mathbf{X}}_l = \mathbf{X} \mathbf{U}_l\) 虽然幅值整体大幅下降，但通道之间仍然存在不可忽视的相对方差差异。若直接对其施加每组共享缩放尺度的 4 比特量化，局部的次级离群值仍会导致该组量化分辨率受损。为此，DiRotQ 引入一个与数据无关的随机正交矩阵 \(\mathbf{R} \in \mathbb{R}^{(m-k) \times (m-k)}\)，将残差特征旋转为 \(\tilde{\mathbf{X}}_l = \hat{\mathbf{X}}_l \mathbf{R}\)。由于 \(\mathbf{U}_l\) 具有标准正交列且 \(\mathbf{R}\) 为正交方阵，二者的复合乘积 \(\mathbf{U}_l \mathbf{R}\) 严格保持正交性，即 \((\mathbf{U}_l \mathbf{R})^\top (\mathbf{U}_l \mathbf{R}) = \mathbf{I}\)。正交变换在完全保持残差信号能量与内积的前提下，将方差均匀发散至所有残差通道，把残差空间内的最大值与中位数比值从 16 倍极致压低至 2 倍，从而极大地契合了基于组（Group Size 64 或 16）共享缩放系数的非对称 INT4 或 NVFP4 量化器。
3. 静态权重逆变换离线融合与 GPTQ 协同：实现零在线额外开销 在包含 AdaLN 等动态条件调制层的 DiT 中，激活旋转无法像 LLM 那样平滑吸收入前序线性层。DiRotQ 针对线性变换 \(\mathbf{Y} = \mathbf{X} \mathbf{W}\) 构造恒等正交分解： $\(\mathbf{Y} = \mathbf{X} \mathbf{V} \mathbf{V}^\top \mathbf{W} = [\mathbf{X}\mathbf{U}_h \quad \tilde{\mathbf{X}}_l] \begin{bmatrix} \mathbf{U}_h^\top \mathbf{W} \\ (\mathbf{U}_l \mathbf{R})^\top \mathbf{W} \end{bmatrix} = [\mathbf{X}\mathbf{U}_h \quad \tilde{\mathbf{X}}_l] \begin{bmatrix} \tilde{\mathbf{W}}_h \\ \tilde{\mathbf{W}}_l \end{bmatrix}\)$ 因为 \(\mathbf{U}_h\)、\(\mathbf{U}_l\) 以及 \(\mathbf{R}\) 在校准后皆为静态常量，其逆变换矩阵 \(\mathbf{U}_h^\top\) 与 \((\mathbf{U}_l \mathbf{R})^\top\) 可以在模型部署前直接与静态权重 \(\mathbf{W}\) 相乘，融合成新的高精权重切片 \(\tilde{\mathbf{W}}_h \in \mathbb{R}^{k \times n}\) 与残差权重切片 \(\tilde{\mathbf{W}}_l \in \mathbb{R}^{(m-k) \times n}\)。融合后的残差权重 \(\tilde{\mathbf{W}}_l\) 进一步借助 GPTQ（设置阻尼系数 \(\lambda=0.01\)，块大小 128）在校准激活加权下进行二阶误差补偿量化为 4 比特 \(Q(\tilde{\mathbf{W}}_l)\)。在推理阶段，仅需通过定制 Triton 算子执行一次激活的前向正交旋转 \(\mathbf{X} \mathbf{V}\)，随后分别输入轻量高精度矩阵乘与 4 比特 INT4/NVFP4 张量核心加速矩阵乘，逆变换的在线计算开销严格为零。
实验关键数据¶
主实验¶
在文生图生成质量评测中，针对 PixArt-\(\Sigma\)（600M）、FLUX.1-dev（12B）和 SANA-1.6B 三大前沿 DiT 模型，在 MJHQ-30K 与 sDCI（5K 评测子集）上开展测试。主要对比指标包括图像分布质量 FID（越低越好）、人类偏好对齐度 Image Reward（IR，越高越好）、感知相似度 LPIPS（越低越好）以及峰值信噪比 PSNR（越高越好）。
| 模型与步数 | 精度配置 | 方法 | MJHQ FID (↓) | MJHQ IR (↑) | MJHQ LPIPS (↓) | MJHQ PSNR (↑) | sDCI FID (↓) | sDCI PSNR (↑) |
|---|---|---|---|---|---|---|---|---|
| PixArt-\(\Sigma\) (20步) | FP16 | 基线 (16-bit) | 16.6 | 0.952 | – | – | 24.8 | – |
| INT W4A8 | PTQ4DiT | 39.8 | 0.572 | 0.531 | 12.8 | 47.1 | 11.0 | |
| INT W4A8 | Q-DiT | 19.6 | 0.874 | 0.500 | 13.8 | 26.1 | 13.0 | |
| INT W4A8 | DiRotQ | 16.5 | 0.942 | 0.222 | 19.5 | 24.5 | 18.4 | |
| INT W4A4 | ViDiT-Q | 21.1 | 0.896 | 0.698 | 10.0 | 26.5 | 9.3 | |
| INT W4A4 | SVDQuant | 18.9 | 0.871 | 0.319 | 17.6 | 25.8 | 16.4 | |
| INT W4A4 | DiRotQ | 15.9 | 0.941 | 0.236 | 19.1 | 23.1 | 18.0 | |
| NVFP W4A4 | SVDQuant | 16.8 | 0.933 | 0.264 | 18.6 | 23.2 | 17.4 | |
| NVFP W4A4 | DiRotQ | 16.0 | 0.946 | 0.224 | 19.6 | 23.0 | 18.4 | |
| FLUX.1-dev (25步) | BF16 | 基线 (16-bit) | 19.9 | 0.939 | – | – | 25.0 | – |
| INT W4A4 | ViDiT-Q | 36.3 | 0.382 | 0.672 | 13.6 | 32.8 | 12.6 | |
| INT W4A4 | SVDQuant | 19.7 | 0.920 | 0.241 | 21.0 | 24.7 | 19.2 | |
| INT W4A4 | DiRotQ | 19.8 | 0.934 | 0.139 | 24.6 | 24.9 | 22.8 | |
| NVFP W4A4 | SVDQuant | 19.9 | 0.930 | 0.180 | 22.9 | 24.8 | 21.2 | |
| NVFP W4A4 | DiRotQ | 19.9 | 0.932 | 0.144 | 24.4 | 25.0 | 22.7 | |
| SANA-1.6B (20步) | BF16 | 基线 (16-bit) | 16.0 | 1.10 | – | – | 22.8 | – |
| INT W4A4 | SVDQuant | 16.1 | 1.08 | 0.222 | 18.1 | 21.5 | 16.8 | |
| INT W4A4 | DiRotQ | 16.0 | 1.10 | 0.086 | 23.5 | 22.9 | 21.4 |
消融实验¶
在 PixArt-\(\Sigma\) 模型上基于 MJHQ-30K 进行各核心模块的累加消融实验（INT W4A4 与 NVFP W4A4 格式）：
| 量化精度 | PCA 高精分支 (\(r=10\%\)) | 残差正交旋转 \(\mathbf{R}\) | 权重 GPTQ 校准 | FID (↓) | IR (↑) | 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 骤降 20.6 |
| INT W4A4 | ✓ | ✓ | ✗ | 19.2 | 0.904 | 0.438 | 14.8 | 残差正交均衡进一步稳定量化 |
| INT W4A4 | ✓ | ✓ | ✓ | 15.9 | 0.941 | 0.236 | 19.1 | 结合 GPTQ 协同优化，反超全精度基线 |
| NVFP W4A4 | ✗ | ✗ | ✗ | 27.7 | 0.710 | 0.515 | 14.7 | NVFP4 朴素量化底座较好但仍失真 |
| NVFP W4A4 | ✓ | ✗ | ✗ | 16.7 | 0.899 | 0.335 | 17.1 | 加入 PCA 分支大幅收敛感知差异 |
| NVFP W4A4 | ✓ | ✓ | ✗ | 16.7 | 0.912 | 0.324 | 17.3 | 正交旋转均衡提升偏好对齐 IR |
| NVFP W4A4 | ✓ | ✓ | ✓ | 16.0 | 0.946 | 0.224 | 19.6 | 全配置达到 19.6 dB PSNR 极致保真 |
关键发现¶
- PCA 分解是质量提升的决定性杠杆：在 INT W4A4 下，仅引入 PCA 分解就将 FID 从 40.3 锐减至 19.7（提升超 50%），说明扩散模型在 4 比特下的退化绝大部分来自于少数高方差离群通道的截断溢出。
- 高精尾部比例在 10% 处达到最佳性价比：对高精子空间比例 \(r\) 在 5% 至 25% 的扫描显示，激活 QSNR 在 \(r=10\%\) 处产生拐点并饱和，继续增加到 25% 带来的增益每 5% 不足 0.5 dB，因此全局统一设定 \(r=10\%\) 可兼顾质量与计算开销。
- 系统端吞吐与显存实测大幅提升：在消费级 NVIDIA RTX 4090（24GB）显卡上部署 12B 的 FLUX.1-dev，DiRotQ 将显存占用由 23.8 GB 压减至 11.3 GB（节省 2.1×）；当 batch size 为 1 时，推理速度相比 BF16 提升 2.3×；当 batch size 为 2 时，BF16 必须依赖慢速 CPU offload，而 DiRotQ 仍能完全装入单卡显存，从而实现了 9.2× 的断崖式端到端延迟降低。
- 首创 VLM-as-a-Judge 评测协议：采用 Qwen2-VL-7B 与 InternVL2-8B 进行两轴（语义一致性 SC、感知质量 PQ，取几何均值）双盲自动化裁判，在 INT4 评测中 DiRotQ 相比 SVDQuant 在整体表现上获得 48.1% vs 39.9%（InternVL2）与 23.3% vs 16.2%（Qwen2-VL）的全面胜率。
亮点与洞察¶
- 空间正交分解直击本质：不同于以往方案在权重端兜圈子（如 SVDQuant 将离群值平滑推向权重后做 SVD），DiRotQ 在激活空间直接利用协方差特征分解，将极值能量剥离为 10% 高精度分支，精准匹配了数值敏感度分布。
- 逆变换离线融合规避 AdaLN 阻断：巧妙避开了 DiT 内部 AdaLN 等动态层无法像 LLM 那样进行全局算子融合的限制，通过保持正向旋转在线、将转置逆旋转离线乘入同层权重，以理论零推理代价化解了架构死结。
- 软硬协同的低比特系统闭环：不仅在算法上实现了 W4A4 数值保真，更基于 Triton 与 Nunchaku 实现了在线旋转与 INT4 矩阵乘的定制融合 Kernel，在真实单卡 RTX 4090 上跑通了 12B 大模型的极速无卸载生成。
局限与展望¶
- 去噪时间步共享基底的折衷：为了能将逆旋转离线吸收入静态权重，当前方案对所有去噪时间步强制共用同一套时间平均 PCA 基底；若未来需要根据去噪早晚期阶段动态切换基底，则需权衡多套权重的显存开销。
- 视频扩散模型的时空一致性挑战：当前评测主要集中在文生图领域，视频 DiT（如 Sora 类架构）在时间维度帧间相关性更强且潜在激活动态更广，旋转机制在视频时空注意力上的扩展仍待验证。
- 极小批量下的正向旋转访存开销：在线正向旋转 \(\mathbf{X} \mathbf{V}\) 虽然通过 Triton 进行了优化，但仍引入了额外的 GEMM/点积操作，在更极端的超轻量模型中其延迟占比可能会有所上升。
相关工作与启发¶
- vs SVDQuant：SVDQuant 依赖平滑技术将激活离群值推入权重后再做低秩 SVD 分解，属于权重空间的间接补偿，容易残留激活量化噪声；DiRotQ 直接在激活空间执行 PCA 旋转分解，在 PixArt-\(\Sigma\) 上将 INT W4A4 的 PSNR 从 17.6 dB 显著提升至 19.1 dB。
- vs QuaRot / SpinQuant：QuaRot 与 SpinQuant 依靠大语言模型中层与层之间可融合正交矩阵的特性消除离群值，但在 DiT 架构中由于 AdaLN 动态调节的阻隔无法直接照搬；DiRotQ 创新性地仅在线执行前向旋转、将逆旋转吸收入同层权重，成功将旋转量化范式迁移至扩散架构。
- vs Q-DiT / PTQ4DiT：早期 DiT 量化工作大多止步于 W4A8 或 W8A8，在面对计算受限的 W4A4 场景时会发生灾难性精度崩塌；DiRotQ 通过双正交分解机制，成为首个将 W4A4 DiT 性能推至可比肩浮点原模高度的实用方案。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 巧妙克服 DiT 动态归一化障碍，提出基于 PCA 的激活旋转与离线逆旋转融合机制，构思新颖严谨。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 3 大主流 DiT 架构、2 大基准数据集、全套传统指标与首创的 VLM 裁判协议，并提供单卡系统级加速实测。
- 写作质量: ⭐⭐⭐⭐⭐ 逻辑结构严密，数学公式与物理直觉衔接自然，图表翔实规范。
- 价值: ⭐⭐⭐⭐⭐ 为消费级 GPU 上低延迟、低显存部署超百亿参数扩散 Transformer 提供了可落地的坚实工业级方案。