Additive Quantization for Language Models (AQLM)

• The paper presents AQLM, which generalizes classic additive quantization to LLMs, achieving Pareto-optimal trade-offs at 2–3 bits per parameter.
• It employs input-adaptive code assignment and joint block-wise optimization via an EM-style process to minimize output distortion during calibration.
• The method enables up to 8× model size reduction with notable speedups on both GPU and CPU, facilitating efficient on-device inference.

Additive Quantization for LLMs (AQLM) is an advanced post-training compression technique developed for extreme quantization of LLMs, such as transformer-based architectures. AQLM generalizes the classic Additive Quantization (AQ) approach, traditionally used in information retrieval, to quantize weight matrices in LLMs to exceptionally low bit-counts—specifically targeting the 2 to 3 bits-per-parameter regime. By integrating input-adaptive code assignment and joint block-wise optimization of quantization parameters, AQLM achieves Pareto-optimal trade-offs between accuracy and model size, making it practical for deployment on resource-constrained devices (Egiazarian et al., 2024).

1. Formal Problem Statement

AQLM addresses the problem of compressing pretrained transformer LLMs by replacing the floating-point weight matrices $W \in \mathbb{R}^{d_{out} \times d_{in}}$ with quantized approximations $\hat{W}$ using only $B$ bits per parameter, with the principal focus on $B \approx 2\dots3$. This compression yields up to $8\times$ reduction in model size compared to FP16 baselines.

Classic AQ encodes groups of model weights as sums of $M$ codebook vectors (centroids) chosen from learned codebooks $\{C^{(m)}\}_{m=1}^M$, with assignments governed by one-hot vectors $b^{(m)}$. Row $w$ is approximated as $w \approx \sum_{m=1}^M C^{(m)} b^{(m)}$, and total bit cost is determined by codebook size and group granularity. The AQ layer-level reconstruction objective is:

$\hat{W}$0

AQLM reframes the objective to preserve layer outputs on a calibration set:

$\hat{W}$1

where $\hat{W}$2 is a matrix of calibration inputs.

2. Algorithmic Innovations

AQLM advances AQ via two central mechanisms:

• Input-adaptive quantization: Code assignments $\hat{W}$3 are data-aware, chosen to minimize output distortion for a specific set of calibration inputs $\hat{W}$4 rather than purely weight-level reconstruction.
• Joint block-wise codebook optimization: Quantization errors from multiple linear layers in a transformer block are addressed collectively by fine-tuning codebooks, scaling parameters $\hat{W}$5, and remaining small parameters $\hat{W}$6 to minimize output mismatch at the block level.

The loss for block-level optimization is:

$\hat{W}$7

and for the full model,

$\hat{W}$8

where $\hat{W}$9 typically is small or zero.

Optimization proceeds via an EM-style process:

• E-step: Updates code assignments $B$0 through beam-search in a Markov Random Field formulation (MRF), leveraging precomputed gram matrices.
• M-step: Refines codebooks $B$1, scales $B$2, and small non-quantized parameters $B$3 using Adam optimizer.

Pseudo-code succinctly describes calibrating a block:

Step	Description
Codebook Initialization	Residual K-means on weight matrix rows
Gram Matrix	Precompute $B$4
E-step	Beam-search code assignment per output unit/group
M-step	Adam updates on $B$5, $B$6, $B$7 to minimize block loss

3. Theoretical Analysis

• Reconstruction Bounds: For $B$8 codebooks of size $B$9, and assignments minimizing MSE,

$B \approx 2\dots3$0

with $B \approx 2\dots3$1 the closest of $B \approx 2\dots3$2 prototypes and $B \approx 2\dots3$3 as $B \approx 2\dots3$4 grow (cf. Babenko & Lempitsky 2014). Empirically, sub-percent layer-level error arises with $B \approx 2\dots3$5, $B \approx 2\dots3$6.

• Pareto Frontier: For LLaMA 2-7B on WikiText2,
  • FP16 compression: 5.12 PPL @ 8 bytes/param
  • QuIP# (2 bit): 8.22 PPL @ 2.02 bits/param
  • AQLM (2 bit): 6.64 PPL @ 2.02 bits/param

AQLM is strictly Pareto-optimal versus all prior 2-bit methods for perplexity versus model size, outperforming certain higher-bit (e.g., 4-bit GPTQ) baselines on smaller models.

4. Implementation and Empirical Performance

AQLM supports high-throughput inference on both GPU and CPU via codebook lookup tables:

• GPU kernel: Precomputes $B \approx 2\dots3$7 lookups for each group, with $B \approx 2\dots3$8 additions per group; achieves $B \approx 2\dots3$9 FP16 speed on RTX 3090 (LLaMA 2-70B).
• CPU kernel: Splits each 16-bit codebook into $8\times$0-bit sub-codebooks so lookups reside in L1/L2 cache; up to $8\times$1 FP32 speedup on a 16-core Intel i9.

Model footprint is reduced by $8\times$2 at 2 bits/parameter relative to FP16, while maintaining or exceeding speed.

Summary of token generation rates:

Device	FP16	AQLM (2 bit)	AQLM (2×8 bit)
RTX 3090, LLaMA-2 7B	41.5 tok/s	32.2 tok/s	32.6 tok/s
Intel i9, LLaMA-2 7B	3.1 tok/s	7.0 tok/s	6.8 tok/s

5. Compression–Accuracy Trade-Offs and Calibration

Several operational variables modulate AQLM’s effectiveness:

• Calibration set size: Gains saturate around 2,000 calibration sequences (useful range: 512–4,096).
• Codebook number $8\times$3, bits $8\times$4, group size $8\times$5: Increasing $8\times$6 improves accuracy with fixed $8\times$7 budget but incurs higher E-step computational cost.
• Block-wise fine-tuning: 100–300 Adam steps increase calibration time by $8\times$8–$8\times$9, securing $M$0–$M$1 PPL reduction.

AQLM thus generalizes well with modest one-shot calibration effort, especially compared to direct PTQ methods.

6. Limitations and Future Prospects

AQLM is the first post-training quantization scheme to reach Pareto optimality below 3 bits/parameter on open LLMs, yielding state-of-the-art perplexity and zero-shot accuracy in the extreme quantization regime. Noted limitations:

• Calibration cost (beam-search E-step) exceeds direct PTQ approaches (e.g., RTN/GPTQ), but remains practical for one-shot application.
• Homogeneous codebooks: Current versions employ fixed codebook architectures; incorporating sparsity or non-uniform (layer-dependent) bit allocation could permit further gains.
• Activation quantization: Extending AQLM to quantized activation flows (quantization-aware inference) could push bit-efficiency below 2 bits.

These results indicate the scalable adaptation of multi-codebook quantization for extreme LLM compression, facilitating efficient on-device inference at high fidelity (Egiazarian et al., 2024).