[Paper Review] Extreme Compression of Large Language Models via Additive Quantization

Vage Egiazarian, Andrei Panferov+4 more|arXiv (Cornell University)|Jan 11, 2024

Topic Modeling4 citations

TL;DR

This paper introduces AQLM, a novel additive quantization method for extreme compression of large language models (LLMs) to 2–3 bits per parameter. By adapting instance-aware, input-adaptive codebook optimization and joint tuning across layer blocks, AQLM achieves state-of-the-art accuracy in the 2-bit regime while enabling fast GPU/CPU inference with up to 8x memory reduction and speedups of up to 3.05x on GPU and 4.07x on CPU.

ABSTRACT

The emergence of accurate open large language models (LLMs) has led to a race towards performant quantization techniques which can enable their execution on end-user devices. In this paper, we revisit the problem of "extreme" LLM compression-defined as targeting extremely low bit counts, such as 2 to 3 bits per parameter-from the point of view of classic methods in Multi-Codebook Quantization (MCQ). Our algorithm, called AQLM, generalizes the classic Additive Quantization (AQ) approach for information retrieval to advance the state-of-the-art in LLM compression, via two innovations: 1) learned additive quantization of weight matrices in input-adaptive fashion, and 2) joint optimization of codebook parameters across each transformer blocks. Broadly, AQLM is the first scheme that is Pareto optimal in terms of accuracy-vs-model-size when compressing to less than 3 bits per parameter, and significantly improves upon all known schemes in the extreme compression (2bit) regime. In addition, AQLM is practical: we provide fast GPU and CPU implementations of AQLM for token generation, which enable us to match or outperform optimized FP16 implementations for speed, while executing in a much smaller memory footprint.

Motivation & Objective

• To address the challenge of extreme LLM compression at 2–3 bits per parameter, where existing methods suffer from significant accuracy drops.
• To extend Multi-Codebook Quantization (MCQ) techniques, particularly Additive Quantization (AQ), to post-training quantization of LLM weights.
• To improve accuracy-vs-size Pareto optimality in low-bit regimes by incorporating input-aware calibration and joint block-level optimization.
• To enable practical deployment of compressed LLMs via optimized GPU and CPU inference kernels with minimal latency overhead.

Proposed method

• Proposes AQLM, a learned additive quantization framework that reformulates the classic AQ optimization problem to minimize layer output error under input token distributions.
• Introduces instance-aware, input-adaptive codebook learning by jointly optimizing codebook parameters across entire layer blocks using calibration data.
• Employs residual K-means initialization to accelerate convergence during codebook training.
• Uses a two-stage training process: initial codebook learning followed by calibration-based fine-tuning of codebook parameters.
• Implements efficient GPU and CPU kernels for AQLM inference, leveraging smaller 8-bit codebooks to improve cache utilization and reduce memory footprint.
• Supports homogeneous quantization format without hybrid sparse-quantized components, simplifying deployment.

Experimental results

Research questions

• RQ1 Can Multi-Codebook Quantization (MCQ) techniques be effectively adapted to post-training quantization of LLM weights for extreme compression?
• RQ2 Does instance-aware, input-adaptive codebook optimization significantly improve accuracy in 2-bit LLM quantization compared to standard direct quantization?
• RQ3 Can joint optimization across layer blocks enhance model performance while maintaining low bit-width compression?
• RQ4 Is AQLM practically viable for real-time inference on commodity hardware, given its computational complexity?
• RQ5 How does AQLM’s performance scale with calibration set size and codebook configuration (e.g., code width, number of codebooks)?

Key findings

• AQLM achieves state-of-the-art performance in the 2-bit regime, outperforming prior methods such as QuIP# on Llama 2 7B, 13B, and 70B models.
• On Llama 2 70B, AQLM achieves a perplexity of 7.98 on Wiki2 with 2.07 bits per parameter, significantly improving accuracy at extreme compression.
• AQLM enables up to 3.05x speedup over FP16 on GPU (Nvidia RTX 3090) and up to 4.07x speedup on CPU (Intel i9, 8 cores) for 70B model inference.
• The method reduces memory footprint by up to 8x compared to FP16, while maintaining or exceeding FP16 inference speed.
• Fine-tuning codebook parameters has the highest impact on accuracy, with calibration set size up to 4096 sequences yielding diminishing returns beyond 256.
• Using multiple smaller codebooks (e.g., 2×8-bit) improves GPU cache utilization and enables higher speedups, albeit with slight accuracy trade-offs.