PaperlayerView exact revision on arXiv

Source-linked AI summary

Extreme Compression of Large Language Models via Additive Quantization

Vage Egiazarian, Andrei Panferov, Denis Kuznedelev, Elias Frantar, Artem Babenko, Dan Alistarh

arXiv:2401.06118v4cs.LGcs.CL

TL;DR

Accurate open LLMs remain difficult to deploy locally because their computational and memory costs are high, especially under extreme 2–3-bit compression. The paper introduces AQLM, an input-adaptive additive quantization method with joint transformer-block optimization, and reports Pareto-optimal compression below 3 bits per parameter alongside practical CPU and GPU inference. AQLM is more computationally expensive to train than direct post-training methods, although efficient implementations reduce its execution cost.

  • Problem

    Open LLMs require reduced computational and memory costs for local inference or fine-tuning, while existing extreme 2-bit quantization can be inferior to smaller models quantized at 3–4 bits per parameter.

  • Method

    AQLM extends Additive Quantization with activation-aware optimization of layer outputs and joint quantization-parameter tuning across transformer layers.

  • Results

    AQLM is the first algorithm to achieve Pareto-optimality below 3 bits per parameter and outperforms prior methods across the standard 2–4-bit compression range, especially at 2 bits.

  • Takeaways & Limitations

    AQLM enables massive LLMs to be executed accurately and efficiently with little memory, with implementations supporting practical CPU and GPU generation.

  • Takeaways & Limitations

    AQLM is more computationally expensive than direct post-training quantization methods because it uses a more complex coding representation.

Abstract

from arXiv · show

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.

1. Introduction

AQLM targets the accuracy–compression limits of extreme LLM quantization by adapting multi-codebook methods to preserve layer outputs, especially below 3 bits per parameter. It also provides efficient CPU and GPU implementations for practical generation.

  • Open LLMs can be run or fine-tuned locally if their computational and memory costs become manageable on commodity hardware.
  • Quantization reduces stored weight bit-width, improving model footprint and memory transfer while introducing an accuracy–compression trade-off.
  • Existing methods often make extreme 2-bit quantization less accurate than using a smaller model quantized at 3–4 bits per parameter for the same byte budget.
  • AQLM extends Additive Quantization by adapting optimization to calibration activations and jointly tuning quantization parameters across transformer layers.
  • Across 2–4 bits per parameter, AQLM outperforms the previous state of the art, with its largest improvements in extreme 2-bit quantization.
  • AQLM implementations match or outperform floating-point speed while reducing memory footprint by up to 8x; reported layer-wise speedups reach ∼30% on GPUs and 4x on CPUs.

2. Background & Related Work

Prior LLM quantization methods mainly use direct weight projections, while related MCQ methods jointly encode vector values through learned codebooks for efficient search. AQLM builds on additive MCQ to optimize LLM layer outputs and jointly train codebooks across layers.

  • LLM Quantization: Early LLM post-training quantization methods used direct round-to-nearest projections and adjusted granularity to balance memory efficiency and accuracy.
  • LLM Quantization: The prior Pareto frontier extended below 4 bits per weight, whereas this work brings it below 3 bits per weight for the first time.
  • LLM Quantization: Recent methods address weight outliers using sparse higher-precision storage, per-channel scaling, or Fisher-informed quantization.
  • LLM Quantization: QuIP and QuIP# rotate weights and map them onto a lattice to reduce worst-case layer error, whereas AQLM learns additive codebooks directly from calibration data.
  • Quantization for Nearest Neighbor Search: Vector quantization compresses multiple vector dimensions jointly by learning codebooks whose candidate vectors encode database vectors.
  • Quantization for Nearest Neighbor Search: Product Quantization splits vectors into subvectors and represents each with a codeword from a separate codebook, enabling distance computation through precomputed lookup tables.
  • Quantization for Nearest Neighbor Search: Additive Quantization represents vectors as sums of codewords from unconstrained codebooks, generally reducing compression error but complicating training for large codebook counts.
  • Quantization for Nearest Neighbor Search: Earlier vector-quantization compression methods introduced prediction errors that prevented deep-model compression, while this work adapts and scales MCQ to LLMs.

3. AQLM: Additive Quantization for LLMs

AQLM extends additive quantization to compress LLM weights while preserving layer outputs under calibration inputs. It combines additive codebooks, input-aware optimization, efficient code search, codebook tuning, and joint block-level fine-tuning, with practical calibration costs.

  • 3.1. Overview: The quantization objective minimizes squared error between original and compressed layer outputs on calibration inputs, rather than preserving weights alone.The formulation uses a weight matrix W, calibration inputs X, and quantized weights cW; the objective is generalized across output units and input groups.
  • 3.1. Overview: AQLM represents each weight group as a sum of vectors selected from multiple learned codebooks, using M·B bits per group plus FP16 codebook storage.Groups contain g consecutive weights; each codebook provides 2^B vectors, and one code from every codebook is summed to reconstruct the group.
  • 3.2. Phase 1: Beam search for codes: Codes are initialized by residual K-means, then updated with beam search that tests code replacements and retains the best configurations under the MSE objective.The beam maintains k candidate configurations and evaluates all 2^B alternatives for the selected code position, with incremental loss updates enabled by precomputed XX^T.
  • 3.3. Phase 2: Codebook update: AQLM tunes codebooks by approximately solving a least-squares problem with Adam and full-batch gradient descent using the precomputed XX^T matrix.The presence of XX^T couples codebook coordinates, making the original AQ closed-form coordinate-wise optimization unsuitable without larger matrix inversions or approximations.
  • 3.4. Phase 3: Fine-tuning for intra-layer cohesion: Joint optimization across transformer blocks addresses interactions among matrix quantization errors, which become especially important at extreme 2-bit compression.The procedure fine-tunes block parameters using calibration data after layer-wise compression, while changing only a few trainable parameters.
  • 3.4. Phase 3: Fine-tuning for intra-layer cohesion: Block fine-tuning remains feasible for billion-parameter models on a single GPU and consumes 10–30% or less of total calibration time.It uses little VRAM for optimizer states and converges after a few iterations from a good initial guess.

4. Experiments

Experiments evaluate AQLM on LLAMA 2 and Mixtral across extreme 2–4-bit compression, measuring perplexity, zero-shot accuracy, Pareto efficiency, calibration effects, and inference speed. AQLM consistently improves over prior PTQ methods, with its largest quality gains at about 2 bits and practical CPU/GPU implementations.

  • 4.1. Compression quality for modern LLMs: AQLM outperforms previous best post-training quantization algorithms across compression settings, with the largest gains at 2–2.1 bits per parameter.The advantage appears in both perplexity on WikiText-2 and C4 and accuracy across zero-shot tasks, where all methods diverge substantially from the uncompressed model.
  • 4.1. Compression quality for modern LLMs: AQLM also outperforms QuIP# on 2-bit Mixtral, including a +3-point gain on the harder Arc Challenge task.The margins are smaller than for LLAMA 2 but remain significant on harder tasks.
  • 4.1. Compression quality for modern LLMs: AQLM is the first method reported as Pareto-optimal below 3 bits per parameter.Its apparent Pareto-optimal point is around 2.5 bits per parameter; 2.76-bit AQLM on 13B outperforms the uncompressed 7B model.
  • 4.2. End-to-end fine-tuning experiments: End-to-end fine-tuning improves both AQLM and QuIP#, with larger benefits at 2 bits and diminishing returns from 3 bits upward.A 2.19-bit fine-tuned AQLM model on 13B is comparable to an uncompressed 7B model on zero-shot tasks.
  • 4.3. Ablation analysis: Increasing calibration sequences from 128 to 4096 gradually improves perplexity, but with diminishing returns.AQLM benefits more from larger calibration sets than GPTQ, whose accuracy saturates around 256 input sequences.
  • 4.4. Inference Speed: Smaller multiple codebooks trade slightly lower accuracy for efficient GPU cache utilization, while AQLM reaches up to 4× faster CPU inference than FP32.The codebook decomposition increases quantization error but still outperforms baselines in accuracy.

5. Conclusion and Future Work

AQLM improves low-bit LLM quantization, especially at 2–3 bits, while supporting efficient inference with reduced memory. The paper also identifies computational cost and unexplored fine-tuning and application settings as boundaries.

  • AQLM uses additive quantization with block-wise optimization to preserve compressed-model performance.
  • AQLM requires more computationally expensive quantization than direct methods such as RTN or GPTQ, although inference remains efficiently implementable.
  • Future work includes better fine-tuning strategies and adapting AQLM to computer vision models or long-sequence attention-cache compression.
  • Knowledge distillation trains the quantized model to mimic the floating-point model using KL divergence between teacher and student outputs.
  • The fine-tuning procedure often significantly improves compressed-model performance.

C. Experimental Configurations

The experiments use specified hardware, calibration data, quantization configurations, and optimization settings across LLAMA, Mistral, and Mixtral models. AQLM calibration is substantially slower than simpler methods, with runtime depending on configuration and parallelization.

  • Experiments use Nvidia A100 or H100 GPUs, while inference speed is evaluated on an Nvidia 3090 GPU and Intel Core i9 13900K CPU.
  • All methods are calibrated on RedPajama-v1 using model-specific context lengths: 4096 for LLAMA 2 and 8192 for Mistral/Mixtral.
  • AQLM uses one codebook for 2 bits, two codebooks for 3 bits, and two codebooks for 4 bits, with groups of 8.
  • Adam optimization uses learning rate 10^-4, β1 = 0.90, β2 = 0.95, and early stopping based on least-square-error improvement.
  • AQLM quantization takes considerably longer to calibrate than RTN or GPTQ, but this affects quantization time rather than inference time.
  • Quantizing a 7B model takes about 1 day on one A100, while a 70B model takes 10–14 days on one GPU or 3–4 days on eight GPUs.
  • Reducing quantization beams can produce 2–4x speedups, at the cost of lower model accuracy.

E. Ablation analysis

Ablations examine initialization, fine-tuning restrictions, calibration-set size, codebook structure, and higher-bit settings. The results emphasize K-means initialization and fine-tuning as important contributors to compressed-model performance.

  • Initialization: Residual K-means initialization requires significantly fewer training iterations than random initialization to reach the desired loss.
  • Fine-tuning: Full-model fine-tuning or fine-tuning only AQ parameters achieves competitive performance, whereas training only RMSNorm scales is comparable to no fine-tuning.
  • Number of samples: Calibration-set experiments vary AQLM and baselines from 128 to 4096 sequences to assess calibration requirements and potential overfitting.
  • Number of codebooks vs groups: At approximately 2-bit compression, WikiText-2 perplexity is evaluated across different codebook and group configurations.

F. Additional experiments

Additional experiments report WikiText-2 perplexity for approximately 2-bit LLAMA 2 quantization and extend evaluation to Mixtral, Mistral7B, and LLAMA 2 models.

  • Table 9 reports WikiText-2 perplexity as a function of groups and codebook number for approximately 2-bit LLAMA 2 7B quantization.
  • Additional experiments cover Mixtral, Mistral7B, and LLAMA 2 models.

F.1. Mixtral

AQLM is evaluated on Mixtral and other quantized LLM settings across standard perplexity and zero-shot accuracy metrics. Results are strongest relative to prior methods in several low-bit regimes, while 4-bit Mixtral performance is similar to QuIP# and close to FP16.

  • F.1. Mixtral: In 4-bit Mixtral, AQLM and QuIP# perform very similarly across all reported metrics and remain close to the uncompressed FP16 model.The evaluation reports WikiText-2 and C4 perplexity plus average zero-shot accuracy.
  • F.1. Mixtral: AQLM outperforms other methods in 4-bit LLAMA 2 perplexity, achieving the best or near-best results.The section also reports results for LLAMA 2 models using two 8-bit codebooks.
  • F.1. Mixtral: Figure 5 compares AQLM with QuIP# on LLAMA 2 7B, 13B, and 70B models.The plotted comparison uses quantized model results across multiple model sizes.
  • F.1. Mixtral: On Mistral 7B, QuIP# slightly outperforms AQLM on most 2-bit benchmarks, while 4-bit results are very close across methods.The comparison covers 2-, 3-, and 4-bit configurations using perplexity and zero-shot accuracy.

G. Pareto optimality

AQLM improves the perplexity-versus-size trade-off for LLAMA 2 models across several model scales. The results also show that model size can change which compressed model is preferable at a fixed memory budget.

  • G. Pareto optimality: AQLM achieves lower WikiText-2 perplexity than QuIP# across LLAMA 2 7B, 13B, and 70B model sizes.Figure 5 plots perplexity against quantized weight size in bytes.
  • G. Pareto optimality: Figure 5 evaluates the AQLM–QuIP# comparison by plotting WikiText-2 perplexity against quantized weight size.This frames the comparison as a model-size versus perplexity trade-off.
  • G. Pareto optimality: Around 3.7 GiB, corresponding to 2.5 bits for LLAMA 2 13B, compressing the 13B model is more advantageous than compressing the 7B model at the same byte size.The comparison is based on WikiText-2 perplexity for AQLM.

H. Estimating model size

The paper estimates compressed-model size from codebooks, codes, and per-unit scales, then relates these costs to practical inference measurements and learned-code structure.

  • H. Estimating model size: Quantized-weight storage consists of codebooks, codes, and per-unit scales, with codebooks and scales stored in half precision.The total memory depends on input and output dimensions, group size, codebook count, and code width.
  • H. Estimating model size: Average bits per parameter are computed from the combined storage costs of codebooks, codes, and scales.The formula accounts for group size, codebook count, code width, and layer dimensions.
  • H. Estimating model size: For a LLAMA 2 70B MLP gate projection, group size 8 and two 8-bit codebooks require 2.002 bits per parameter.The example uses din = 8192 and dout = 28672; storage is typically dominated by codes.
  • I. End-to-End Inference Speed: Figure 7 visualizes code-selection frequencies and the two leading principal components of a learned codebook for a specific layer.The layer shown is 5.self_attn.q_proj.
  • I. End-to-End Inference Speed: Inference speed is benchmarked by generating 128 tokens with batch size 1 on a compiled graph using an RTX 3090 GPU and an Intel i9 CPU.The reported metric is average generated tokens per second.

K. Evaluation on MMLU and GSM8k

The paper extends evaluation beyond perplexity and simple zero-shot tasks to MMLU and GSM8k, and also examines block tuning for low-bit quantization.

  • K. Evaluation on MMLU and GSM8k: MMLU and GSM8k are used to evaluate LLAMA 2 models quantized at 2–2.1 bits per parameter.The evaluation includes end-to-end finetuning configurations marked with a star.
  • K. Evaluation on MMLU and GSM8k: The paper motivates MMLU evaluation because WikiText-2, C4, and simple zero-shot tasks may not exhaustively represent real-world LLM abilities.MMLU covers problems from 57 domains.
  • K. Evaluation on MMLU and GSM8k: Block tuning can be applied to scalar quantization by optimizing full-precision quantization scales, but tuned GPTQ remains far below AQLM at similar bit widths.The reported observation concerns low-bit-width GPTQ quality.
Loading 2401.06118v4…