KroQuant: Kronecker结构块变换实现扩散变换器高效后训练量化

KroQuant: Kronecker-Structured Block Transforms for Efficient Post-Training Quantization of Diffusion Transformers

精选理由

这篇论文发了KroQuant,用Kronecker结构做激活变换,W4A4量化又快又准,比SVDQuant和LoRaQ结果好,还省参数。

AI 摘要

KroQuant提出一种Kronecker结构可逆变换,对每个32元素激活块进行在线量化,参数少于per-channel缩放的一半。在MI350 GPU上,KroQuant量化核比SmoothQuant核快14%。在PixArt-Σ、SANA和FLUX.1-schnell上以W4A4(MXFP4e2)量化,KroQuant在MJHQ-30K和SDCI上输出比SVDQuant和LoRaQ更接近FP参考,同时保持或提升图像质量。该方法通过离线LoRaQ权重校准吸收残差权重量化误差。

原文 · arXiv cs.LG

KroQuant: Kronecker-Structured Block Transforms for Efficient Post-Training Quantization of Diffusion Transformers

Post-training quantization (PTQ) of diffusion transformers (DiTs) to W4A4 severely degrades output quality, because activations entering each linear layer contain outliers that 4-bit formats cannot represent. The standard fix applies an invertible linear transform to the activations and its inverse to the weights before quantizing both. Normalization layers between blocks force this transform to run online at every denoising step, making its inference computation cost the binding design constraint. Existing options trade quantization quality for inference cost: per-channel scaling (SmoothQuant) is computationally cheap but impacts the magnitude of the channels, which can harm quantization accuracy; fixed Hadamard transforms yield better quantization accuracy but require large block sizes that incur a high online cost; learned full-$d$ invertible transforms calibrate best but entail a prohibitive dense $d \times d$ matrix multiplication (GEMM) per layer per step. We propose KroQuant, a PTQ method that applies a learned Kronecker-structured invertible transform to each 32-element block of the activation, storing less than half the parameters of per-channel scaling. The block-local structure runs as small tensor-core GEMMs, and on an MI350 GPU the KroQuant quantizer kernel is up to $14\%$ faster than the SmoothQuant kernel. Offline LoRaQ weight calibration then absorbs the residual per-weight quantization error. On PixArt-$Σ$, SANA, and FLUX.1-schnell at W4A4 (MXFP4e2), KroQuant produces outputs closer to the FP reference than SVDQuant and LoRaQ on MJHQ-30K and SDCI, while preserving or improving image quality.