提出KroQuant方法，高效实现扩散模型的4位量化，兼顾精度与速度。
KroQuant: Kronecker-Structured Block Transforms for Efficient Post-Training Quantization of Diffusion Transformers
- 用分块的克罗内克结构变换替代传统全矩阵变换，降低计算开销
- 在W4A4下图像质量优于SVDQuant和LoRaQ，保持或提升生成效果
- 适合追求高效率且需高保真度的扩散模型部署场景
扩散Transformer（DiTs）进行后训练量化（PTQ）至W4A4时，因线性层输入激活值含异常值，导致输出质量严重下降。现有方法通过可逆线性变换调整激活与权重，但归一化层迫使该变换在每个去噪步骤中在线执行，成为性能瓶颈。标准方案如通道级缩放（SmoothQuant）计算成本低但影响通道幅度，固定哈达玛变换精度高但块大小大、在线开销高；学习型全维变换校准最优，却需每层每步执行高成本密集矩阵乘法（GEMM）。本文提出KroQuant：对每个32元素激活块应用学习型克罗内克结构可逆变换，参数量少于通道级缩放一半。其局部结构以小规模张量核心GEMM运行，实测在MI350 GPU上比SmoothQuant快达14%。离线使用LoRaQ校准权重残差误差。在PixArt-Σ、SANA和FLUX.1-schnell上，采用W4A4（MXFP4e2）时，KroQuant在MJHQ-30K与SDCI数据集上的输出质量优于SVDQuant和LoRaQ，同时维持或提升图像质量。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。