AQLM Extreme Compression via Additive Quantization

Chapter 1: Vector Quantization Primer

Before we can understand AQLM's multi-codebook trick, we need to understand the simpler version: vector quantization (VQ). The idea is ancient — it dates back to the 1980s in signal compression — but it is the foundation of everything that follows.

Instead of quantizing each weight independently (scalar quantization), VQ groups weights into vectors and quantizes the entire vector at once. Think of it this way: you have a codebook — a dictionary of K prototype vectors. To compress a weight vector, you find the closest codebook entry and store only its index.

From K-Means to Codebooks

Training a VQ codebook is just K-means clustering. You gather all weight vectors from the model, run K-means with K clusters, and the cluster centroids become your codebook entries. Each weight vector is then assigned to its nearest centroid.

Storage: if your codebook has K = 2^B entries (B-bit codes), each vector of g weights costs B bits to store (just the index). So the bits per parameter is B / g. With B = 8 and g = 8, that is 1 bit per parameter — extreme compression.

The Storage Arithmetic

Let us make the VQ savings concrete. A weight matrix in Llama 2 7B's MLP is 4096 × 11008. In FP16, that is 4096 × 11008 × 2 = 90.2 MB. With VQ using g = 8, B = 8 (256-entry codebook):

• Indices: (4096 × 11008 / 8) groups × 1 byte per index = 5.6 MB
• Codebook: 256 entries × 8 floats × 2 bytes = 4 KB (negligible)
• Total: ~5.6 MB — a 16× compression from 90.2 MB

That is 1 bit per parameter. But the accuracy is terrible — 256 prototype vectors cannot capture the diversity of millions of weight groups. We need more expressiveness without more bits.

Product Quantization (PQ)

Pure VQ has a problem: if you want fine-grained quantization, you need a huge codebook. With g = 8 dimensions and B = 16 bits, the codebook has 65,536 entries of 8 floats each — that is 2 MB per codebook. And you need one per layer.

Product Quantization (PQ) fixes this by splitting each vector into sub-vectors and using a separate, smaller codebook for each sub-vector. A vector of 8 weights is split into 2 sub-vectors of 4, each encoded with its own 8-bit codebook (256 entries). Total: 16 bits for 8 weights = 2 bits per parameter, but the codebooks are tiny.

Chapter 2: The Additive Quantization Framework

Now the main idea. Take a weight matrix W ∈ R^{d_out × d_in} from a linear layer. Split its rows into groups of g consecutive weights. Each group is a vector w ∈ R^g. We want to approximate w using M codebooks C₁, ..., C_M, each containing 2^B vectors in R^g.

The approximation is a sum:

where b_m ∈ {0, 1, ..., 2^B − 1} is the index into the m-th codebook. Each codebook contributes one vector, and they are added together to reconstruct the weight group. This is Additive Quantization (AQ).

The Optimization Objective

AQLM is not just about finding the nearest codebook entry. It is input-aware. Given a calibration dataset with input activations X, the objective is:

This minimizes the squared error on the outputs of the layer, not on the weights directly. Why? Because not all weights are equally important. A weight that multiplies a large activation matters far more than one that multiplies near-zero. The input matrix X (collected from calibration data) encodes this importance.

Expanding the Objective

Substituting the additive representation into the objective and expanding:

where ⟨A, B⟩_XX^T = ⟨AXX^T, B⟩_F is the Frobenius inner product weighted by the input correlation matrix. The key insight: XX^T can be pre-computed once from the calibration data. This makes the objective efficient to evaluate — no need to pass activations through the model repeatedly.

Why Not Just a Bigger Codebook?

A natural question: why use M = 2 codebooks of 256 entries each (total 512 vectors to store) when you could use 1 codebook of 65,536 entries? Both use the same number of bits per weight group (16 bits). The answer is threefold:

• Codebook memory. A single 2¹⁶-entry codebook with g = 8 dimensions occupies 65,536 × 8 × 2 = 1 MB in FP16. Two 2⁸-entry codebooks occupy only 2 × 256 × 8 × 2 = 8 KB. That is 128× smaller, and small codebooks fit in GPU shared memory for fast lookups.
• Optimization landscape. Jointly optimizing 65K codebook entries with K-means is a hard problem — many local minima. Optimizing two sets of 256 entries is much easier and converges faster.
• Generalization. Smaller codebooks are shared across more weight groups. This parameter sharing acts as regularization — the codebooks must represent common patterns, not memorize individual weights.

In practice, AQLM finds that the 1×16 configuration (single large codebook) gives slightly better accuracy, while 2×8 (two small codebooks) gives faster inference. The paper reports both.

Representation Cost

Each weight group stores M indices of B bits each. Additionally, AQLM learns a per-output-unit scale factor s ∈ R (stored in FP16). The reconstructed weight is:

The scale factors add a negligible cost (one 16-bit float per row of the weight matrix) but help compensate for the limited codebook expressiveness.

Concrete Data Flow

Let us trace the full storage layout for one linear layer of Llama 2 7B. The q_proj weight matrix is 4096 × 4096.

The codebooks are shared across all 512 weight groups per output dimension. The indices dominate storage; the codebooks and scales are negligible. Compression ratio: 8.4× (equivalent to ~1.9 bits per parameter for this layer, before accounting for the codebook overhead).

Chapter 3: Beam Search for Code Assignment

AQLM has three phases: (1) find the best codebook indices b for each weight group, (2) optimize the codebook entries C, (3) fine-tune at the block level. Phase 1 is the hardest — it is a discrete combinatorial problem.

For a single weight group, we need to choose M indices (one per codebook), each from 2^B options. The brute-force search over all (2^B)^M combinations is intractable. With B = 8 and M = 2, that is 65,536 combinations per group — feasible. But with M = 4, it is 4.3 billion. We need a smarter search.

The Beam Search Algorithm

AQLM uses beam search, adapted from the additive quantization literature (Babenko & Lempitsky, 2014). The algorithm maintains a beam of k best partial assignments and extends them one codebook at a time:

Worked Example

Suppose g = 2 (2D weight groups for simplicity), M = 2 codebooks, each with 4 entries (B = 2 bits). We want to approximate w = [0.7, −0.3].

Codebook C₁ has entries: {[0.5, 0.0], [0.0, −0.5], [−0.3, 0.4], [0.8, 0.2]}. Codebook C₂ has entries: {[0.1, −0.2], [0.3, 0.1], [−0.1, −0.4], [0.2, −0.1]}.

Step 1: Score all 4 entries from C₁ against w (using the MSE objective). Suppose the scores rank: entry 3 ([0.8, 0.2]) best, entry 0 ([0.5, 0.0]) second. With beam width k = 2, we keep both.

Step 2: For each beam, try all 4 entries from C₂. That gives 2 × 4 = 8 candidates:

The winner: C₁[3] + C₂[2] = [0.8, 0.2] + [−0.1, −0.4] = [0.7, −0.2]. Error = 0.01. The target w = [0.7, −0.3] is nearly recovered by summing two codebook vectors.

Why Beam Search, Not Greedy?

A greedy approach picks the single best entry from C₁, then the single best from C₂ given that choice, and so on. This is beam search with k = 1. The problem: early choices constrain later ones. If the best C₁ entry points the approximation slightly wrong, C₂ may not be able to fix it. Beam width k lets us maintain k alternative "hypotheses" and recover from suboptimal early decisions.

In practice, AQLM uses a beam width of k = 1 for the initial assignment (since codebooks are subsequently optimized anyway). The beam search runs over all d_out output units in parallel — each output dimension is independent in the loss function, so this is a massively parallel operation on GPU.

Complexity

For each weight group: beam search considers k · 2^B candidates per codebook step, and there are M steps. Total candidates evaluated: M · k · 2^B. Compare to brute force: (2^B)^M = 2^BM. With B = 8, M = 2, k = 1: beam search evaluates 512 candidates. Brute force would check 65,536. With M = 4: beam search evaluates 1024; brute force would check 4.3 × 10⁹.

Chapter 4: Codebook Tuning

Phase 1 (beam search) found the best indices assuming fixed codebooks. Phase 2 flips the problem: fix the indices, optimize the codebook entries.

With the codes b fixed, the objective becomes a continuous optimization over the codebook vectors C₁, ..., C_M and the scale factors s:

where Ŵ = diag(s) · [∑_m C_mb_m] is the reconstructed weight matrix. This is differentiable with respect to C_m and s, so we can use standard gradient-based optimization.

Computing the Gradient

The loss is quadratic in the codebook entries. Let us trace the gradient for a single codebook entry c = C_m[k] that is used by some set of weight groups S_k (all groups where b_m = k). The gradient is:

This is a weighted average of the residual errors across all weight groups that share this codebook entry, weighted by the input correlation matrix XX^T. The gradient is cheap to compute because XX^T is pre-computed and shared, and each weight group contributes one additive term.

Adam Optimizer

AQLM uses Adam with learning rate 10⁻⁴ and full-batch gradient descent on the calibration data (128 samples). Each codebook update phase runs for 100 steps. The authors found the results are robust to these hyperparameters — varying the learning rate or number of steps does not significantly change the outcome.

Scale Factor Initialization

The per-output-unit scale s_i is initialized to the L2 norm of the i-th row of W: s_i = ||W_i||₂. This gives each row a learned magnitude while the codebooks handle the direction. The scales are updated alongside the codebooks via the same Adam optimizer.

Alternating Optimization

Phases 1 and 2 alternate: update codes (beam search) → update codebooks (Adam) → update codes → update codebooks → ... until the loss stops improving (or for a fixed number of rounds). Each round refines both the discrete assignments and the continuous codebook entries.

Initialization: Residual K-Means

Good initialization is critical. AQLM uses residual K-means (Chen et al., 2010):

1. Run K-means on all weight groups to get C₁. Assign each group to its nearest centroid.
2. Compute the residual: r = w − C₁[b₁] (the part of each weight group not captured by the first codebook).
3. Run K-means on the residuals to get C₂. This codebook specializes in what C₁ missed.
4. Compute the new residual: r' = w − C₁[b₁] − C₂[b₂]. Repeat for C₃, etc.

This is equivalent to Residual Quantization (RQ) and gives a reasonable starting point. The subsequent alternating optimization (beam search + Adam) then jointly refines all codebooks, escaping the limitations of the sequential residual approach.

Chapter 5: Block-Level Fine-Tuning

Phases 1 and 2 compress each layer independently. But in a real transformer, quantization errors in one layer's output become the input to the next layer. Errors compound. At 2 bits, this compounding is devastating.

Phase 3 of AQLM fixes this with block-level fine-tuning. Instead of optimizing one layer at a time, it optimizes an entire transformer block (self-attention + MLP, typically 4–8 linear layers) jointly.

The Procedure

After quantizing all layers within a block using phases 1 and 2, AQLM:

1. Records the original block output Y_block = block(X_block) using FP16 weights
2. Replaces the block's weights with their quantized versions
3. Defines the fine-tuning loss: || Y_block − block̂(X_block) ||²
4. Optimizes the codebook entries C_m, scale factors s, and all non-quantized parameters (RMSNorm scales and biases) via Adam, backpropagating through the entire quantized block
5. Keeps the codes b_m fixed — only continuous parameters are tuned

Practical Details

Fine-tuning uses the PyTorch autograd engine. The calibration data is the same 128 samples from RedPajama (sequence length 4096) used in phases 1 and 2. The optimizer is Adam with a schedule. Fine-tuning the transformer blocks takes only 10–30% of the total calibration time — it converges quickly because it starts from a good initialization (the per-layer quantization from phases 1–2).

The algorithm processes the model block by block, from the first transformer layer to the last. After fine-tuning block k, it saves the updated codebooks and moves to block k+1, using the quantized output of block k as input. This sequential processing means each block adapts to the actual (quantized) inputs it will see at inference time.

Algorithm 1: Full AQLM Pipeline

pseudocode
Input: model, calibration data (128 samples)
for layer = 1 to num_layers:
  W = layer.weight                # original FP16 weights
  X = layer_inputs(calibration)    # collect activations
  C, s = initialize(W)             # K-means + residual K-means

  while loss improves:
    C, s = adam_update(XX^T, W, C, b, s) # Phase 2: codebook tune
    b = beam_search(XX^T, W, C, b, s) # Phase 1: code update

  layer.weight = AQLMFormat(C, b, s)

for block = 1 to num_blocks:         # Phase 3: block fine-tune
  theta = trainable_params(block)   # codebooks + scales + norms
  while loss improves:
    L = ||block(X) - Y_orig||^2
    theta = adam(theta, grad(L))

Chapter 6: The Bits Arithmetic

Let us work through the exact compression math. This is where AQLM's configurations map to concrete bits-per-parameter numbers.

The Formula

For a weight group of g parameters encoded with M codebooks of 2^B entries each, the cost per group is:

Plus there are two small overheads: (1) the codebook entries themselves, and (2) the per-output scale factors. For large models, these are negligible compared to the indices.

AQLM Configurations

What Does Not Count

Following standard practice, the "bits per parameter" metric excludes parameters that are kept in full precision: the embedding layer and the LM head. These are a small fraction of total parameters for large models (e.g., 1.3% for Llama 2 70B). All benchmarks compare methods at the same average bitwidth, so this is a fair comparison.

Chapter 7: Codebook Lookup Visualization

Let us see the entire AQLM dequantization process in action. This is what happens during inference when the model needs to reconstruct a weight group to compute a matrix-vector product.

The Dequantization Pipeline

GPU Kernel Architecture

The kernel operates in two phases per thread block:

1. Load codebooks into shared memory. For the 2×8-bit configuration (M = 2 codebooks, 256 entries each, g = 8 floats per entry), each codebook is 256 × 8 × 2 bytes = 4 KB in FP16. Two codebooks fit in 8 KB — well within the 48–164 KB of shared memory available on modern GPUs.
2. Stream codes + accumulate. Each thread loads a code pair (b₁, b₂), looks up both codebook entries in shared memory (no global memory access), sums them, multiplies by the scale factor, computes the dot product with the input activation group, and atomically accumulates into the output.

The key insight: codebook lookups from shared memory cost ~5 cycles, while global memory loads cost 200–400 cycles. Since the codebooks are reused across all weight groups in a column, the shared memory approach amortizes the load cost across thousands of groups.

Inner Product Trick

For the 1×16 configuration with a single large codebook, there is an additional speedup. The dot product between a weight group and the input can be rewritten as:

You can pre-compute the dot products of all codebook entries with the current input: dp[k] = C[k] · x for k = 0, ..., 2¹⁶ − 1. Then reconstruction + dot product becomes a single table lookup: dp[b]. This turns the computation from g multiplications and additions per group into a single indexed load. However, with 65K entries, the lookup table is large (256 KB per group of 8 input dimensions), so this trick is mainly useful for CPU inference where caches are larger.

Inference Speed

The GPU speedups are modest for 7B (the model already fits in FP16) but dramatic for 70B, where the FP16 model requires multi-GPU while the 2-bit model fits on a single card. On CPU, the inner product pre-computation trick delivers consistent 3–4× speedups.

Chapter 8: Results

The moment of truth. How does AQLM compare to GPTQ, AWQ, SpQR, and QuIP# on actual language model benchmarks?

2-Bit Regime (Llama 2)

This is where AQLM shines brightest. All methods are compressing to approximately 2 bits per parameter. Perplexity is measured on WikiText-2 (lower is better).

3-Bit Regime (Llama 2)

Calibration Data

All results above use a tiny calibration set: 128 sequences of length 4096 tokens from the RedPajama dataset. That is roughly 524K tokens — a vanishingly small fraction of the pre-training data. The authors found that:

• Increasing calibration samples beyond 128 gives diminishing returns (less than 0.1 PPL improvement on 7B)
• Calibration from the same domain as the eval set helps, but the model is not sensitive to the exact dataset
• AQLM benefits more from larger calibration sets than GPTQ does, because the codebook optimization has more parameters to tune and thus more capacity to absorb information from data

End-to-End Fine-Tuning Results

After the paper's initial release, the authors added end-to-end fine-tuning (following QuIP#'s approach): fine-tune the entire model to minimize KL divergence, not just per-block MSE. This further improves results:

With end-to-end fine-tuning (marked *), AQLM and QuIP# reach near-parity at 2-bit. Both erase the "2-bit is useless" narrative. The 2-bit 70B model achieves 3.83 perplexity — closer to FP16's 3.12 than GPTQ's 3-bit result of 4.40.

Chapter 9: Connections

AQLM sits at the intersection of information retrieval (additive quantization) and LLM compression (post-training quantization). Here is how it connects to the broader landscape.

Method Comparison

Key Takeaways

• 2-bit LLMs are viable. AQLM proved that with the right quantization structure, you can compress to 2 bits and retain most of the model's capability. The gap between 2-bit AQLM and FP16 is smaller than the gap between GPTQ at 3 bits and FP16.
• Vector quantization beats scalar at extreme compression. The codebook structure captures correlations between weights that per-element rounding cannot.
• Input-aware objectives matter. Minimizing output MSE (weighted by activations) is fundamentally better than minimizing weight MSE, especially at low bitwidths.
• Block-level fine-tuning is cheap and powerful. A few hundred Adam steps on the codebooks, using only 128 calibration samples, can recover a significant fraction of the quantization loss.
• Calibration data is small. 128 sequences of length 4096 from RedPajama. No need for large fine-tuning datasets.

Limitations

• Calibration time. AQLM is slower to quantize than GPTQ (hours vs. minutes for a 7B model) because of the alternating beam search + codebook optimization.
• Kernel complexity. The multi-codebook lookup kernels are more complex than simple integer dequantization, requiring careful GPU kernel engineering.
• Not RTN-compatible. Round-to-nearest (RTN) methods can be applied without any calibration data. AQLM requires a calibration set, which is a stronger assumption.

Using AQLM in Practice

python
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load a pre-quantized AQLM model from HuggingFace
model_id = "ISTA-DASLab/Llama-2-7b-AQLM-2Bit-1x16-hf"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype="auto",
    device_map="auto"    # fits on a single 24GB GPU
)

# Generate text -- the dequantization happens inside the kernel
inputs = tokenizer("The key insight of additive quantization is", return_tensors="pt")
outputs = model.generate(**inputs, max_new_tokens=100)
print(tokenizer.decode(outputs[0]))

What Came Next

After AQLM, the authors and others pushed further:

• PV-Tuning (Tseng et al., 2024) improved the fine-tuning phase, showing that better optimization of the codebooks can recover additional accuracy.
• QuIP# v2 added end-to-end fine-tuning matching AQLM, demonstrating that the fine-tuning innovation (not just the quantization structure) was key to the accuracy gains.
• QLoRA + AQLM: The community found that AQLM-quantized models can be further fine-tuned with LoRA adapters, enabling task-specific adaptation on top of extreme compression.
• HuggingFace integration: AQLM was integrated into the Transformers library, making it accessible via AutoModelForCausalLM.from_pretrained() with automatic kernel dispatch.

The open-source AQLM codebase at github.com/Vahe1994/AQLM provides quantization scripts and pre-quantized model weights for Llama 2, Mistral, and other model families.

"The question of optimal data compression is really a question of understanding." — Gregory Chaitin