Balanced codebook refinement improves low-bit large language model quantization
BARQ: Balanced Codebook Refinement for Low-Bit LLM Quantization
Machine Learning
Summary
Large language models need a lot of computer memory and storage, which makes them hard to use efficiently. The authors address this by improving a technique called codebook-based quantization that shrinks the model size by assigning groups of model parameters to representative codes. They noticed that some codes don’t get used evenly, which limits the method’s effectiveness. To fix this, they propose BARQ, a method that balances how parameters are assigned to codes, improving the model’s accuracy at smaller sizes. Their method performs better than previous approaches when tested on multiple language models.
What this means in practice
- •For machine learning engineers: Reduce memory and storage requirements when deploying large language models on limited hardware without greatly sacrificing model quality.
- •For ai hardware developers: Design chips or hardware accelerators that support balanced codebook quantization schemes for improved efficiency in running language models.
Authors
Chenhang Cui, Xu Xie, Linrui Xu, Xiaohao Liu, Xingyu Zhu, Fei Shen, Tat-Seng Chua
Abstract
As large language models (LLMs) grow in parameter count, model storage and parameter memory traffic have become major bottlenecks to efficient deployment. Codebook-based weight quantization reduces these costs, but imbalanced nearest-codeword assignments during fitting can leave some codewords insufficiently updated, limiting effective codebook utilization. To address this limitation, we propose Balanced Assignment Refinement for Quantization (BARQ), which improves quantization quality through balanced fitting of existing codebooks. Specifically, we first compute joint soft assignments between weight blocks and codewords through entropically regularized optimal transport with uniform marginals and curvature-weighted reconstruction costs, ensuring equal positive fitting mass for every codeword in the exact solution. We then refine the codewords through an assignment-weighted barycentric update, which we prove minimizes the fitting objective for fixed assignments. For finite Sinkhorn iterations, the implemented update retains this optimality provided all codeword masses exceed the denominator floor. Finally, we discard the soft assignments and use the refined codebook for standard hard nearest-codeword encoding, with our analysis establishing sufficient conditions for reducing hard-quantization distortion and evaluation loss. Across multiple LLMs, BARQ achieves lower perplexity and higher mean zero-shot accuracy than the evaluated baselines at comparable bit budgets. The code is available at https://github.com/chenhangcuisg-code/BARQ.