Contraction-Gauge Preconditioning for Quantized Matrix Multiplication

Abstract: We study low-precision computation of C=AB with both factors quantized. We derive an exact finite-dimensional identity for the expected squared product error under independent, zero-mean entrywise errors with known variance fields; it holds exactly for non-overloading subtractive dither and for independent stochastic rounding, and we empirically assess deterministic round-to-nearest (RTN). Using the product-preserving equivalence AB=(AT)(T^{-1}B), we formulate contraction-gauge preconditioning: jointly choosing a factor representation and its sharing pattern before quantization. Preconditioning can reduce product error but may require extra transformed, quantized copies of the opposite operand: a shared transform needs one copy, a block-specific transform up to one per block. Within the bounded family of positive diagonal gauges (folds), a geometric program computes a globally optimal shared fold and a linear program decides whether the identity fold is already optimal. For other families we derive computable selection statistics -- tail index for scaling, profile spread for partitioning, coherence and weighted-Gram energy for rotations, slice-energy covariance for hierarchy depth -- with upper bounds for ranking heuristic candidates. Across twelve linear products from a trained three-block image classifier, median within-product rank correlations between dither-model predictions and deterministic-RTN errors are 0.937 at 8 bits and 0.918 at 4 bits. The GP fold cuts held-out product error over the identity fold by 18.0% (8-bit) and 20.5% (4-bit) in geometric mean, beats a SmoothQuant-style grid baseline at both precisions and on ten of twelve products, and lowers composed logit MSE by 15.4% and 26.4%. We thus provide exact stochastic product-error accounting, certified selection within the diagonal family, and a common objective for evaluating reusable transform candidates under RTN.

• The paper develops an exact expected product-error identity and uses it to optimize diagonal folds, rotations, partitions, and hierarchical gauges before quantization.
• Experiments show that geometric-program folds, profile-aware clustering, and targeted rotations can substantially reduce error, including a 59.82× rotation gain on spiky synthetic data.
• On twelve classifier products, the approach reduced held-out RTN error by 18.0% at 8 bits and 20.5% at 4 bits, while lowering composed logit MSE by 15.4% and 26.4%.

Overview

This paper develops a unified framework for selecting and sharing transforms in low-precision matrix multiplication $C = AB$ where both factors are quantized. The central object is a contraction-gauge equivalence: for any invertible $T \in \mathrm{GL}(K)$ acting on the shared summation (contraction) dimension, $AB = (AT)(T^{-1}B)$, so transform selection is a choice of representative within this equivalence class made before quantization. The authors—Sao, Miniskar, Valero-Lara, Teranishi, and Seal of Oak Ridge National Laboratory—couple this classical preconditioning identity to an exact finite-dimensional expected product-error identity, derive optimization and selection rules for several structural families of gauges (diagonal folds, orthogonal rotations, partial Householder rotations, block-diagonal hierarchies), and validate the resulting criteria on controlled synthetic data and on twelve internal products of a trained image classifier under deterministic round-to-nearest (RTN) quantization.

Two contributions organize the work. First, an exact stochastic accounting of expected squared product error under mutually independent, zero-mean entrywise errors with known variance fields, which holds verbatim for non-overloading subtractive dither and independent stochastic rounding, and serves as a high-resolution approximation for deterministic RTN. Second, a representation-reuse descriptor, the opposite-factor quantized-copy count $n_{\mathrm{opp}}$, that prices the storage cost of using distinct gauges across output blocks: a shared gauge requires one transformed-and-quantized copy of the opposite factor, whereas block-specific gauges can require one per block.

The product-error identity

Under the noise model $\hat A = A + E_A$, $\hat B = B + E_B$ with mutually independent zero-mean entries and variance fields $v^A_{ik}$, $v^B_{kj}$, the paper proves

$$\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 \Big(\sum_i v^A_{ik}\Big)\Big(\sum_j v^B_{kj}\Big).$$

The first two leading terms weight each entry's error by the opposite factor's energy through which it propagates; the third is a bilinear simultaneous-error term that is $O(c^2)$ at fine precision but material at coarse precision or large contraction dimension. This is the exact finite-dimensional counterpart of one-sided objectives such as GPTQ's $T \in \mathrm{GL}(K)$0, treating both operands and their cross term. A weighted-norm corollary replaces raw energies with $T \in \mathrm{GL}(K)$1- and $T \in \mathrm{GL}(K)$2-weighted sensitivities, connecting the framework to Hessian-weighted quantization and iterative solvers. The design chain is then explicit: the gauge $T \in \mathrm{GL}(K)$3 determines factor representations, grouping determines per-group ranges $T \in \mathrm{GL}(K)$4, the dither law $T \in \mathrm{GL}(K)$5 converts ranges to variances, and the identity scores the result. A Hanson–Wright-based appendix extends the mean to a conservative single-run certificate under variance-normalized sub-Gaussian errors, though the constant $T \in \mathrm{GL}(K)$6 is left unspecified and calibration requires empirical work.

The model's validity is exact only under non-overloading subtractive dither or independent stochastic rounding with residue-dependent variances. For deterministic RTN the paper treats the identity as an approximation and measures its transfer empirically; a truncated characteristic-function diagnostic $T \in \mathrm{GL}(K)$7 flags lattice-aligned residues where the approximation degrades. A notable remark shows that transform ordering can be quantizer-dependent: on a constructed four-dimensional pair, dither ranks a Hadamard rotation first while deterministic INT8 RTN ranks the identity first—a concrete caveat against transferring stochastic rankings to RTN deployments without calibration data.

Diagonal folds as a geometric program

For a positive diagonal contraction gauge ("fold") $T \in \mathrm{GL}(K)$8 shared over a rectangular output domain $T \in \mathrm{GL}(K)$9, the leading-error objective with epigraph range variables $AB = (AT)(T^{-1}B)$0 is shown to be a posynomial subject to monomial constraints—hence a geometric program whose log-domain reformulation is convex, so every solution is globally optimal. Finite box bounds on $AB = (AT)(T^{-1}B)$1 guarantee attainment; without compactification, disjoint-support degenerate inputs admit no finite infimum. The program exhibits a one-dimensional scale-gauge non-identifiability $AB = (AT)(T^{-1}B)$2, fixed by a normalization such as $AB = (AT)(T^{-1}B)$3. Convexity also yields an exact linear-program test deciding whether the identity fold is already optimal via directional derivatives at $AB = (AT)(T^{-1}B)$4.

This certified joint optimization distinguishes the approach from SmoothQuant-style grids, AWQ searches, and gradient-based learned scales. A two-channel counterexample demonstrates that coordinatewise heuristics—the range rule $AB = (AT)(T^{-1}B)$5 and the norm rule—can be strictly suboptimal (a 4.22% leading relative gap in the displayed instance), because a shared fold couples every transformed row range to every transformed column range. Complementing the shared-fold optimum, a row-local fold benchmark gives the energy-matched infimum $AB = (AT)(T^{-1}B)$6 per row, attained when the fold follows coordinate range rather than energy; the spread statistic $AB = (AT)(T^{-1}B)$7 isolates the loss from sharing one fold across a block.

Partitioning theory

Choosing which rows share a fold is treated as partitioning magnitude profiles. The paper proves that scalar-norm sorting—sorting by a single per-row norm and cutting contiguous blocks—can incur an expected $AB = (AT)(T^{-1}B)$8 loss relative to the profile-pure optimum even under uniformly random tie-breaking, because equal-norm rows may have disjoint supports. In contrast, clustering in log-magnitude coordinates with regularizer $AB = (AT)(T^{-1}B)$9 yields blocks whose regularized spread satisfies $n_{\mathrm{opp}}$0 whenever the block diameter is at most $n_{\mathrm{opp}}$1 in the $n_{\mathrm{opp}}$2 log-metric; Gonzalez farthest-point traversal gives a constant-factor guarantee, and common support plus a positive lower bound transfers control to the unregularized spread. For rank-one profiles, sorting makes the optimum contiguous and an interval dynamic program finds the exact $n_{\mathrm{opp}}$3-block solution in $n_{\mathrm{opp}}$4. Latent-dimension and tropical (max-plus template) refinements extend the guarantee to low-dimensional profile structure, with the approximation residual entering the bound explicitly.

Rotations, reflectors, and hierarchies

For orthogonal gauges, block-local coherence factors $n_{\mathrm{opp}}$5 measure how much a rotation flattens squared quantizer ranges. Haar-random or Rademacher-signed Hadamard transforms control all block vectors simultaneously within a logarithmic factor $n_{\mathrm{opp}}$6 with probability $n_{\mathrm{opp}}$7, so a spiky baseline gains $n_{\mathrm{opp}}$8 while an already flat profile offers nothing—an important boundary condition on when rotation helps. A conditional substitution analysis explains when rotation removes structure a subsequent fold would exploit: majorization arguments apply only under uniform opposite-factor weights, and random rotations raise the expected coherence of flat profiles from $n_{\mathrm{opp}}$9 to $\hat A = A + E_A$0, so candidate chains must be ranked by the full product-error objective.

When the weighted Gram matrix $\hat A = A + E_A$1 has a low-rank spectral head, a construction using at most $\hat A = A + E_A$2 Householder reflectors maps the top-$\hat A = A + E_A$3 eigenspace to an incoherent target frame, yielding the bound $\hat A = A + E_A$4: the head is discounted logarithmically while the tail retains full weight. A deterministic Hadamard target trades the log discount for a worst-case $\hat A = A + E_A$5 factor. Notably, the top-$\hat A = A + E_A$6 eigenspace minimizes this bound only when the head coefficient is below one; above that threshold the estimate supplies no preference, and a matching spectral lower bound on the reflector budget remains open.

For hierarchical (block-diagonal) gauges, the single-level decision reduces exactly to slice-energy covariance: the ratio of hierarchical to flat leading-error surrogates equals $\hat A = A + E_A$7, which is below one precisely when the centered inner product of the two factors' normalized slice-energy profiles is negative. A telescoping identity accumulates node increments across levels, each increment being negative exactly when child profiles anti-correlate at that scale; since increments can change sign across scales, optimal stopping is bottom-up rather than greedy. Slice design itself reduces to balanced graph partitioning with weights $\hat A = A + E_A$8, which the paper proves weakly NP-complete even for two slices with $\hat A = A + E_A$9 (via reduction from Partition), though a pseudo-polynomial dynamic program exists for that restricted case.

Quantizer design after the transform

With the transform fixed, three quantizer choices receive closed-form treatment. Under a fixed bit-width sum $\hat B = B + E_B$0 and the high-rate law $\hat B = B + E_B$1, the cross term is constant and the interior optimum obeys $\hat B = B + E_B$2, allocating more bits to the operand with larger product-weighted error coefficient; groupwise weighted-storage constraints instead yield coupled water-filling equations. Clipping admits an exact identity separating deterministic overload bias (weighted by the unclipped opposite factor) from granular error (weighted by the clipped factor), with a convex diagonal surrogate whose optimal threshold is the unique root of a balancing equation. Companding analysis derives the product-weighted optimal point density $\hat B = B + E_B$3, with the caveat that nonlinear expansion generally destroys the zero-mean independence property required by the exact identity, so compressed-domain dither introduces bias.

Experimental validation

Controlled synthetic experiments isolate each mechanism. Per-vector scaling gain grows polynomially as $\hat B = B + E_B$4 under Pareto tail index $\hat B = B + E_B$5—measured log-log slopes of $\hat B = B + E_B$6 for $\hat B = B + E_B$7 against predicted $\hat B = B + E_B$8—and only logarithmically for Gaussian factors. On constructed instances, the domain-shared fold GP achieves a modeled objective $\hat B = B + E_B$9 lower than the identity fold, and asymmetric bit allocation predicts a 1.7-bit gap that the measured 9/7 split confirms, beating the symmetric 8/8 split by $v^A_{ik}$0 in squared relative Frobenius error. Log-magnitude clustering attains $v^A_{ik}$1 the profile-pure reference versus $v^A_{ik}$2 for scalar-norm sorting. Rotation gain reaches $v^A_{ik}$3 on injected coordinate outliers while remaining statistically indistinguishable from zero ($v^A_{ik}$4) on the flat Gaussian control; three Householder reflectors match a full rotation ($v^A_{ik}$5 ratio) on rank-3 outlier data. The hierarchy-to-flat measured crossover occurs at criterion value $v^A_{ik}$6, closely tracking the surrogate crossing. The reuse trade-off is priced directly: on heavy-tailed operands, a four-copy block-GP design lowers error by $v^A_{ik}$7 relative to a one-copy shared rotation, while four independent block rotations show no resolved gain for the same copy count.

The trained-classifier study evaluates thirteen candidates (identity, eleven SmoothQuant-style $v^A_{ik}$8 values, and the numerically fitted GP fold) across twelve products of a ViT-like digit classifier at 8 and 4 bits, with symmetric RTN, per-token activation scales, and no clipping. Median within-product Spearman correlations between dither-model predictions and held-out deterministic-RTN errors are 0.937 at 8 bits and 0.918 at 4 bits. Relative to the identity fold, the GP fold reduces held-out product error by 18.0% at 8 bits and 20.5% at 4 bits in geometric mean, improving ten of twelve products at each precision and beating even a test-selected $v^A_{ik}$9 oracle in geometric mean (worst GP-to-oracle ratios $v^B_{kj}$0 and $v^B_{kj}$1). Calibration-set dither prediction selects the held-out winner on ten of twelve products with geometric-mean regret of $v^B_{kj}$2 and $v^B_{kj}$3. Quantizing all twelve products simultaneously reduces composed logit MSE by 15.4% and 26.4% relative to the identity-fold model. The simultaneous-error term contributes at most $v^B_{kj}$4 of predicted GP error on average, supporting the leading-objective fit at these precisions.

Limitations and open questions

The paper is explicit about scope and assumptions. The exact identities hold only under non-overloading subtractive dither or independent stochastic rounding; deterministic RTN results rest on empirical transfer, calibrated here on one compact vision model with 128 calibration images, and the authors state plainly that applying the protocol to pretrained LLMs and real token activations remains to be done. The coherence, partial-rotation, and hierarchy results are upper-bound surrogates rather than realized-error guarantees, and the clipping threshold rule uses a diagonal surrogate that drops cross-coordinate overload correlations. Several problems are left open: exact max-range optimization over the full $v^B_{kj}$5 family with condition-number constraints, jointly convex fold–threshold selection, a matching spectral lower bound on the Householder budget, exact dynamic programs for tropical-profile partitioning, calibration-aware RTN theory, coupled finite-rate interfaces, and kernel measurements translating the copy-count descriptor into runtime and memory traffic. The repository housing code and artifacts was private at writing time pending release clearance.

Conclusion

The paper organizes quantized matrix-product design around a single exact objective and an explicit sharing choice. Its certified core—the product-error identity, the bounded domain-shared fold geometric program, the identity-fold optimality test, and the telescoping hierarchy criterion—is exact under stated stochastic assumptions, while five computable statistics (tail index, profile spread, block coherence, weighted-Gram spectrum, slice-energy covariance) screen heuristic candidates across the scaling, grouping, rotation, partial-rotation, and hierarchy families. The trained-classifier results demonstrate that dither-derived rankings transfer to deterministic RTN with high fidelity and consistent error reductions, establishing a reusable evaluation workflow: score stochastic models exactly, then validate RTN candidates on representative calibration data.