🔢 mathematics

Optimal Scalar Quantization for Matrix Multiplication: Closed-Form Density and Phase Transition

This paper derives the optimal entrywise scalar quantization strategy for matrix multiplication by establishing a closed-form point density that minimizes mean-squared error, revealing a correlation-driven phase transition where the optimal density shifts from unimodal to bimodal as the correlation between matrix entries exceeds a critical threshold.

Original authors: Calvin Ang, Sungyoon Kim, Mert Pilanci

Published 2026-03-23

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Calvin Ang, Sungyoon Kim, Mert Pilanci

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

Imagine you are running a massive kitchen where you need to multiply two giant lists of ingredients (Matrices A and B) to create a final recipe (Matrix C). In the world of modern AI, these lists are so huge that they don't fit on your counter (your computer's memory). To solve this, you decide to shrink the lists by rounding off the numbers, a process called quantization.

Usually, when people shrink these lists, they try to make sure the individual numbers look as close to the original as possible. They ask, "Is my rounded number 5.1 or 5.2?"

This paper says: "Stop worrying about the individual numbers! Worry about the final dish."

The authors argue that if you are going to multiply these lists together, you should round the numbers in a way that makes the final multiplication result as accurate as possible, even if the individual numbers look a bit weird.

Here is the breakdown of their discovery using simple analogies:

1. The "Filter" Effect

Imagine you are trying to guess the total weight of a bag of apples (Matrix A) multiplied by the price per apple (Matrix B).

Old Way: You try to guess the weight of every single apple perfectly.
New Way: The authors realized that if an apple is very heavy, a small error in its weight matters a lot. But if an apple is tiny, a small error doesn't matter much.
The Insight: They found that the "importance" of a number depends on what it is being multiplied by. If a number in Matrix A is going to be multiplied by a huge number in Matrix B, you need to be very precise with that number. If it's multiplied by a tiny number, you can be sloppy.

They created a mathematical rule that tells you exactly how to round the numbers based on this "partner" relationship.

2. The "Shape-Shifting" Density (The Phase Transition)

This is the most magical part of the paper.

Imagine you are trying to pack marbles (data points) into a box.

Scenario A (Low Correlation): If the two lists of numbers are unrelated (like the height of people and the price of tea), the best way to pack your marbles is to put the most marbles in the middle (where most numbers live) and fewer on the edges. It's a classic "hill" shape.
Scenario B (High Correlation): If the two lists are strongly related (like height and weight), the authors discovered something surprising. The best way to pack the marbles changes shape!
- Instead of one big hill in the middle, the "perfect" packing shape splits into two hills with a valley in the middle.
- The Analogy: Imagine you are trying to catch fish in a river. If the fish are swimming randomly, you put your net in the middle of the river. But if the fish are swimming in two distinct schools on the left and right banks, you move your net to catch both schools, leaving the middle empty.

The paper proves that this "splitting" happens exactly when the relationship between the two numbers gets stronger than a specific threshold (about 57% correlation). It's a sudden "phase transition," like water turning into ice, but for data packing.

3. Why This Matters for AI (The "LLM" Connection)

Large Language Models (like the ones that write this response) rely on massive matrix multiplications to think.

The Problem: These models are getting too big for our computers. We need to shrink them (quantize them) to run them on phones or laptops.
The Current Fix: We use standard rounding methods (like INT8 or FP8) that treat every number the same.
The Paper's Solution: By using their new "smart rounding" method, we can shrink the model more without losing as much intelligence.
- They tested this on famous models (GPT-2 and Qwen).
- The Result: Their method kept the model's "brain" working better than standard methods, especially in the "attention" parts of the model where the AI decides what words to focus on.

Summary

Think of this paper as a new set of instructions for packing a suitcase for a trip.

Old Rule: Pack your clothes neatly so every shirt looks perfect.
New Rule: Pack your clothes based on what you are doing at the destination. If you are going hiking, pack your heavy boots tightly and leave your fancy dress loose. If you are going to a gala, do the opposite.

The authors figured out the exact mathematical formula for "what to pack" when the destination is Matrix Multiplication. They found that when the data is highly connected, the best packing strategy actually leaves a gap in the middle, creating a two-hump shape that standard methods completely miss. This allows us to run bigger, smarter AI models on smaller, cheaper hardware.

1. Problem Formulation

The paper addresses the problem of entrywise scalar quantization of two matrices, $A \in \mathbb{R}^{m \times k}$ and $B \in \mathbb{R}^{k \times n}$ , prior to their multiplication.

Goal: Minimize the Mean Squared Error (MSE) of the resulting product matrix, defined as $E \triangleq \mathbb{E}[\|AB - \hat{A}\hat{B}\|_F^2]$ , where $\hat{A}$ and $\hat{B}$ are the quantized versions of $A$ and $B$ .
Constraint: The quantization is performed independently on entries of $A$ and $B$ using scalar quantizers with $K_X$ and $K_Y$ levels, respectively.
Context: Unlike traditional quantization which minimizes reconstruction error of the operands ( $\|A - \hat{A}\|^2$ ), this work focuses on the downstream distortion (the product error). This is critical for modern low-precision inference (e.g., INT8, FP8, NVFP4) in Large Language Models (LLMs) where throughput and memory efficiency are paramount.
Assumption: The paper assumes a pair-i.i.d. inner-product model, where the pairs of entries $(A_{i\ell}, B_{\ell j})$ feeding into each dot product are independent and identically distributed (i.i.d.) copies of a generic random pair $(X, Y)$ with finite moments.

2. Methodology: High-Resolution Analysis

The authors derive a sharp asymptotic characterization of the optimal quantization strategy in the high-resolution regime ( $K_X, K_Y \to \infty$ ).

A. Decoupling the Matrix Objective

The core theoretical insight is that the matrix multiplication MSE can be decoupled into two independent weighted scalar quantization problems.

Error Expansion: By expanding the error term $D = XY - \hat{X}\hat{Y}$ , the authors show that the dominant terms in the MSE are driven by the conditional second moments of the operands.
Weighted MSE: The problem reduces to minimizing:
$\mathbb{E}[w_X(X)(X - \hat{X})^2] + \mathbb{E}[w_Y(Y)(Y - \hat{Y})^2]$
where the weights are the conditional second moments:
- $w_X(x) = \mathbb{E}[Y^2 \mid X=x]$
- $w_Y(y) = \mathbb{E}[X^2 \mid Y=y]$
Asymptotic Scaling: The authors prove a sharp $K^{-2}$ scaling law for the MSE, identifying the exact leading constant.

B. Optimal Point Densities

Using high-rate quantization theory (specifically Bennett's integral and Hölder's inequality), they derive the asymptotically optimal point density $\lambda^*(x)$ for the companding quantizers:
$\lambda^*_X(x) \propto (f_X(x) w_X(x))^{1/3}$
$\lambda^*_Y(y) \propto (f_Y(y) w_Y(y))^{1/3}$
where $f_X$ and $f_Y$ are the marginal densities. This indicates that the optimal quantizer places more points where the product of the marginal probability and the conditional variance is high.

C. Special Case: Correlated Gaussian

The paper specializes the theory to the case where $(X, Y)$ follows a bivariate Gaussian distribution with correlation $\rho$ .

Closed-Form Density: They derive a closed-form expression for the optimal density:
$\lambda^*(u) \propto \exp\left(-\frac{u^2}{6}\right) \left((1-\rho^2) + \rho^2 u^2\right)^{1/3}$
where $u = x/\sigma_X$ .
Phase Transition: A major theoretical contribution is the discovery of a correlation-driven phase transition in the shape of the optimal density:
- Unimodal ( $|\rho| \leq 1/\sqrt{3}$ ): The density has a single peak at the origin.
- Bimodal ( $|\rho| > 1/\sqrt{3}$ ): The density develops two symmetric peaks at $u_{peak} = \pm\sqrt{3 - 1/\rho^2}$ , with a local minimum at the origin.
- This implies that for highly correlated inputs, the optimal quantizer should avoid quantizing values near zero and instead focus on the "tails" where the product error is most sensitive.

3. Key Contributions

High-Rate Characterization: Proved a sharp $K^{-2}$ asymptotic expansion for matrix multiplication MSE with an exact leading constant, decomposing the matrix problem into two weighted scalar problems.
Optimal Densities: Identified the optimal companding point densities driven by conditional second moments.
Phase Transition Discovery: For correlated Gaussians, derived a closed-form density and proved a unimodal-to-bimodal phase transition at $|\rho| = 1/\sqrt{3}$ .
Optimal Bit Split: Derived the optimal allocation of bits between matrices $A$ and $B$ based on their respective conditional moment integrals.
Empirical Validation: Validated the theory on synthetic data, quantized least squares, and real-world LLM activations (GPT-2 and Qwen3).

4. Experimental Results

The authors evaluated their "MatMul-Opt" quantizer against standard baselines (Uniform, Lloyd-Max, Gaussian Compander, $\mu$ -Law, NF4, NVFP4).

Synthetic Matrix Multiplication: The proposed optimal quantizer consistently outperformed all baselines in relative Frobenius norm error across various correlation levels $\rho$ .
Quantized Least Squares: In solving $\min_W \|XW - Y\|_F^2$ , the method showed that tuning the correlation parameter $\rho$ (even if it differs from the ground truth) significantly improves solution accuracy compared to naive marginal quantization, especially at low bit rates.
LLM Activation Quantization:
- Applied to Key (K) and Query (Q) activations in GPT-2 and Qwen3 models.
- Results: The $\rho$ -tuned quantizer achieved lower relative Frobenius error on attention logits compared to INT8 and FP8 baselines.
- Win Rates: For GPT-2 models, the proposed method won in 96.5% to 100% of attention heads against INT8 and FP8.
- Observation: The optimal $\rho$ values found via tuning were often higher than estimated correlations, suggesting model misspecification (e.g., due to rotary embeddings in Qwen3) but confirming the robustness of the correlation-aware approach.

5. Significance and Impact

Theoretical Advancement: This work bridges the gap between classical high-rate quantization theory and the specific bilinear distortion of matrix multiplication, providing the first closed-form optimal densities for this specific downstream task.
Practical Relevance: As hardware moves toward ultra-low precision (4-bit, 2-bit), understanding that the optimal quantizer for products differs fundamentally from the optimal quantizer for inputs is crucial. The discovery of the bimodal phase transition suggests that standard uniform or Gaussian-matching quantizers are suboptimal for highly correlated data common in neural networks.
Efficiency: The method offers a path to higher accuracy in LLM inference without increasing memory footprint, simply by re-designing the quantization codebooks based on the statistical dependence of the operands.

In summary, the paper establishes that optimal matrix multiplication quantization requires accounting for the correlation between operands, leading to non-standard, potentially bimodal quantization densities that significantly outperform existing industry standards.