🤖 machine learning

High-Rate Quantized Matrix Multiplication II

This paper investigates high-rate quantized matrix multiplication for weight-only LLM post-training quantization by demonstrating how covariance-aware waterfilling improves upon equal-rate allocation, analyzing the basis-free and near-optimal performance of the WaterSIC scheme, and showing that GPTQ with random rotation achieves comparable near-optimal results.

Original authors: Or Ordentlich, Yury Polyanskiy

Published 2026-05-14

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Or Ordentlich, Yury Polyanskiy

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: Compressing the Brain of AI

Imagine a massive Artificial Intelligence (AI) model, like a giant library of knowledge, trying to solve a problem. To do this, it performs a mathematical operation called Matrix Multiplication (MatMul) billions of times. Think of this as the AI's "thinking" process.

However, these "weights" (the numbers inside the AI) take up a huge amount of memory. To make the AI run faster and on smaller devices, engineers compress these numbers, a process called Quantization. It's like shrinking a high-resolution photo into a smaller file size so it loads faster.

This paper is the second part of a study on how to do this shrinking most efficiently. While the first part looked at shrinking without any prior knowledge, this paper focuses on a scenario where we do have some prior knowledge: we know the statistical "shape" of the data the AI is processing.

The Core Problem: The "Weight-Only" Puzzle

In many modern AI systems (like LLMs), the data flowing through the system (activations) is kept in full precision, but the "weights" (the static knowledge) are compressed.

The Goal: We want to compress the weights ( $W$ ) as much as possible without making the AI's answers ( $Y$ ) too wrong.
The Catch: The "wrongness" (distortion) depends on how the weights interact with the incoming data. If the incoming data has a specific pattern (like a long, thin oval shape), compressing the weights in a standard, square-grid way is inefficient. It's like trying to fit a long, thin suitcase into a square box; you waste a lot of space.

The Old Way: "One Size Fits All" (GPTQ)

Currently, popular methods like GPTQ treat every part of the weight matrix the same. They use a standard grid (like graph paper) to round off numbers.

The Analogy: Imagine you are packing a suitcase with items of different sizes. The old method uses a grid of identical boxes for everything. You put a tiny pebble in a big box and a giant rock in a big box. You waste space on the pebble and might not fit the rock perfectly.
The Flaw: This method doesn't look at the specific shape of the data. It assumes the data is perfectly round (isotropic). If the data is actually an oval, this method is suboptimal.

The New Theory: "Waterfilling" (The Ideal Solution)

The authors look at the math of Weighted Mean Squared Error (WMSE). They use a concept called Waterfilling.

The Analogy: Imagine you have a landscape with hills and valleys (representing the importance of different parts of the data). You want to pour a fixed amount of "water" (your limited bits/bandwidth) onto this landscape.
- The Waterfilling strategy says: Pour the water into the deepest valleys first. These are the parts of the data that matter most or are most sensitive. You give them more "resolution" (more bits).
- The shallow hills get less water (fewer bits).
- This ensures you get the best possible picture for the amount of water you have.

Mathematically, this is the "Information-Theoretic Limit"—the absolute best you can possibly do.

The Practical Solution: WaterSIC

The problem with Waterfilling is that it's hard to calculate in real-time. The authors propose a practical algorithm called WaterSIC.

How it works: It uses a technique called Successive Interference Cancellation (SIC). Imagine you are trying to hear a conversation in a noisy room. You listen to the loudest voice first, understand it, and then "cancel it out" so you can hear the quieter voices better.
The Innovation: WaterSIC applies this idea to the weights. It looks at the data's shape (using a mathematical tool called Cholesky decomposition) and adjusts the "grid size" for each part of the weight.
- For parts of the data that are "stiff" or important, it uses a finer grid (more bits).
- For parts that are "loose" or less important, it uses a coarser grid (fewer bits).

Key Findings

Near-Perfect Efficiency: The authors prove that WaterSIC is incredibly close to the theoretical "Waterfilling" limit. It is only about 0.25 bits away from the absolute best possible performance. This is a tiny gap, meaning the method is almost perfect.
Rotation Immunity: One of the coolest findings is that WaterSIC is "basis free."
- The Analogy: Imagine you have a map. If you rotate the map, the North Pole moves. Some compression methods break or get worse if you rotate the data (like rotating the map). WaterSIC, however, works just as well regardless of how the data is rotated. It adapts to the shape of the data, not the direction it's facing.
GPTQ is Surprisingly Good (with a twist): The paper also found that the standard GPTQ method (which doesn't use the fancy water-filling adjustments) performs surprisingly well if you randomly rotate the data first.
- It turns out that the standard way AI data is currently organized is actually "lucky" (it's close to the optimal shape). But if you mess with it (rotate it), GPTQ gets worse, while WaterSIC stays strong.

Summary

This paper introduces WaterSIC, a smarter way to compress AI weights.

Old Way: Use a standard grid for everything (GPTQ).
New Way: Look at the shape of the data and adjust the grid size for each part (WaterSIC).
Result: The new method is nearly perfect, theoretically unbeatable, and robust even if the data is shuffled or rotated. It bridges the gap between complex math theory and practical, fast AI compression.

The authors conclude that while current methods are good, there is still room to improve, especially for very low-bit compression (where you have very few bits to work with) and for making the "shaping" of data faster on computer chips.

Technical Summary: High-Rate Quantized Matrix Multiplication II

Problem Statement
This paper investigates the theoretical and practical limits of weight-only quantization for Large Language Models (LLMs). The specific setting assumes that the activation matrix $X$ is available in full precision during inference, while the weight matrix $W$ must be quantized to a low-resolution representation $\hat{W}$ to reduce memory bandwidth and storage requirements. The objective is to minimize the expected weighted mean squared error (WMSE) distortion:
$D = \frac{1}{n} \mathbb{E} \| X^\top (W - \hat{W}) \|_F^2 = \frac{1}{n} \text{tr}((W - \hat{W})^\top \Sigma_X (W - \hat{W}))$
where $\Sigma_X = \mathbb{E}[XX^\top]$ is the covariance matrix of the activations, which is assumed to be known (via calibration) but the decoder (inference engine) does not have access to $\Sigma_X$ or its basis, only the quantized bits.

Methodology and Theoretical Framework
The authors approach the problem through two lenses: information-theoretic limits and practical algorithmic design.

Information-Theoretic Limits (High-Rate Regime):
- Isotropic vs. Anisotropic Coding: The paper contrasts "isotropic" quantization (using a standard codebook oblivious to $\Sigma_X$ ) with "anisotropic" quantization (adapting to $\Sigma_X$ ).
- Waterfilling Solution: Drawing on classical source coding theory, the authors derive that the optimal distortion $D^*$ in the high-rate regime is achieved by a reverse waterfilling strategy. This allocates quantization rates unevenly across coordinates based on the eigenvalues of $\Sigma_X$ , resulting in a distortion proportional to the geometric mean of the eigenvalues: $D^* \propto |\Sigma_X|^{1/n} 2^{-2R}$ .
- The Gap: Standard isotropic coding yields distortion proportional to the arithmetic mean of the eigenvalues. The gap between the optimal and isotropic solutions is governed by the Arithmetic Mean-Geometric Mean (AM-GM) inequality.
- Oblivious Decoder: Crucially, the paper demonstrates that even if the decoder is $\Sigma_X$ -oblivious (cannot apply the SVD basis), an adaptive encoder using an isotropic codebook can asymptotically achieve the waterfilling distortion limit in the high-rate regime, provided the rounding to the codebook is optimal.
Practical Algorithms (Successive Interference Cancellation):
- The paper analyzes the GPTQ (and equivalent LDLQ) algorithm, which uses a fixed rate per entry and a successive interference cancellation (SIC) approach (Babai's algorithm) based on the Cholesky decomposition $\Sigma_X = U^\top U$ .
- Theoretical Limit of GPTQ: The authors show that standard GPTQ (with equal spacing) achieves distortion proportional to the arithmetic mean of the squared diagonal elements of $U$ ( $U_{i,i}^2$ ). This performance is sensitive to the basis of $\Sigma_X$ ; specifically, it performs worst in the PCA basis (where $U$ is diagonal) and better in random bases.
- WaterSIC Algorithm: Motivated by the waterfilling theory, the authors propose WaterSIC. This algorithm adjusts the quantization step size (spacing) for each coordinate based on the diagonal elements of the Cholesky factor $U$ . Specifically, it allocates more bits to coordinates with larger $U_{i,i}$ values.
- Performance: WaterSIC achieves distortion proportional to the geometric mean of $U_{i,i}^2$ . Since the product of $U_{i,i}$ equals $\sqrt{|\Sigma_X|}$ , WaterSIC attains a distortion within a multiplicative factor of $\frac{2\pi e}{12}$ (approx. 1.423, or 0.25 bits/entry) of the information-theoretic lower bound.

Key Results

Near-Optimality of WaterSIC: The WaterSIC algorithm is shown to be "basis-free" (its performance depends only on the determinant of $\Sigma_X$ , making it immune to random rotations of the input space) and is within 0.25 bits/entry of the theoretical optimum in the high-rate regime.
GPTQ Performance: While standard GPTQ is suboptimal in specific bases, the authors find that for Llama-3-8B layers, GPTQ with a random rotation of the weights performs surprisingly well, often within 0.1 bits of WaterSIC. This suggests that the standard basis used in LLMs (likely influenced by Adam optimizer dynamics) is not the worst-case scenario for GPTQ, and random rotation can mitigate basis sensitivity.
Empirical Validation: Experiments on Llama-3-8B layers using Wikitext-2 calibration data confirm that WaterSIC with entropy coding closely tracks the theoretical waterfilling limit, while standard GPTQ shows variable performance depending on the layer and basis.

Significance and Claims
The paper claims to bridge the gap between abstract information theory and practical LLM quantization:

Theoretical Justification: It provides a rigorous information-theoretic justification for why per-coordinate rate allocation (waterfilling) is superior to uniform allocation in weight-only quantization.
Algorithmic Improvement: It introduces WaterSIC as a practical, low-complexity algorithm that realizes these theoretical gains, offering a provable path to near-optimal distortion.
Insight into Existing Methods: It clarifies the behavior of GPTQ, explaining its sensitivity to the basis of $\Sigma_X$ and demonstrating that random rotations can make it nearly optimal, thereby validating its widespread use while highlighting the potential for further gains via WaterSIC.
Limitations and Future Work: The authors modestly note that their high-rate approximations assume high resolution. They identify that at low rates (e.g., < 2 bits), other factors like dimensionality reduction and specific weight structures become critical. Furthermore, they highlight that while entropy coding (EC) is theoretically optimal, its computational cost on GPUs may be prohibitive, and efficient shaping methods for GPUs remain an open challenge. They also note that the 0.25-bit gap of WaterSIC stems from the sub-optimality of the integer lattice $\mathbb{Z}^n$ , which could be reduced by using higher-dimensional lattices, though this complicates entropy coding.