KroQuant：基於克羅內克積結構的擴散Transformer高效後訓練量化方法
針對擴散Transformer（DiT）在W4A4低比特量化中因激活值異常值導致圖像質量嚴重下降的問題，本文提出KroQuant方法。傳統後訓練量化（PTQ）方案在計算效率與量化精度間難以兼顧：SmoothQuant計算快但精度受損，固定Hadamard變換精度高但在線計算成本高，全秩可逆變換精度最佳但引入巨大的矩陣乘法開銷。KroQuant創新性地採用學習到的克羅內克積結構可逆變換，對每個32元素的激活塊進行處理，其參數量少於每通道縮放方法。該方法利用小塊張量核心GEMM實現高效推理，在MI350 GPU上比SmoothQuant快14%。結合離線LoRaQ權重校準吸收殘餘誤差，KroQuant在PixArt-Σ、SANA和FLUX.1-schnell模型上，於W4A4（MXFP4e2）格式下，在MJHQ-30K和SDCI基準上生成的圖像質量優於SVDQuant和LoRaQ，顯著逼近FP參考模型，實現了效率與精度的雙重突破。
扩散Transformer（DiT）作为当前图像生成领域的核心架构，其计算复杂度与显存占用一直制约着端侧部署与高效推理。尽管后训练量化（PTQ）被视为降低模型成本的有效途径，但在将DiT量化至极低的W4A4（权重4位、激活4位）格式时，面临严峻挑战。核心痛点在于进入线性层的激活值往往包含极端异常值，这些异常值超出了4位浮点或整型格式的有效表示范围，导致量化误差急剧放大，进而使得生成图像出现严重伪影或细节丢失。虽然理论上可以通过对激活值应用可逆线性变换、对权重应用其逆变换来平滑分布，从而提升量化精度，但DiT结构中各模块间的归一化层强制要求该变换必须在每个去噪步骤中在线执行。因此，推理阶段的计算开销成为设计该变换时的硬性约束。
现有的解决方案陷入了两难境地：以SmoothQuant为代表的每通道缩放方法虽然计算廉价，但会改变通道的幅度分布，损害量化精度；固定结构的Hadamard变换能提供较好的量化效果，但需要较大的块尺寸，导致在线计算成本高昂；而学习得到的全秩d×d可逆变换虽然校准效果最佳，但每层每步都需要执行密集的矩阵乘法（GEMM），其计算开销在实际应用中几乎不可接受。本研究旨在打破这一效率与精度的权衡僵局，提出一种既轻量又高精度的量化方案。为解决上述困境，本文提出了KroQuant方法，其核心创新在于引入了一种基于克罗内克积结构（Kronecker-structured）的可逆变换，专门针对激活值的块局部特性进行优化。具体而言，KroQuant将激活值划分为每个包含32个元素的独立块，并对每个块应用学习到的克罗内克积结构可逆变换。这种结构设计极具巧思：它不仅在参数存储上极具优势，其参数量甚至少于传统的每通道缩放方法（SmoothQuant），从而减少了模型整体的存储负担。
在计算实现上，由于采用了块局部结构，该变换可以转化为小型的张量核心GEMM操作，而非庞大的稠密矩阵乘法。这种设计使得KroQuant能够充分利用现代GPU的张量核心硬件加速能力。在NVIDIA MI350 GPU上的实测数据显示，KroQuant的量化器内核执行速度比SmoothQuant内核快高达14%。此外，为了进一步弥补量化过程中产生的残余误差，特别是针对权重的量化偏差，本文还结合了离线LoRaQ（Low-Rank Activation Quantization）权重校准技术。通过在离线阶段对权重进行精细校准，KroQuant能够有效吸收由激活值变换和量化带来的分布偏移，确保最终量化模型在保持低比特精度的同时，最大化还原原始模型的表达能力。
在实验验证环节，研究团队在多个主流扩散Transformer模型上进行了全面评估，包括PixArt-Σ、SANA以及FLUX.1-schnell。所有实验均严格采用W4A4量化格式，具体为MXFP4e2（4位指数2位尾数的混合精度格式），以模拟极端的低比特场景。评估基准涵盖了MJHQ-30K和SDCI（Stable Diffusion Community Index）等公认的高质量图像生成评测数据集。关键结果表明，KroQuant生成的图像在视觉质量上显著优于现有的SVDQuant和LoRaQ方法，且更接近未量化的FP（全精度）参考模型。在MJHQ-30K和SDCI数据集上的定量指标显示，KroQuant在保持或提升图像质量的同时，并未引入额外的推理延迟。
消融实验进一步证实了克罗内克积结构的有效性：相较于全秩变换，块局部结构在几乎不损失精度的前提下大幅降低了计算复杂度；而相较于固定变换，学习到的结构能够自适应地捕捉不同层和不同数据分布下的异常值特性。这些结果不仅验证了KroQuant在精度上的优越性，也证明了其在计算效率上的可行性，特别是在资源受限的硬件环境中，其14%的速度优势具有极高的实用价值。从行业意义与潜在影响来看，KroQuant的提出为扩散模型的高效部署开辟了新的技术路径。随着AIGC应用从云端向边缘设备、移动端甚至嵌入式系统延伸，模型的大小与推理速度成为决定产品可行性的关键因素。KroQuant通过巧妙的矩阵结构创新，证明了无需牺牲大量精度即可实现极低比特量化，这为DiT架构在消费级硬件上的普及扫清了关键障碍。对于开源社区而言，该方法提供了一种可复现、高效的PTQ范式，有助于推动更多研究者关注低比特量化中的结构优化问题。在工业落地方面，KroQuant带来的14%推理加速与显著的存储节省，可直接转化为云服务成本的降低或终端设备电池续航的提升。此外，其结合离线校准的策略也为后续研究提供了借鉴，即通过解耦激活变换与权重校准，分别优化不同部分的量化误差。未来，这一思路可延伸至其他类型的Transformer架构或更极致的比特压缩场景，有望成为下一代高效生成式AI模型的标准组件之一，推动AI生成技术向更普惠、更高效的方向发展。