4-bit量化扩散模型,通过旋转感知保留关键激活,显著提升生成质量。
DiRotQ: Rotation-Aware Quantization for 4-bit Diffusion Transformers

- 用PCA识别激活主成分,仅对次要成分做4位量化。
- 在MJHQ-30K数据集上FID达15.9,PSNR达19.1dB,优于之前方法。
- 适用于追求高效推理的图像生成系统,尤其适合4位部署场景。
扩散变压器(DiTs)虽生成质量顶尖,但推理时内存与计算开销巨大。虽然将模型压缩至4位精度(PTQ)能大幅提升效率,却常导致质量严重下降。现有方法如平滑、混合精度、旋转和低秩残差等虽部分缓解此问题,仍与FP16/BF16性能存在明显差距。本文提出DiRotQ,一种针对4位激活的旋转感知量化框架。该方法通过主成分分析(PCA)识别主导激活变异的低秩子空间,保留该子空间系数为高精度,其余成分量化至4位。推理时使用校准获得的正交变换对激活进行旋转,逆旋转则离线融合至层权重中。结合基于GPTQ的权重量化,DiRotQ在PixArt-Σ模型上于MJHQ-30K数据集实现FID 15.9(越低越好)与PSNR 19.1 dB(越高越好),超越先前最优方案SVDQuant(FID 18.9,PSNR 17.6)。此外,首次引入基于视觉语言模型的评判协议,更全面评估感知质量与提示对齐性。系统层面,基于Triton开发自定义核函数,使12B FLUX.1-dev模型内存减少2.1倍,推理速度较BF16基线提升2.3倍,部署于24 GB RTX 4090 GPU。
原文摘要 · Abstract (English)
Diffusion Transformers (DiTs) achieve state-of-the-art image generation quality but incur substantial memory and computational costs at inference. While aggressive Post-Training Quantization (PTQ) to 4-bit precision offers significant efficiency gains, it typically results in severe quality degradation. Existing approaches, including smoothing-based methods, mixed-precision schemes, rotation techniques, and low-rank residual methods, partially mitigate this issue but still leave a noticeable gap to FP16/BF16 performance. In this work, we introduce DiRotQ, a W4A4 PTQ framework that mitigates this degradation through rotation-aware activation quantization. DiRotQ identifies a low-rank subspace capturing dominant activation variance via Principal Component Analysis (PCA), preserving coefficients in this subspace at higher precision while quantizing the remaining components to 4-bit. Activations are rotated into the PCA basis at inference time using calibration-derived orthogonal transformations, while the inverse rotation is fused into the layer weights offline. Combined with GPTQ-based weight quantization, DiRotQ achieves an FID (lower is better) of 15.9 and PSNR (higher is better) of 19.1 dB on PixArt-Σ over the MJHQ-30K dataset, outperforming the prior state-of-the-art SVDQuant (FID 18.9, PSNR 17.6) under the same INT W4A4 setting. Beyond standard metrics, we introduce a VLM-as-a-Judge evaluation protocol for diffusion model quantization, the first such evaluation in this setting, providing a more holistic assessment of perceptual quality and prompt alignment under aggressive compression. On the systems side, we implement a Triton-based custom kernel to enable efficient end-to-end inference, reducing memory usage of the 12B FLUX.1-dev model by 2.1x and delivering 2.3x speedup over the BF16 baseline, on a 24 GB RTX 4090 GPU.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。