pQuant: Towards Effective Low-Bit Language Models via Decoupled Linear Quantization-Aware Training
Authors: Wenzheng Zhang, Bingzheng Liu, Yang Hu, Xiaoying Bai, Wentao Zhang, Bin Cui
Organizations: School of Computer Science, Peking University · College of Future Information Technology, Fudan University · Center for Information Research, Academy of Military Sciences · Advanced Institute of Big Data · Center of Machine Learning Research, Peking University
Quantization-Aware Training from scratch has emerged as a promising approach for building efficient large language models (LLMs) with extremely low-bit weights (sub 2-bit), which can offer substantial advantages for edge deployment. However, existing methods still fail to achieve satisfactory accuracy and scalability. In this work, we identify a parameter democratization effect as a key bottleneck: the sensitivity of all parameters becomes homogenized, severely limiting expressivity. To address this, we propose pQuant, a method that decouples parameters by splitting linear layers into two specialized branches: a dominant 1-bit branch for efficient computation and a compact high-precision branch dedicated to preserving the most sensitive parameters. Through tailored feature scaling, we explicitly guide the model to allocate sensitive parameters to the high-precision branch. Furthermore, we extend this branch into multiple, sparsely-activated experts, enabling efficient capacity scaling. Extensive experiments indicate our pQuant achieves state-of-the-art performance in extremely low-bit quantization.
Figures & tables
Figure 1: An overview of performance (Perplexity on WikiText2) and bit-width achieved by our pQuant and other extremely low-bit methods on 1.3B model.
Figure 2: Weight log-sensitivities in the final FFN layer of LLaMA3-3B and BitNet-3B. Matrices are down sampled via max pooling for visualization; darker blue indicates higher sensitivity. Red boxes highlight regions of peak sensitivity. Notably, in the 1-bit weights of BitNet-3B, no pronounced sensitivity variation is observed.
Figure 3: Computational flow of pQuant’s core modules. (a) pQuant replaces all linear layers with quantized counterparts, with 8-bit branch to enable dynamic scaling when needed. (b-d) Computation in three representative pQuant modules. FP16 weights are retained solely during training to ensure numerical stability and discarded post-training. Here, r denotes the dimension of weights in 8-bit branch, where r≪Dmodel .
Figure 4: Final training loss with varying numbers of parameters. pQuant with N=8 demonstrates excellent scalability, whereas 1-bit BitNet do not.
Figure 5: (a) Sensitivity analysis of the 1-bit and 8-bit branches in the down-projection layer of the final FFN block in the 700M pQuant model. (b) Ablation study of pQuant on 700M-parameter model, assessing the impact of feature scaling and the number of active 8-bit branches on final loss. The mid-training loss drop stems from the two-phase learning rate schedule (see Appendix B.2 ).
Figure 6: The memory footprint. Memory footprint of model weights transferred during a single forward pass in inference. pQuant is smaller than BitNet1.58 and significantly lower than LLaMA-2.
Figure 7: Perplexity of pQuant under different configurations. Lower perplexity indicates better performance. Left: Perplexity decreases as the number of 8-bit branches N in pQuant increases. Right: Comparison with alternative quantization methods. Native Mix retains 8% of parameters in high precision on top of 1-bit BitNet, rather than using our decoupled architecture. Channel-wise quantizes weights per output channel of the matrix multiplication. Group-wise quantizes weights in groups of 64 elements.
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: Computation time across components in a Transformer block, measured on an Apple M2 with a 7B model and input sequence length of 256. pQuant is 38% faster than BitNet1.58 and 82% faster than FP16 LLaMA-2.
Figure 9: Two-Stage Scheduler. During mid-training, both the learning rate (LR) and weight decay (WD) are decayed to accelerate convergence of 1-bit parameters and prevent sign flips caused by excessive weight oscillation around quantization thresholds.
Figure 10: Training loss curve of two model. pQuant exhibits greater training stability. In contrast, BitNet frequently suffers from gradient explosion during training, often requiring checkpoint reloading and restarts.
Figure 11: Illustration of the Quantization-Aware Training process. Before quantization, matrix multiplications are performed in 16-bit precision. During QAT, simulated quantization are inserted to enable the model to learn and compensate for quantization-induced errors. In the final inference phase, quantization factors are fused into the weights, only use the quantized weights.