PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Ian W. Kennedy, Nafise Sadat Moosavi

arXiv:2604.08118v1cs.CLcs.LG

TL;DR

Extreme additive quantization is attractive for memory-constrained deployment but can fail at 2-bit precision despite extensive search and fine-tuning. The paper identifies codebook initialisation as the dominant bottleneck, characterises it with the representational ratio ρ, and proposes OA-EM, whose advantages persist after PV-tuning across tested settings. The study is scoped to free-form additive quantization and primarily the 3B–8B model range.

  • Problem

    At 2-bit precision, additive quantization can fail severely, while the role of codebook initialisation across compression rates and subsequent fine-tuning has not been examined.

  • Method

    OA-EM uses output-aware expectation-maximisation with Hessian-weighted Mahalanobis distance from calibration activations to initialise additive-quantization codebooks.

  • Results

    Across compression rates, search budgets, and three model architectures, OA-EM consistently produces lower perplexity than greedy initialisation after PV-tuning and dominates the quality–compute frontier.

  • Takeaways & Limitations

    Initialisation strongly influences the optimisation basin of compressed models, with the effect extreme in the undercomplete regime where ρ is much greater than 1.

  • Takeaways & Limitations

    The method applies to free-form additive quantization, does not directly transfer to lattice- or trellis-based methods, and the experiments focus on the 3B–8B range.

Abstract

from arXiv · show

Additive quantization enables extreme LLM compression with O(1) lookup-table dequantization, making it attractive for edge deployment. Yet at 2-bit precision, it often fails catastrophically, even with extensive search and finetuning. We show that the dominant bottleneck is codebook initialisation. Greedy sequential initialisation frequently places the model in poor optimisation regions that subsequent beam search and PV-tuning struggle to overcome. We analyse this behaviour through the representational ratio \r{ho} = N/KM, which characterises the relationship between weight groups and codebook capacity, and propose OA-EM, an output-aware EM initialisation method using Hessian-weighted Mahalanobis distance. Across compression rates, search budgets, and three architectures (Llama 3.2 3B, Llama 3.1 8B, Qwen 2.5 3B), OA-EM consistently produces better solutions after PV-tuning and dominates the quality-compute frontier. The severity of the bottleneck scales with \r{ho}: moderate at 3 bpp but extreme at 2 bpp, where poor initialisation can degrade perplexity by orders of magnitude. More broadly, our results highlight the importance of optimisation geometry in compressed model spaces, where initialisation can dominate subsequent search and fine-tuning.

1 Introduction

At extreme additive-quantization compression, codebook initialisation can determine the optimisation basin more strongly than subsequent search. OA-EM uses output-aware EM to improve initialisation, with benefits that persist across search settings, compression rates, and model architectures.

  • Mechanism: The representational ratio ρ = N/KM relates weight groups to codebook capacity and predicts when initialisation becomes critical.When ρ > 1, weight groups compete for limited codebook capacity, making initial placement especially important.
  • Motivation: 2 bpp greedy initialisation yields WikiText-2 perplexity 352.39 at beam 4 and 60.61 at beam 8, versus 7.28 for FP16.At beam 16, perplexity remains 46.01, showing that wider search does not fully recover the poor starting point.
  • Method: OA-EM refines codebook initialisation with iterative EM and Hessian-weighted Mahalanobis distance from calibration activations, targeting output reconstruction error.It is proposed as a replacement for greedy initialisation.
  • Results: OA-EM remains better after PV-tuning in every tested Llama 3.2 3B configuration, even as a 43-point perplexity gap narrows to 0.23.The comparison spans beam widths and epoch budgets.
  • Results: Increasing beam width from 8 to 16 improves OA-EM post-PV perplexity from 11.53 to 11.49 but worsens the greedy baseline from 11.76 to 12.01.The authors report this as asymmetric search scaling, with additional search helping primarily after good initialisation.
  • Results: OA-EM dominates the quality–compute frontier, achieving lower perplexity and better downstream accuracy at matched compute budgets.The evaluation covers multiple compression rates, search budgets, and model architectures.

2 Related Work

Prior work improves quantizer geometry, codebook learning, or post-quantization tuning, but does not examine how additive-quantization initialisation interacts with compression and later fine-tuning. This paper positions OA-EM as an orthogonal improvement to the initialisation stage.

  • Structured and free-form quantization: Structured methods such as QuIP#, QTIP, and GLVQ constrain codebook geometry, while free-form additive methods learn unconstrained codebooks.Structured approaches can avoid the studied initialisation problem, whereas free-form methods expose it.
  • Additive quantization: Free-form additive quantization enables extreme compression but leaves codebook initialisation as a potential bottleneck.AQLM introduced multi-codebook additive quantization for LLM compression, and PV-tuning later added end-to-end fine-tuning.
  • Research gap: Existing additive-quantization work does not examine how initialisation quality interacts with compression rate or persists through fine-tuning.This gap is especially relevant at extreme compression, where poor initialisation can cause failure despite sufficient representational capacity.
  • EM-based methods: Prior EM-based methods apply EM to vector quantization, whereas OA-EM operates in multi-codebook additive LLM quantization and is output-aware.OA-EM optimises activation-space reconstruction and serves as an initialisation stage for downstream beam search.
  • Positioning: The paper improves initialisation rather than quantizer geometry or post-quantization fine-tuning, showing that it can determine the persistent optimisation basin.The representational-ratio and basin-persistence analyses may also inform learned-codebook methods combining structured and free-form components.

3 Method

The method identifies greedy codebook initialisation as a bottleneck in additive quantization and introduces OA-EM to improve centroid placement using output-aware, Hessian-weighted optimisation. It combines this initialisation with AQLM’s residual representation and downstream beam search or fine-tuning.

  • AQLM background: AQLM represents each weight group as a sum of codewords and minimises output reconstruction error on calibration data.At 2 bpp with group size g=8, AQLM uses M=2 codebooks with K=256 entries each.
  • AQLM background: Residual k-means initialises codebooks sequentially, while beam search refines assignments over combinations of codewords.Beam search maintains b active candidates per codebook stage, costing O(MbK) per weight group.
  • Initialisation bottleneck: Greedy sequential assignment can irreversibly commit to a poor first-codebook entry because later codebooks may not compensate for that choice.The problem worsens when representational capacity is limited, because codebooks are jointly coupled rather than independently optimised.
  • Initialisation bottleneck: At 2 bpp, wider beams increase quantization time from 6.1h to 16.9h while reducing WikiText-2 perplexity from 352.39 to 46.01.The passage argues that beam search cannot reshape poorly placed centroids, even when assignment search becomes more extensive.
  • OA-EM: OA-EM refines sequential codebook initialisation with Hessian-weighted Mahalanobis distances for centroid optimisation and assignment.Its M-step targets Hessian-weighted reconstruction error, using a damped block-diagonal Hessian approximation derived from calibration activations.
  • OA-EM: OA-EM targets the dominant greedy-error terms by allocating centroids toward output-sensitive groups when representational capacity is scarce.This reduces first-codebook displacement where it contributes most to the reconstruction objective while preserving the free-form additive-quantization paradigm.

4 Experimental Setup

The experiments evaluate AQLM quantization across three model architectures, multiple compression and search settings, and perplexity and zero-shot benchmarks. They isolate initialisation effects before and after PV-tuning using controlled hardware and fixed-seed runs.

  • Models: Three models are evaluated: Llama 3.2 3B, Llama 3.1 8B, and Qwen 2.5 3B.All models use AQLM at 2 bpp with M=2 and g=8; Llama 3.2 3B is also tested at 3 bpp with M=3.
  • Calibration: Calibration uses 128 C4 sequences of length 4096.This calibration setup follows common post-training quantization practice.
  • Search configurations: Beam width varies over b ∈ {4, 8, 16} and maximum epochs over e ∈ {5, 100}, spanning a 2.8× quantization-time range.On Llama 3.2 3B, the reported range is 6.1h to approximately 17h, with early stopping at 0.01 relative MSE.
  • OA-EM configuration: OA-EM runs for 3 EM rounds with 100 Adam steps per round and learning rate η = 10^-4 annealed to 0.1η.Both E-step and M-step use the damped block-diagonal Hessian approximation.
  • PV-tuning: PV-tuning uses the same Adam-based configuration across quantization settings to isolate initialisation effects.The procedure uses η_ft = 3 × 10^-4, batch size 32, 10K samples, and 5 epochs, selecting the best WikiText-2 checkpoint.
  • Evaluation: Evaluation reports WikiText-2 and C4 perplexity alongside zero-shot results on ARC-Easy, ARC-Challenge, HellaSwag, PIQA, WinoGrande, and LAMBADA.Zero-shot evaluation uses the LM Evaluation Harness.
  • Evaluation: Table 1 reports pre-PV-tuning results at 2 bpp on Llama 3.2 3B, including beam width, maximum epochs, and quantization time on one A100.The table is intended to compare search configurations before fine-tuning.
  • Hardware: All experiments use a single fixed random seed, with one A100 80 GB GPU except Llama 3.1 8B PV-tuning on one B200 192GB GPU.This is an experimental-scope constraint on hardware and run replication.

5 Initialisation and Basin Persistence

The study tests whether codebook initialisation determines the optimisation basin reached by additive quantization. OA-EM remains advantageous before and after PV-tuning, across search configurations, while broader search benefits good initialisation more reliably.

  • Initialisation Effects Before PV-Tuning: At 3 bpp, OA-EM reduces WikiText-2 perplexity from 9.52 to 8.87 (−0.65), while improving LAMBADA accuracy from 0.673 to 0.687.C4 perplexity increases slightly from 13.39 to 13.51, while LAMBADA perplexity drops from 4.87 to 4.60.
  • Initialisation Effects Before PV-Tuning: At 2 bpp, greedy initialisation remains severely degraded before PV-tuning, whereas OA-EM stays stable across beam widths at 16.82–17.39 WikiText-2 perplexity.Greedy results improve from 352.39 at b=4 to 46.01 at b=16, but remain far worse than OA-EM.
  • Initialisation Effects Before PV-Tuning: With b=8, OA-EM reaches 17.39/18.00 perplexity in 9.2h versus 60.61/18.64 in 9.9h for greedy initialisation.The paired values are WikiText-2/C4 perplexities.
  • Basin Persistence After PV-Tuning: PV-tuning compresses the 43-point perplexity gap to 0.23, yet OA-EM remains better in every tested configuration.This persistence indicates that fine-tuning mitigates but does not erase the effect of initialisation.
  • Basin Persistence After PV-Tuning: After PV-tuning, increasing beam width from 8 to 16 improves OA-EM from 11.53 to 11.49 but worsens greedy initialisation from 11.76 to 12.01.The contrasting responses suggest different optimisation trajectories.
  • Basin Persistence After PV-Tuning: OA-EM improves the quality–time frontier: at b=4, it achieves 11.53 perplexity in 6.1h versus greedy’s 11.76 in 9.9h at b=8.This is 38% less quantization time for better perplexity.

6 Downstream Task Performance

OA-EM improves downstream accuracy across the settings where initialisation is non-trivial, with the clearest gains on Llama 3.2 3B and broader but smaller gains on Llama 3.1 8B. Perplexity remains the more consistent signal across architectures.

  • Cross-model results: OA-EM matches or improves average accuracy in every setting where the initialisation bottleneck is non-trivial.Table 3 summarizes post-PV-tuning downstream performance across models and beam configurations.
  • Llama 3.2 3B: On Llama 3.2 3B, OA-EM wins average accuracy across all four beam configurations, with a +1.7pp gain at b=4.The clearest gain occurs where the greedy baseline has minimal search to compensate for poor initialisation.
  • Llama 3.1 8B: On Llama 3.1 8B, OA-EM wins 4 of 6 accuracy tasks with a 0.7-point average improvement.This result extends the downstream advantage beyond the 3B model.
  • Qwen 2.5 3B: On Qwen 2.5 3B, the baseline has a small downstream advantage of 0.606 versus 0.603 average accuracy, despite OA-EM’s perplexity advantage.The paper notes that zero-shot evaluations at the 3B scale are high-variance and treats perplexity as the more reliable signal.

7 Generality Across Compression Rates and Architectures

The initialisation bottleneck varies with compression rate and architecture. Its severity increases sharply with ρ, while model-specific weight statistics moderate whether high ρ produces catastrophic degradation.

  • Compression-rate gradient: At 3 bpp (ρ ≈0.07), the pre-PV gap of 0.65 compresses to 0.12 after PV-tuning, while OA-EM wins 5/6 downstream tasks.OA-EM also improves ARC-Easy by 3.5 points in this regime.
  • Compression-rate gradient: At 2 bpp (ρ ≈18), the 43-point pre-PV gap compresses to 0.23 after PV-tuning, yet OA-EM wins every beam configuration and most downstream tasks.The larger representational ratio corresponds to a much more severe initialisation bottleneck.
  • Basin persistence: PV-tuning improves models within their existing optimisation basin rather than moving them between basins.This conclusion is supported by persistent pre- and post-PV differences across compression rates and architectures.
  • Architecture dependence: On Llama 3.1 8B, the baseline reaches 18.86 perplexity versus 60.61 for the 3B baseline, showing that ρ is necessary but not sufficient for catastrophic failure.The paper attributes the 8B model’s resilience to smoother per-layer weight statistics.

8 Domain-Dependent Degradation

Quantization degradation increases with domain distance from the C4 calibration set in the undercomplete 2 bpp regime. OA-EM mitigates this pattern by allocating codebook capacity according to output sensitivity.

  • Domain-distance gradient: Baseline degradation rises from 1.04× in-domain on C4 to 3.49× far out-of-domain on WikiText-2.The results use 2 bpp, b=8, and pre-PV-tuning evaluation.
  • Mechanism: When ρ ≫1, limited codebook capacity is concentrated on calibration-important groups, increasing reconstruction error under a different evaluation distribution.The layer-wise objective implicitly weights weight groups by their calibration importance.
  • OA-EM response: OA-EM mitigates domain-shift degradation by distributing capacity according to Hessian-weighted output sensitivity rather than calibration frequency alone.The resulting codebooks better preserve linguistic representations that remain important across domains.

9 Conclusion

The paper argues that codebook initialisation determines the optimisation basin in additive quantization, with the effect governed by representational ratio ρ. Across three models, it reports persistent basin differences, practical efficiency gains, and scope boundaries for the method.

  • 9 Conclusion: Initialisation determines the optimisation basin, with severity governed by the representational ratio ρ.In undercomplete settings, greedy initialisation can trap optimisation; in overcomplete settings, its effect attenuates but remains.
  • 9 Conclusion: OA-EM at beam 4 takes 6.1h and outperforms the greedy baseline at beam 16, which takes 16.9h.The reported comparison corresponds to a 2.8× speedup.
  • 9 Conclusion: Basin persistence holds across Llama 3.2 3B, Llama 3.1 8B, and Qwen 2.5 3B.The models span two architecture families and multiple model sizes.
  • 9 Conclusion: The downstream accuracy signal is clearest on Llama 3.2 3B, while Qwen’s baseline slightly leads average accuracy at 0.606 versus 0.603.The authors note that larger models would require multi-GPU infrastructure and that the evaluation uses English-only models and benchmarks.
  • 9 Conclusion: The method targets free-form additive quantization and does not directly transfer to lattice- or trellis-based methods.The paper does not compare absolute perplexity against GLVQ or QTIP and evaluates English-only models and benchmarks.
  • 9 Conclusion: The work improves efficiency of existing openweight LLMs without introducing new training data or models.Compression reduces deployment costs and environmental impact but does not address biases in the original models.

A Overcomplete Regime: 3 bpp Detailed Results

At 3 bpp, OA-EM improves several pre- and post-PV-tuning outcomes on Llama 3.2 3B, while the initialisation bottleneck is comparatively mild. The reported results also show downstream gains and a quality–compute advantage for OA-EM.

  • A Overcomplete Regime: 3 bpp Detailed Results: At 3 bpp, Table 6 reports pre-PV-tuning perplexity for Llama 3.2 3B at beam 8 and 100 epochs.The table is identified as the 3 bpp pre-PV-tuning perplexity evaluation.
  • A Overcomplete Regime: 3 bpp Detailed Results: After PV-tuning, the perplexity gap shrinks from 0.65 to 0.12 points.WikiText-2 changes from 8.66 versus 8.54, while C4 changes from 11.43 versus 11.45.
  • A Overcomplete Regime: 3 bpp Detailed Results: OA-EM wins 4 of 6 downstream tasks and ties on 1 after PV-tuning.The largest reported gains are ARC-Easy at +3.5pp and LAMBADA accuracy at +1.6pp.
  • A Overcomplete Regime: 3 bpp Detailed Results: Average accuracy increases by 0.7 percentage points, from 0.647 to 0.654.The only notable regression is WinoGrande at −1.2pp.
  • A Overcomplete Regime: 3 bpp Detailed Results: OA-EM wins or ties average accuracy in every reported 2 bpp beam setting.Its largest advantage is at beam 4, where it reaches +1.7pp over the greedy baseline.
  • A Overcomplete Regime: 3 bpp Detailed Results: At 2 bpp, every OA-EM configuration produces a better final model than any greedy configuration at equal or lower compute.The comparison is the reported quality–compute frontier result for Llama 3.2 3B.

D Proof of Proposition 1

The supplied passages combine a proof residual-mismatch statement with table descriptions of post-PV-tuning evaluations across models and beam settings. They establish the reported comparison structure but provide limited proof-specific detail.

  • D Proof of Proposition 1: The proof states that the residual mismatch is nonnegative because j* minimises over r*.This is the stated justification for the residual term in the derivation.
  • D Proof of Proposition 1: At 2 bpp and beam 4, OA-EM wins or ties every downstream metric in the Llama 3.2 3B evaluation.The table is a post-PV downstream evaluation with 100 epochs.
  • D Proof of Proposition 1: The 2 bpp evaluation includes post-PV downstream tables for Llama 3.2 3B at beams 8 and 16.The supplied table labels identify the corresponding beam configurations.
  • D Proof of Proposition 1: The supplied results include post-PV downstream evaluation for Llama 3.1 8B at beam 8 and 100 epochs.The associated table is identified as Table 13.
  • D Proof of Proposition 1: The supplied results include post-PV perplexity and downstream evaluation for Qwen 2.5 3B at beam 8 and 100 epochs.Tables 14 and 15 provide the corresponding perplexity and downstream evaluations.
  • D Proof of Proposition 1: Table 16 consolidates the 2 bpp quality–compute Pareto analysis on Llama 3.2 3B using post-PV metrics and six accuracy tasks.Avg denotes the mean of the six accuracy tasks.
Loading 2604.08118v1…