WUSH: Near-Optimal Adaptive Transforms for LLM Quantization

AI summary

Overview

Research area: Machine learning systems — post-training quantization (PTQ) of large language models, specifically data-aware linear transforms applied to weights and activations before low-bit quantization.

Technical level: Advanced. The core contribution is a closed-form optimality derivation involving Cholesky decompositions, singular value decompositions, and trace inequalities, alongside systems-level kernel work.

Scope: The paper derives closed-form, data-dependent, non-orthogonal blockwise transforms (WUSH) that provably minimize quantization error under round-to-nearest (RTN) AbsMax block quantizers for both floating-point (FP) and integer (INT) formats, and demonstrates their accuracy and throughput benefits on LLM weight-activation quantization.

What This Paper Is About

Quantizing LLM weights and activations to 4 bits is a standard way to shrink and speed up inference, but a handful of extreme outlier values stretch the dynamic range and make low-bit errors much worse. Existing fixes rely on fixed, data-agnostic transforms such as Hadamard rotations, which work well empirically but had no clear optimality justification for quantization. This paper asks what the provably optimal linear blockwise transform looks like when the quantizer is specified, and shows it is a non-orthogonal, data-aware transform that keeps Hadamard as one ingredient.

Key Contributions

1. Closed-form optimal blockwise transforms. The authors derive an analytical solution for the linear transforms that minimize the blockwise output loss of joint weight-activation quantization under standard RTN AbsMax-scaled block quantizers, instead of learning transforms by gradient descent on calibration data as SpinQuant and FlatQuant do.

2. The WUSH construction. The optimal transform combines a normalized Hadamard backbone with a data-dependent second-moment component built from the Cholesky factors of the weight and activation second moments and the SVD of their product, yielding a generally non-orthogonal, data-aware transform. The paper notes that Hadamard is the only data-agnostic ingredient in the optimal formulation, which explains its prior empirical success.

3. Optimality proofs for FP and INT formats. Theorem 4.1 establishes that the optimal configuration is U' = H, S' = S^(1/2), R = I for floating-point types, and that the same configuration is optimal for integer types up to a d^(o(1)) factor for zero-mean multivariate Gaussian/Laplacian data and within a d factor for any distribution.

4. Practical integration and speed. WUSH is combined with both RTN and GPTQ (via an interleaved schedule that couples weight updates with transform construction) and is implemented as a fused GPU kernel whose throughput matches optimized blockwise Hadamard kernels.

Main Findings

• Accuracy gains over Hadamard baselines. On Llama-3.1-8B-Instruct in MXFP4, WUSH improves W4A4 accuracy by +2.8 average points with RTN and +0.7 average points with GPTQ relative to the strongest Hadamard-based baselines. The paper reports improvements of up to +2.8 average points (MXFP4 RTN) and +0.7 points (MXFP4 GPTQ) over the Hadamard-based baseline.

• Throughput. WUSH delivers up to 5.8× per-layer speedups over BF16 via FP4 matrix multiplication, while its kernel throughput matches that of optimized blockwise Hadamard kernels despite using a distinct data-aware transform per block.

• Format gap narrowing. The paper states that WUSH substantially narrows the gap between NVFP and MXFP formats, and that benefits are consistent when combined with GPTQ.

• Orthogonal transforms cannot help under the FP model. Under the paper's smooth FP error model, the objective evaluates to the same value for any orthogonal transform, including the identity and the Hadamard matrix; WUSH can reduce the error by at most d times in the extreme outlier scenario. Both non-trivial choices U' = H and S' = S^(1/2) are essential — setting either to identity yields the same suboptimal trace term as orthogonal transforms.

• Asymmetry in transform roles. The construction satisfies T_xvsh = T_wush^(-T), so the activation-side transform and weight-side transform are tied; the weight-side transform is absorbed into pre-quantized weights, leaving the activation-side block transform as the only online inference overhead.

• Offline cost is comparable to GPTQ. Per-block offline time and memory are O(d² d_batch) and O(d d_batch); per-layer costs are O(d d_in d_batch) time and O(d_in d_batch) memory. Standard GPTQ's Hessian computation for a layer requires O(d_in² d_batch) time and O(d_in d_batch) memory. When combined with GPTQ, WUSH adds only negligible overhead on top of GPTQ.

• Negligible storage overhead. With 4-bit weights (plus groupwise 8-bit scales) and 16-bit transforms, the relative storage overhead is 4d/d_out per layer, which the paper describes as negligible for typical LLM dimensions.

• Not reported. The provided content does not report calibration dataset sizes, the number of calibration tokens, or benchmark results beyond the Llama-3.1-8B-Instruct MXFP4 figures cited.

Methodology in Plain English

The authors frame quantization as an optimization problem: given a weight matrix W and calibration activations X, find transforms that make the quantized product match the unquantized product as closely as possible. They constrain the transforms to be block-diagonal, with block size matching the quantization group size, so that no block is larger than a power of two and each block can be handled independently.

The first approximation is that each block's loss can be minimized separately, which reduces a global problem to many small ones. The second is to model quantization error probabilistically: for floating-point formats, the error of a number is modeled as proportional to its own magnitude; for integer formats, the error is modeled as proportional to the maximum absolute value in the group. This lets them write the expected error as an expression in terms of the transform, then minimize it.

The minimization yields the answer in closed form. Each block's transform is built by taking the Cholesky factor of the weight second moment and of the activation second moment, computing the SVD of their product to get orthogonal matrices U and V and a singular-value matrix S, and then assembling the transform as a Hadamard matrix times S^(-1/2) times U^T times the transposed weight Cholesky factor. Damping is added to the second-moment diagonals when the blocks are rank-deficient. Everything is done on small d × d matrices.

For deployment, weights are pre-transformed and quantized offline (either with plain RTN or with GPTQ, where WUSH and GPTQ are interleaved because GPTQ's iterative weight updates change the second-moment statistics the transform depends on). At inference time, only the activation side needs a dynamic transform, which is fused with activation quantization in a GPU kernel.

Why This Matters

This is a rare case where a widely used heuristic — the Hadamard rotation — is shown to be one component of a provably optimal construction rather than merely a good empirical default. That reframes how transform-based quantization methods should be designed: instead of searching over rotations with gradient descent on calibration data, the paper shows an analytical solution exists under standard quantizer assumptions. It also unifies the treatment of FP and INT formats, which previously motivated different mitigation strategies.

Real-world applications:

• Serving LLMs at 4-bit precision. W4A4 quantization with WUSH can reduce memory footprint and raise per-layer throughput, which matters for hosting large models on constrained accelerator memory.
• On-device and edge inference. Lower-precision matrix multiplication is directly useful where power and memory bandwidth are the binding constraints.
• Datacenter cost reduction. The reported 5.8× per-layer speedup over BF16 via FP4 MatMul translates into throughput per accelerator, which maps to serving cost.
• Integration with existing PTQ pipelines. Because WUSH plugs into GPTQ with negligible overhead and reuses the same Hessian information, existing calibration pipelines can adopt it without new infrastructure.

Industry relevance: The authors include affiliations with Red Hat AI and ETH Zürich alongside ISTA, and the work targets the emerging FP4 block formats (MXFP, NVFP) that hardware vendors are actively shipping. A closed-form transform that fits into an existing GPTQ-style calibration pass and ships with a fused kernel is directly actionable for inference-stack engineers.

Future Directions

• Extending the optimality analysis to other quantizer families. The derivation assumes RTN with AbsMax scales and block-diagonal transforms; whether the same construction remains optimal for other scaling rules or adaptive rounding schemes is open.

• Reducing the INT gap. For integer types, the same configuration is optimal only up to a d^(o(1)) factor for Gaussian/Laplacian data and within a d factor for arbitrary distributions, leaving room to tighten the guarantee or design INT-specific transforms.

• Removing the interleaved GPTQ schedule. Because GPTQ's weight updates and WUSH's transform construction are coupled, the current method requires a specific computational schedule; decoupling or amortizing this could simplify deployment.

• Extending beyond the block-diagonal constraint for activations and further kernel optimization. The paper constrains transforms to be block-diagonal to keep online activation transformation cheap, and reports kernel throughput on par with optimized Hadamard kernels — whether larger blocks or different granularities pay off is an open question.

Target Audience

Researchers and engineers working on LLM quantization, efficient inference, and low-precision numerics — particularly those already familiar with GPTQ, QuaRot, SpinQuant, or FlatQuant and with the MXFP/NVFP format landscape. The theoretical sections assume comfort with linear algebra (SVD, Cholesky, trace inequalities) and probabilistic error modeling, so readers wanting only the practical takeaway should focus on the introduction, the algorithm listings, and the empirical results.

Authors’ abstract

Quantizing LLM weights and activations is a standard approach for efficient deployment, but a few extreme outliers can stretch the dynamic range and amplify low-bit quantization errors. Prior transform-based mitigations (e.g., Hadamard rotations) are fixed and data-agnostic, and their optimality for quantization has remained unclear. We derive closed-form optimal linear blockwise transforms for joint weight-activation quantization under standard RTN AbsMax-scaled block quantizers, covering both integer and floating-point formats. The resulting construction, WUSH, combines a Hadamard backbone with a data-dependent second-moment component to form a non-orthogonal transform that is provably near-optimal for FP and INT quantizers under mild assumptions while admitting an efficient fused GPU implementation. Empirically, WUSH improves W4A4 accuracy over the strongest Hadamard-based baselines (e.g., on Llama-3.1-8B-Instruct in MXFP4, it gains +2.8 average points with RTN and +0.7 with GPTQ) while delivering up to 5.8$\times$ per-layer throughput over BF16 via FP4 MatMul. Source code is available at https://github.com/IST-DASLab/WUSH.