ICML2026

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

Kaicheng Yang, Xun Zhang, Haotong Qin, Yucheng Lin, Kaisen Yang, Xianglong Yan, Yulun Zhang

5 citations

Abstract

Diffusion Transformers (DiTs) have emerged as a powerful backbone for image generation, offering superior scalability over U-Nets. However, their practical deployment is hindered by significant computational costs. While Quantization-Aware Training (QAT) shows promise, its application to DiTs is challenged by the high sensitivity and complex distributions of activations. Identifying activation quantization as the primary bottleneck for low-bit settings, we propose RobuQ , a systematic QAT framework. We first establish a strong ternary weight (W1.58A4) baseline. Building on this, we introduce RobustQuantizer , which utilizes the Hadamard transform to convert unknown per-token distributions into normal distributions. Furthermore, we propose AMPN , the first A ctivation-only M ixed- P recision N etwork pipeline, applying ternary weights globally while allocating layer-specific activation precisions to eliminate information bottlenecks. Extensive experiments demonstrate that RobuQ achieves state-of-the-art performance on ImageNet-1K , representing the first stable image generation with activations quantized to an average of 2 bits. Code is available at https://github.com/racoonykc/RobuQ.