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

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 Activation-only Mixed-Precision Network 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.

Diffusion Transformers are powerful image generation models, but their high computational cost makes deployment difficult. We started this project because, on challenging datasets such as ImageNet, existing quantization methods still struggle to reduce activation precision while maintaining good generation quality. In particular, although low-bit weights have become increasingly feasible, activation quantization remains the main bottleneck for pushing Diffusion Transformers to extremely low precision. To address this problem, we propose RobuQ, a framework focused on robust low-bit activation quantization for Diffusion Transformers. Our method uses the Hadamard transform to make activation distributions easier to quantize and designs an efficient quantizer based on this transformed distribution. We further introduce mixed activation precision to allocate more bits only to sensitive layers, removing activation bottlenecks without increasing the whole model precision. With these techniques, RobuQ pushes Diffusion Transformers to W1.58A2 precision, using ternary weights and about 2-bit activations on average, while maintaining stable ImageNet generation. This work shows that activation precision can be aggressively reduced in Diffusion Transformers without severe quality collapse, making low-bit generative models more practical for efficient deployment.