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

Post-training quantization (PTQ) of diffusion transformers (DiTs) to 4-bit precision (W4A4) is hindered by activation outliers that standard quantizers cannot represent. While existing methods like SmoothQuant use per-channel scaling to mitigate this, they lack the expressivity to fully resolve inter-channel correlations. Conversely, dense invertible transforms are too computationally expensive for the online inference requirements of DiTs. This paper introduces KroQuant to bridge this gap by providing a learned, computationally efficient activation transform.

KroQuant employs a block-diagonal invertible transform where each 32-element activation block is processed by a learned matrix. To keep this efficient, the authors parameterize each 32x32 block as a Kronecker product of five 2x2 unit-determinant LU-decomposed matrices. This structure is significantly more expressive than per-channel scaling, requires fewer parameters than dense matrices, and can be implemented as a high-performance tensor-core GEMM. The method is designed as a drop-in replacement for the activation-transform stage in existing weight-splitting pipelines like LoRaQ, allowing for offline weight calibration to further absorb residual errors.

KroQuant consistently produces image outputs closer to full-precision references than SVDQuant and standard LoRaQ across multiple models (PixArt-Σ, SANA, and FLUX.1-schnell). On the MJHQ-30K and SDCI benchmarks, it achieves superior LPIPS and FID scores, indicating better preservation of visual quality. Furthermore, the custom KroQuant quantizer kernel is up to 14% faster than the SmoothQuant kernel on an MI350 GPU, demonstrating that the increased expressivity does not come at the cost of inference latency.

As diffusion models scale to billions of parameters, efficient quantization is essential for practical deployment. KroQuant demonstrates that block-local, learned transforms can effectively handle the activation outliers that plague 4-bit DiT quantization without violating the strict online computation constraints imposed by the normalization layers between transformer blocks.