Optimal Formats for Weight Quantisation

TL;DR

This paper investigates optimal formats for weight quantisation by modeling the end-to-end effect of quantisation as KL divergence under a memory constraint and linking it to classical quantisation theory via Fisher information. It presents non-uniform, block-based quantisers, notably the $ oot3 ext{p}$ family and Signmax scaling, and shows that variable-length encoding—either through block scaffolding or compression—consistently yields superior performance over fixed-length codes. It also develops a principled per-tensor bit-width allocation rule guided by Fisher information and validates the framework with simulations and extensive experiments on large language models, including QAT scenarios. The findings provide theoretical insight and practical guidelines for designing quantisation formats that approach the efficiency of lossless compression while preserving model fidelity on downstream tasks.

Abstract

Weight quantisation is an essential technique for enabling efficient training and deployment of modern deep learning models. However, the recipe book of quantisation formats is large and formats are often chosen empirically. In this paper, we propose a framework for systematic design and analysis of quantisation formats. By connecting the question of format design with the classical quantisation theory, we show that the strong practical performance of popular formats comes from their ability to represent values using variable-length codes. We frame the problem as minimising the KL divergence between original and quantised model outputs under a model size constraint, which can be approximated by minimising the squared quantisation error, a well-studied problem where entropy-constrained quantisers with variable-length codes are optimal. We develop non-linear quantisation curves for block-scaled data across multiple distribution families and observe that these formats, along with sparse outlier formats, consistently outperform fixed-length formats, indicating that they also exploit variable-length encoding. Finally, by using the relationship between the Fisher information and KL divergence, we derive the optimal allocation of bit-widths to individual parameter tensors across the model's layers, saving up to 0.25 bits per parameter when applied to large language models.

Figures (35)

• Figure 1: The trade-off between average bits per parameter and top-$k$ KL divergence for Llama 3.1 8B. To approach optimal performance, some form of variable-length encoding is needed: lossless compression, block (or channel) absmax scaling or $0.1\%$ sparse outlier removal. The shaded line width is $\pm 2$ standard error over evaluation data. See \ref{['fig:xp_tradeoff_overview_grid']} for other models.
• Figure 2: 4-bit quantisation curve gradients for Normal, Laplace, Student-t distributions, with strong agreement between cube root density and Lloyd-Max, which optimises codepoints against samples directly. The legend shows relative quantisation error $R$ for "(cube root quantiser; Lloyd-Max)" for data matching the quantiser. (Left) RMS-scaled formats. The cube root density rule breaks down with the heavy tails of Student-t. (Right) absmax-scaled formats. The discrepancy at the extremes occurs because the cube root quantiser has a special case for $\pm 1$, whereas Lloyd-Max treats it as a single distribution to quantise.
• Figure 3: 3-bit $\sqrt[3]\mathrm{p}$ codepoint distributions for normally distributed data, illustrating RMS, absmax and signmax scaling methods and symmetric/asymmetric variants (with $B=64$ for block formats). The principal benefit of asymmetric variants is that they have an encoding for $\color{blue}{\symbfit{0}}$. INT formats are asymmetric, while most floating-point formats are symmetric but represent zero twice ($\pm 0$).
• Figure 4: The error/size tradeoff for different data distributions (column) and optimal quantisers (hue). Surprisingly, block absmax quantisers can outperform tensor RMS formats for iid data, even though there is no inherent block structure. However this situation is reversed when adding optimal compression, implying that block absmax quantisers exploit some form of variable-length coding.
• Figure 5: A 2D histogram of bits $\beta_i$ used to encode parameter $i$ from the first MLP down-projection from Llama 3.1 8B, illustrating how different schemes achieve variable-length encoding. (Left) sparse outliers create a distinct step between regular values and the $0.1\%$ largest absolute values. (Center) block absmax can be seen as using the bfloat16 scale to represent the block maximum and fewer bits for everything else. The histogram has an overlap, since the maximum is per-block not global. (Right) lossless compression on a uniform grid with $\beta_i = -\log_2 p_i$, where $p_i$ is the proportion of parameters assigned to that quantisation bucket.