💬 NLP

Initialisation Determines the Basin: Efficient Codebook Optimisation for Extreme LLM Quantization

This paper identifies codebook initialization as the critical bottleneck in extreme 2-bit LLM quantization and proposes OA-EM, an output-aware EM initialization method using Hessian-weighted Mahalanobis distance, which consistently outperforms existing approaches by placing models in superior optimization basins for subsequent fine-tuning.

Original authors: Ian W. Kennedy, Nafise Sadat Moosavi

Published 2026-04-10

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Ian W. Kennedy, Nafise Sadat Moosavi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Squeezing a Giant into a Backpack

Imagine you have a massive, incredibly detailed encyclopedia (a Large Language Model like Llama 3) that you want to carry around in your pocket on a small phone. The problem is, the book is too heavy and thick. To fit it, you have to compress it.

In the world of AI, "compression" usually means reducing the precision of the numbers inside the model.

4-bit compression is like summarizing the encyclopedia into a thick paperback. It's still very accurate.
2-bit compression is like trying to fit that same encyclopedia onto a single index card. You only have 4 possible symbols to represent every single word. This is extremely hard, and usually, the result is a garbled mess.

This paper tackles the specific problem of 2-bit compression. The authors found that the reason this usually fails isn't because the math is too hard to solve later; it's because the starting point is wrong.

The Problem: The "Bad Map" Analogy

Think of training a compressed AI model like trying to find the lowest point in a vast, foggy mountain range (the "loss landscape"). Your goal is to find the deepest valley (the best, most accurate model).

The Old Way (Greedy Initialization):
Imagine you are dropped into this mountain range blindfolded. You are told to just walk downhill as fast as you can.
- The Trap: If you are dropped on the side of a small, shallow hill, you will walk down to the bottom of that small hill and stop. You think you've found the bottom of the world, but you're actually stuck in a tiny valley while the real deep ocean trench is miles away.
- In the paper, this is called Greedy Sequential Initialization. The AI picks the "best" first piece of the puzzle without looking ahead. Once it makes that first bad choice, it gets stuck in a "bad basin" (a shallow valley).
The Frustrating Reality:
The researchers tried to fix this by sending out more "search parties" (Beam Search) to look for a better path.
- The Result: It's like sending 16 people to explore the mountain instead of 1. But if they all started at the top of the wrong hill, sending more people just helps them explore that specific hill more thoroughly. They still never find the deep ocean trench.
- At 2-bit precision, the "hill" is so steep and the "valley" so shallow that even massive search efforts couldn't save the model. The model became useless (perplexity scores skyrocketed).

The Solution: OA-EM (The "Smart GPS")

The authors realized the problem wasn't the search; it was the drop-off point. They proposed a new method called OA-EM.

The Analogy: Instead of dropping the AI blindfolded on a random hill, OA-EM acts like a Smart GPS.
How it works: Before the AI starts walking, the GPS looks at the terrain and the specific destination (the "output" the AI needs to produce). It calculates the best possible starting spot that leads toward the deepest valley, not just the nearest one.
The Secret Sauce: It uses something called "Hessian-weighted distance." In plain English, this means it pays extra attention to the parts of the model that matter most for getting the right answers, rather than just the parts that look big or loud.

The "Representational Ratio" (The Crowded Room)

The paper introduces a concept called the Representational Ratio ( $\rho$ ). Let's explain this with a party analogy.

The Scenario: You have a room (the Codebook) with a limited number of chairs (representations). You have a bunch of guests (Weight Groups) who need to sit down.
3-bit Compression (Overcomplete): You have 100 guests and 1,000 chairs. Everyone can sit comfortably. Even if you put a guest in a slightly wrong chair, there are plenty of empty chairs nearby to swap into later. The system is forgiving.
2-bit Compression (Undercomplete): You have 100 guests but only 5 chairs.
- The Crisis: Now, every single chair matters. If you put the wrong guest in the wrong chair at the start, there is no room to fix it later. The guests are fighting for space.
- The Finding: The paper shows that at 2-bit, the "chairs" are so scarce that a bad start ruins the whole party. The "Greedy" method puts the wrong people in the chairs, and the "Search" method can't fix it because there's nowhere to move.

The Results: Better Start, Better Finish

The researchers tested this on three different AI models (Llama 3.2, Llama 3.1, and Qwen 2.5).

Before Fine-Tuning:
- Old Method: The model was a disaster (Perplexity of 352, meaning it was guessing randomly).
- New Method (OA-EM): The model was already decent (Perplexity of 17) just by starting in the right place.
After Fine-Tuning (The "PV-Tuning"):
- They applied a standard "fine-tuning" process to both models to see if they could recover.
- The Shock: The old model improved, but it was still stuck in its shallow valley. The new model (OA-EM) started in the deep valley and stayed there.
- The Winner: The OA-EM model was consistently better, even when the old model was given 3x more computing time to search for a solution.

Why This Matters for You

Speed and Cost: The new method (OA-EM) with a small search budget produced a better model than the old method with a huge search budget. This means we can make powerful AI models run on cheap, small devices (like phones or Raspberry Pis) much faster and cheaper.
The Lesson: In AI, starting well is more important than searching hard. If you start in the wrong place, no amount of effort will get you to the best destination.

Summary in One Sentence

This paper discovered that when squeezing AI models down to tiny sizes, the biggest mistake is starting in the wrong spot; by using a smarter "GPS" to choose the starting point, they can create high-quality models that run on small devices without needing expensive, time-consuming search methods.

1. Problem Statement

The paper addresses the challenge of extreme quantization (2-bit precision) for Large Language Models (LLMs) using Additive Quantization (AQLM). While AQLM allows for $O(1)$ lookup-table (LUT) dequantization—making it ideal for edge devices where memory bandwidth is the bottleneck rather than compute—it often fails catastrophically at 2-bit precision, even with extensive search and fine-tuning.

The Core Issue: The dominant bottleneck is not the search algorithm (e.g., beam search) or the fine-tuning process (PV-tuning), but the initialization of the codebooks.
The Failure Mode: Standard greedy sequential initialization (fitting codebooks one by one) frequently places the model in a poor optimization region (a "bad basin"). Once trapped, subsequent beam search and fine-tuning cannot recover the model's performance, leading to orders-of-magnitude degradation in perplexity.
Theoretical Insight: The authors introduce the Representational Ratio ( $\rho = N/KM$ ), where $N$ $N$ is the number of weight groups, $K$ $K$ is the codebook size, and $M$ $M$ is the number of additive codebooks.
- Overcomplete ( $\rho < 1$ ): There are more representable points than weight groups; initialization errors can often be absorbed.
- Undercomplete ( $\rho > 1$ ): Weight groups compete for limited capacity. At 2-bit precision (e.g., Llama 3.2 3B), $\rho \approx 18$ , making the system highly sensitive to initial placement. A 256x reduction in capacity compared to 3-bit leads to a qualitative shift in behavior where greedy initialization fails.

2. Methodology: OA-EM

To address the initialization bottleneck, the authors propose OA-EM (Output-Aware Expectation-Maximisation).

Concept: Unlike standard $k$ -means or greedy residual fitting which minimize Euclidean distance in weight space, OA-EM minimizes output reconstruction error directly.
Mechanism:
- It operates within the AQLM residual framework but refines codebook initialization using an iterative EM process.
- Distance Metric: It replaces Euclidean distance with a Hessian-weighted Mahalanobis distance. The Hessian ( $H_i$ ) is approximated from calibration activations ( $X_i^T X_i + \lambda I$ ), capturing the sensitivity of the model's output to weight perturbations.
- M-Step (Centroid Optimization): Updates centroids to minimize the Hessian-weighted reconstruction error: $L_{EM} = \frac{1}{N} \sum e_i^T H_i e_i$ . This ensures centroids are placed where they matter most for the model's output, rather than just where weight magnitudes are large.
- E-Step (Reassignment): Reassigns weight groups to the nearest centroid based on the Mahalanobis distance.
Integration: OA-EM serves as a drop-in replacement for the initial greedy step before beam search and PV-tuning.

3. Key Contributions

Identification of the Initialization Bottleneck: The paper demonstrates that in the undercomplete regime ( $\rho \gg 1$ ), poor initialization traps the model in a suboptimal basin that subsequent search and fine-tuning cannot escape.
Representational Ratio ( $\rho$ ): A theoretical metric that predicts when quantization becomes sensitive to initialization. It explains why 2-bit quantization fails while 3-bit often succeeds.
OA-EM Algorithm: A novel initialization method using Hessian-weighted EM that aligns codebook placement with output sensitivity, significantly reducing reconstruction error.
Basin Persistence: Empirical evidence showing that the benefits of OA-EM persist even after end-to-end PV-tuning, indicating that initialization determines the final optimization basin.

4. Experimental Results

The authors evaluated OA-EM on Llama 3.2 3B, Llama 3.1 8B, and Qwen 2.5 3B at 2-bit and 3-bit precision.

Pre-PV-Tuning Performance (Llama 3.2 3B @ 2-bit):
- Greedy Initialization: WikiText-2 perplexity was catastrophic (e.g., 352.39 with beam width 4). Even increasing beam width to 16 only reduced it to 46.01.
- OA-EM: Achieved a perplexity of 16.82 with beam width 4, comparable to the best greedy results with beam width 16, but with significantly less compute.
Post-PV-Tuning (Basin Persistence):
- PV-tuning improved all models, but the gap remained.
- Greedy (b=8): Post-PV perplexity: 11.76.
- OA-EM (b=8): Post-PV perplexity: 11.53.
- Crucially, OA-EM with a narrow beam (b=4) achieved 11.53 perplexity, outperforming the greedy baseline with a wide beam (b=8) at 11.76, despite using 38% less quantization time.
Downstream Tasks: OA-EM consistently matched or improved zero-shot accuracy on tasks like ARC, HellaSwag, and PIQA, particularly on Llama 3.2 3B where the initialization bottleneck was most severe.
Domain Robustness: The performance gap between greedy and OA-EM scaled with domain distance. Greedy initialization memorized calibration statistics (C4), failing on out-of-domain data (WikiText-2), while OA-EM maintained robustness.

5. Significance and Implications

Optimization Geometry over Search Intensity: The paper challenges the prevailing strategy of "throwing more compute at the problem" (wider beams, more epochs). It argues that improving initialization quality is more effective than increasing search intensity in extreme compression regimes.
Edge Deployment: By enabling high-quality 2-bit quantization with narrow beam searches, OA-EM makes extreme compression viable for consumer GPUs, mobile devices, and microcontrollers where LUT-based dequantization is critical.
General Principle: The concept of the Representational Ratio and Basin Persistence provides a new lens for understanding vector quantization. It suggests that for any learned-codebook method relying on greedy initialization, the initial placement dictates the final solution quality, regardless of subsequent fine-tuning.
Efficiency: The method offers a "free" speedup; the cheapest OA-EM configuration produces better models than the most expensive greedy configuration.

In summary, the paper establishes that for extreme LLM quantization, initialization is the primary determinant of success. By shifting from weight-space distance to output-aware, Hessian-weighted initialization, OA-EM unlocks the potential of free-form additive quantization at 2-bit precision, preserving the efficiency of LUT-based inference while achieving state-of-the-art perplexity.