arXiv:2509.23582cs.CV2025-09中稿 · ICML被引 5

提出RobuQ框架,让DiT模型激活值低至平均2比特仍稳定生成高质量图像

RobuQ: Pushing DiTs to W1.58A2 via Robust Activation Quantization

  • 设计鲁棒量化器,用哈达玛变换稳定激活分布,突破低比特限制
  • 实现激活值平均2比特的量化,首次在ImageNet-1K上保持竞争力
  • 首创仅量化激活的混合精度方案,适配大模型部署与边缘计算

扩散Transformer(DiTs)作为图像生成的新骨干架构,展现出优于U-Net的可扩展性与性能。但其实际部署受限于巨大的计算与内存开销。尽管量化感知训练(QAT)在U-Net中表现良好,应用于DiTs时面临激活敏感性与分布复杂性的挑战。本文识别出激活量化是推动DiTs进入极低比特设置的主要瓶颈。为此,提出面向DiTs的系统性QAT框架RobuQ:首先建立三值权重(W1.58A4)的DiT基线;进而提出稳健量化器(RobustQuantizer),理论证明哈达玛变换可将未知的逐令牌分布转化为正态分布,为该方法提供基础;此外提出首个仅激活混合精度网络(AMPN),全网使用三值权重,各层分配不同激活精度以消除信息瓶颈。在无条件与有条件图像生成任务上,RobuQ在亚4比特配置下达到当前最优性能。据我们所知,RobuQ是首个在ImageNet-1K等大规模数据集上实现激活值平均2比特量化且仍能稳定生成高质量图像的方案。代码与模型将开源。

原文摘要 · Abstract (English)

Diffusion Transformers (DiTs) have recently emerged as a powerful backbone for image generation, demonstrating superior scalability and performance over U-Net architectures. However, their practical deployment is hindered by substantial computational and memory costs. While Quantization-Aware Training (QAT) has shown promise for U-Nets, its application to DiTs faces unique challenges, primarily due to the sensitivity and distributional complexity of activations. In this work, we identify activation quantization as the primary bottleneck for pushing DiTs to extremely low-bit settings. To address this, we propose a systematic QAT framework for DiTs, named RobuQ. We start by establishing a strong ternary weight (W1.58A4) DiT baseline. Building upon this, we propose RobustQuantizer to achieve robust activation quantization. Our theoretical analyses show that the Hadamard transform can convert unknown per-token distributions into per-token normal distributions, providing a strong foundation for this method. Furthermore, we propose AMPN, the first Activation-only Mixed-Precision Network pipeline for DiTs. This method applies ternary weights across the entire network while allocating different activation precisions to each layer to eliminate information bottlenecks. Through extensive experiments on unconditional and conditional image generation, our RobuQ framework achieves state-of-the-art performance for DiT quantization in sub-4-bit quantization configuration. To the best of our knowledge, RobuQ is the first achieving stable and competitive image generation on large datasets like ImageNet-1K with activations quantized to average 2 bits. The code and models will be available at https://github.com/racoonykc/RobuQ .

模型量化扩散模型低比特推理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。