Justquant improves 4 bit activation quantization without complex operations
JustQuant: You Don't Need Smoothing, SVD, or Rotation for 4-Bit Activation Quantization
Artificial Intelligence
Summary
Large generative AI models are powerful but require a lot of computing power. Reducing the precision of these models can speed them up but is tricky, especially for certain internal data called activations at 4-bit precision. The authors show that existing methods rely on complicated steps that make running these models harder on hardware. They propose JustQuant, a simpler approach that shifts complexity to the training phase instead of the running phase, using a new training technique called Theseus QAD. This approach improves model accuracy while avoiding difficult operations during use.
What this means in practice
- •For machine learning engineers: Streamline deployment of 4-bit quantized generative models by using simpler operators that are easier to support on standard hardware.
- •For hardware accelerator designers: Design accelerators without needing to support complex operators like smoothing or rotations for low-bit activation quantization.
Authors
Kaicheng Yang, Kaisen Yang, Chunyu Liu, Xianglong Yan, Haotong Qin, Junyi Wu, Tianao Zhang, Xun Zhang, Shaoqiu Zhang, Youbang Sun, Yulun Zhang
Abstract
Recent generative models have become increasingly powerful, but their inference cost continues to grow. Model quantization offers a promising way to compress these models and accelerate inference. However, at 4 bits, activation quantization is substantially more challenging than weight quantization. Recent post-training quantization (PTQ) and quantization-aware training (QAT) methods have made progress in 4-bit activation quantization by introducing smoothing, SVD branches, rotations, mixed precision, or advanced formats such as NVFP4. These additional operators and data types impose demanding requirements on inference engines and hardware, limiting the broad adoption of low-precision models. Can quantization be achieved using only plain low-bit operators? To answer this question, we propose JustQuant, a simple yet effective framework that moves the complexity of low-bit quantization from deployment-time operators into the training process. We first revisit model quantization from the perspective of knowledge distillation and show that a key reason existing PTQ and QAT methods fail is that they typically exploit supervision at only a single level. We then introduce Theseus QAD, a quantization-aware distillation method that progressively applies multi-level supervision, analogous to the gradual replacement process in the Ship of Theseus. Extensive experiments on DiT and diffusion large language models show two distinct regimes. For smaller models, Theseus QAD can serve as a lightweight warm-up stage that substantially improves subsequent QAT with plain operators, while naive QAD may collapse in the same setting. For larger models, Theseus QAD provides a stronger distillation training path than ordinary QAD. Across both regimes, JustQuant improves low-bit quantization quality while avoiding the complex operators required by many existing PTQ methods.