Reproduce “1-Bit Wonder: Improving QAT Performance in the Low-Bit Regime through K-Means Quantization” on OpenResearch

About this paper

Sohir Maskey, Constantin Eichenberg, Johannes Messner, Douglas Orr

Aleph Alpha Research and Graphcore Research developed a K-means quantization-aware training approach for large language models, allowing 1-bit and 4-bit models to achieve competitive or superior generative performance compared to a BF16 baseline while operating within the same memory budget. This approach enabled a 31B 1-bit model to outperform a 4B BF16 model on most generative tasks, validating the benefit of increasing parameter count over precision for fixed memory.

Key insights

• K-means based nonlinear quantization consistently yielded lower loss and greater stability in the ultra-low-bit regime (e.g., 1-bit) compared to uniform integer quantization.
• Precision-aware scaling laws indicated that for a fixed inference memory budget, prioritizing a larger number of parameters at lower precision provides better performance than fewer parameters at higher precision.
• Nonlinear k-means formats provided higher effective capacity per memory bit than uniform formats, particularly for bit-widths up to 4.25 bits.

Summary by alphaXiv.