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

2026-07-23

AI Summary

This work addresses the severe degradation in generation quality observed in diffusion Transformers under 4-bit post-training quantization (W4A4), primarily caused by activation outliers, and the high inference overhead of existing linear transformation methods. To this end, the authors propose KroQuant, which introduces Kronecker-structured transformations into post-training quantization for the first time. KroQuant applies learned invertible transforms to every 32-dimensional activation block and integrates LoRaQ-based offline weight calibration with the MXFP4e2 numerical format. This approach significantly reduces both parameter count and computational cost while improving quantization accuracy. Experiments demonstrate that KroQuant achieves superior W4A4 image generation quality over SVDQuant and LoRaQ on PixArt-Σ, SANA, and FLUX.1-schnell, and its quantized kernel runs 14% faster than SmoothQuant on MI350 GPUs.

Abstract

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.

Problem

• post-training quantization
• diffusion transformers
• outlier activations
• inference cost
• low-bit quantization

Innovation

• Kronecker-structured transform
• post-training quantization
• diffusion transformers
• low-bit quantization
• efficient inference