🤖 machine learning

Contraction-Gauge Preconditioning for Quantized Matrix Multiplication

This paper introduces contraction-gauge preconditioning, a method that optimizes factor representations and sharing patterns before quantization to minimize product error in low-precision matrix multiplication, supported by an exact stochastic error model and validated by significant error reductions across various precisions and benchmarks.

Original authors: Piyush Sao, Narasinga Miniskar, Pedro Valero-Lara, Keita Teranishi, Sudip Seal

Published 2026-07-22

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Piyush Sao, Narasinga Miniskar, Pedro Valero-Lara, Keita Teranishi, Sudip Seal

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 trying to bake a giant, perfect cake, but you only have a tiny, crumbly measuring cup. In the world of artificial intelligence, the "cake" is a massive calculation called matrix multiplication, where two huge grids of numbers are smashed together to create a new result. The "measuring cup" is the computer's memory, which often forces these numbers to be rounded off to fit into smaller, simpler boxes (a process called quantization). If you round too aggressively, the cake collapses; the AI gets confused, makes mistakes, or stops working entirely.

For a long time, scientists have tried to fix this by just being more careful with how they round, or by grouping numbers together. But there's a sneaky problem: the numbers aren't all the same. Some are huge and loud (outliers), while others are tiny and quiet. If you treat them all the same, the loud ones drown out the quiet ones, or the rounding errors pile up in a way that ruins the final taste. The big question has been: How do we rearrange these numbers before we squeeze them into the small boxes, so that when they are smashed together later, the result is still deliciously accurate?

This paper, written by a team of researchers at Oak Ridge National Laboratory, acts like a master chef's new recipe book for rearranging those ingredients. They discovered a precise mathematical rule—a "product-error identity"—that tells them exactly how much error will happen when they mix two rounded grids together. Think of it as a crystal ball that predicts exactly how much the cake will crumble based on how the ingredients were chopped up.

Using this crystal ball, they developed a method called "Contraction-Gauge Preconditioning." Imagine you have a stack of unevenly sized bricks (the numbers in your first grid) and a stack of unevenly sized mortar bags (the numbers in your second grid). If you try to build a wall with them as they are, the wall will be wobbly. The researchers found a way to stretch and shrink the bricks and mortar bags before you start building, so they fit together perfectly, even if you have to use a slightly rougher mortar later. They call this stretching and shrinking a "fold."

The paper proves that finding the perfect way to fold these numbers is like solving a specific type of math puzzle called a "geometric program." They showed that you can solve this puzzle to find the absolute best way to rearrange the numbers for a specific job, rather than just guessing. They also figured out that sometimes, to get the best result, you might need to make a few extra copies of your ingredients (the "opposite factor") to use different folding tricks in different parts of the wall. They created a way to count these extra copies so engineers know exactly how much extra memory they will need.

The team tested their new recipe in two ways. First, they built fake, controlled scenarios in a computer lab to see if their math predictions matched reality. They found that their "crystal ball" was incredibly accurate, predicting the errors with over 90% correlation. Second, they took a real, trained AI that recognizes handwritten digits (like the ones you might write on a tablet) and tried to shrink its brain down to use less memory. When they used their new "GP fold" method, the AI made significantly fewer mistakes than when they used the old, standard ways of rounding. Specifically, at 8-bit precision, the error dropped by 18%, and at 4-bit precision (which is very tiny), the error dropped by over 20%.

The paper also rules out some common shortcuts. They showed that simply sorting numbers by their size (like lining up bricks from biggest to smallest) isn't always the best strategy; sometimes, you need to look at the shape of the numbers to group them correctly. They also found that while spinning the numbers around (rotating them) can help, it doesn't always work if the numbers are already flat and even. Their method is a smart, adaptable system that checks the specific ingredients of your problem and chooses the best rearrangement strategy, whether that's a simple stretch, a complex spin, or a mix of both.

In short, this paper doesn't just say "rounding is hard." It gives a precise, mathematically proven toolkit for engineers to rearrange their data so that even when they are forced to use low-precision, low-memory tools, the final result remains sharp, accurate, and reliable. It turns the messy art of guessing how to shrink AI models into a precise science of measuring and optimizing.

Technical Summary: Contraction-Gauge Preconditioning for Quantized Matrix Multiplication

Problem Statement

The paper addresses the challenge of low-precision matrix multiplication ($C = AB$) where both factors $A \in \mathbb{R}^{m \times K}$ and $B \in \mathbb{R}^{K \times n}$ are quantized. A primary difficulty in this domain is that quantization loss does not propagate uniformly; an error in a specific entry $A_{ik}$ affects the output row $C_{i,:}$ in proportion to the energy of the corresponding row in $B$ ( $B_{k,:}$ ), and vice versa. Existing methods, such as channel grouping or orthogonal transforms, often lack a unified, product-level objective that accounts for these asymmetric error propagations and the cost of storing multiple transformed copies of operands. The authors aim to derive exact error identities under stochastic assumptions and formulate a "contraction-gauge" framework to optimize the representation of factors before quantization while explicitly managing the trade-off between error reduction and the storage/traffic cost of transformed copies.

Methodology

1. Theoretical Foundation: Product-Error Identity
The authors derive an exact finite-dimensional identity for the expected squared product error under mutually independent, zero-mean entrywise errors with known variance fields. For quantized factors $\hat{A} = A + E_A$ and $\hat{B} = B + E_B$ , the expected squared Frobenius norm of the error is:
$\mathbb{E}\|\hat{A}\hat{B} - AB\|_F^2 = \sum_{i,k} v^A_{ik} \|B_{k,:}\|_2^2 + \sum_{k,j} v^B_{kj} \|A_{:,k}\|_2^2 + \sum_{k} \left(\sum_i v^A_{ik}\right)\left(\sum_j v^B_{kj}\right)$
This identity is exact for non-overloading subtractive dither and independent stochastic rounding with residue-dependent variances. It serves as the primary objective function, weighting errors by the energy of the opposite factor.

2. Contraction-Gauge Preconditioning
The method leverages the algebraic equivalence $AB = (AT)(T^{-1}B)$ for any invertible $T \in GL(K)$ . The authors define a "contraction gauge" as the choice of $T$ applied to the shared contraction dimension $K$ before quantization.

Gauge Domains: A domain is a fixed output block (e.g., a row block) assigned a single gauge.
Reuse Descriptor ( $n_{opp}$ ): The cost of preconditioning is measured by the number of distinct transformed, quantized copies of the opposite factor required. If a single gauge is shared across all domains, $n_{opp}=1$ ; if distinct gauges are used per block, $n_{opp}$ increases.

3. Structural Families and Optimization Rules
The paper proposes specific optimization strategies for different families of gauges:

Diagonal Folds (Scaling): For a shared domain, finding the optimal diagonal scaling $D$ such that $(AD, D^{-1}B)$ minimizes the error is formulated as a Geometric Program (GP). This allows for globally optimal selection of shared folds. A linear program (LP) is derived to test if the identity fold is already optimal.
Partitioning: To balance accuracy and copy count, the paper argues against scalar-norm sorting (which can lose a linear factor in performance due to random ties) and proposes clustering in log-magnitude coordinates to group rows with similar profiles.
Rotations: For orthogonal gauges, the paper uses block-local coherence to bound potential gains. It shows that random rotations (e.g., Hadamard) can reduce coordinate maxima within a logarithmic factor.
Partial Rotations (Householder): When weighted energy is low-rank, the authors propose using a small number of Householder reflectors ( $t$ ) to align the top eigenspace of a weighted Gram matrix to a flat target, costing $O(tK)$.
Hierarchies: For block-diagonal orthogonal gauges, a slice-energy anti-correlation criterion is derived. If the energy profiles of $A$ and $B$ across slices are anti-correlated, a hierarchical refinement reduces error; otherwise, a flat rotation is preferred. A telescoping identity allows for bottom-up depth selection.

4. Quantizer Design

Bit Allocation: Under a fixed bit-width sum constraint ( $b_A + b_B = b_{sum}$ ), the optimal split is asymmetric: $b_A - b_B = \frac{1}{2} \log_2(P_A/P_B)$ , where $P_A$ and $P_B$ are product-weighted error coefficients.
Clipping and Diagnostics: The paper introduces a truncated characteristic-function diagnostic ( $\Xi$ ) to detect lattice alignment in deterministic Round-to-Nearest (RTN) rounding, which can violate the independence assumptions of the noise model.

Key Contributions

Exact Error Accounting: Derivation of an exact expected squared product-error identity for quantized matrix multiplication under dither and stochastic rounding, explicitly separating leading error terms from bilinear cross-terms.
Contraction-Gauge Framework: A unified formulation for preconditioning that couples the choice of transform $T$ , its sharing pattern (gauge domains), and the resulting representation reuse cost ( $n_{opp}$ ).
Globally Optimal Fold Selection: Proof that the bounded shared-fold problem is a Geometric Program, solvable to global optimality, and the derivation of an exact LP test for identity-fold optimality.
Decision Statistics: Development of computable statistics to guide transform selection without computing exact optima for all families:
- Tail index for scaling.
- Profile spread for partitioning.
- Block coherence and weighted-Gram spectrum for rotations.
- Slice-energy covariance for hierarchy depth.
Reuse-Aware Design: Introduction of $n_{opp}$ as a metric to quantify the storage and traffic cost of different gauge sharing patterns, enabling a trade-off analysis between error reduction and implementation cost.

Results

Controlled Synthetic Experiments:

Scaling: Log-magnitude clustering for partitioning achieved near-optimal spread (1.015x profile-pure reference) compared to scalar-norm sorting (2.116x).
Folds: The domain-shared GP fold reduced modeled error by a factor of ~161x compared to the identity fold and 4.86x compared to separately optimized one-sided bounds.
Bit Allocation: Asymmetric bit splits (e.g., 9/7 instead of 8/8) achieved 1.7x lower error than symmetric splits, matching theoretical predictions.
Rotations: Randomized Hadamard rotations provided up to 60x error reduction on heavy-tailed data with injected outliers, while Householder reflectors with $t=3$ matched full rotation performance on rank-3 data.
Hierarchies: The anti-correlation criterion successfully predicted the crossover point where hierarchical refinement outperforms flat rotation.

Trained Classifier Validation:
Using a three-block image classifier (ViT-like) on handwritten digit images:

Correlation: The dither-model predictions showed high median Spearman correlation with deterministic RTN errors (0.937 at 8 bits, 0.918 at 4 bits).
Performance: Using the GP fold instead of the identity fold reduced held-out product error by 18.0% at 8 bits and 20.5% at 4 bits (geometric mean).
Comparison: The GP fold outperformed a SmoothQuant-style grid search (selected via calibration) on 10 of 12 products and achieved lower error than a test-selected oracle in geometric mean.
End-to-End: Quantizing all 12 products together with the GP fold reduced logit MSE by 15.4% (8-bit) and 26.4% (4-bit) relative to the identity-fold model.

Significance and Claims

The paper claims to provide a rigorous, stochastic foundation for quantized matrix multiplication that moves beyond heuristic scaling or ad-hoc grouping. By establishing an exact error identity and linking it to a "contraction-gauge" framework, the authors offer:

Certified Selection: Globally optimal selection within the bounded diagonal family via Geometric Programming.
Unified Objective: A common objective (product-weighted error) to evaluate diverse transform families (folds, rotations, hierarchies) and their sharing patterns.
Practical Trade-offs: A mechanism to explicitly balance error reduction against the cost of storing multiple transformed copies ( $n_{opp}$ ), which is critical for deployment.

The authors emphasize that while the exact identities hold for dither and stochastic rounding, the derived statistics and selection rules transfer effectively to deterministic Round-to-Nearest quantization, as validated by the trained classifier experiments. The work does not claim to solve the full $GL(K)$ optimization problem but provides a structured, executable procedure for selecting high-performing candidates within restricted, practical families.