PaperlayerView exact revision on arXiv

Source-linked AI summary

ShamAN-Q: Shampoo Augmented NanoQuant for Sub-1-bit LLM Weights

Jonathan Mei, Sang Hyub Kim, Oliver Knitter, Chi Chen, Martin Roetteler

arXiv:2609.38521v1cs.LGcs.AIstat.ML

TL;DR

Sub-1-bit quantization must account for reconstruction errors across many weights, while full Fisher curvature is too large to use directly. ShamAN-Q uses a tractable Kronecker curvature metric and redistributes rank at the same bit budget; on Qwen3-Base, it lowers WikiText-2 perplexity while matching or improving zero-shot accuracy. The Fisher is local to the weights, and one-bit reconstruction error lies outside its accurate-curvature neighborhood.

  • Problem

    NanoQuant’s diagonal reconstruction geometry does not capture relationships between channels, while the full mn×mn Fisher matrix is prohibitive for an m×n weight.

    Source: §1 · p. 2

  • Method

    ShamAN-Q fits a dense Kronecker curvature metric, uses it for reconstruction and rank allocation, and retains NanoQuant’s discrete projection and packed format.

    Source: §1 · p. 2

  • Results

    At ≲1-bit budgets, WikiText-2 perplexity falls from 27.56 to 22.96 (0.6B), 19.21 to 16.72 (1.7B), and 14.29 to 13.80 (4B), while zero-shot accuracy matches or improves.

    Source: §1 · p. 2

  • Takeaways & Limitations

    ShamAN-Q preserves NanoQuant’s packed format while using a dense metric and redistributed ranks at the same total number of bits.

    Source: §1 · p. 2, §6 · p. 9

    One-bit reconstruction error is about 0.4–0.5 of the weight norm, outside the neighborhood where the Fisher is an accurate curvature, so ShamAN-Q tempers the transported statistics before ADMM.

    Source: §4.3 · p. 7

Abstract

from arXiv · show

We introduce ShamAN-Q, a sub-1-bit post-training quantization method that extends NanoQuant by replacing each its diagonal reconstruction geometry with a tractable dense curvature metric, using a general paradigm popularized by the Shampoo optimizer. For each linear weight, ShamAN-Q fits a Kronecker product to the empirical Fisher information matrix of a small calibration set by Kullback--Leibler minimization, forming a Mahalanobis reconstruction loss from the result. The continuous ADMM updates from NanoQuant become solutions to Sylvester equations, while its discrete projection and deployment format remain unchanged. Because the curvature is local to a given set of weights, ShamAN-Q re-measures the input curvature statistic for each layer immediately before layer factorization, periodically refreshing all statistics on the partially quantized model. ShamAN-Q also redistributes the uniform rank from NanoQuant across layers at the same total number of bits. On Qwen3-Base, ShamAN-Q lowers WikiText-2 perplexity at $\approx$1 bpw from 27.56 to 22.96 (0.6B), 19.21 to 16.72 (1.7B), and 14.29 to 13.80 (4B) while matching or improving zero-shot accuracy on the Eleuther LM Evaluation Harness. On 0.6B, ShamAN-Q at $\approx$0.8 bpw matches the published perplexity of NanoQuant at $\approx$1.0 bpw.

1 INTRODUCTION

ShamAN-Q targets sub-1-bit quantization, where reconstruction error affects many weights and channel correlations matter. It replaces NanoQuant’s diagonal geometry with dense curvature statistics and redistributes rank at the same bit budget, improving reported Qwen3-Base results.

1 INTRODUCTION

  • Decoding is often bandwidth-bound because each generated token requires transferring the active parameters; at batch size B = 1, matrix–vector multiplication moves O(n^2) data for O(n^2) computation.The paper motivates shrinking stored weights as a way to address this decoding cost.Source: §1 · pp. 1–2
  • NanoQuant’s diagonal preconditioning captures channel-wise second-order information but not relationships between channels, while the full mn×mn Fisher matrix is intractable for an m×n weight.ShamAN-Q uses input- and output-side statistics to approximate that information tractably.Source: §1 · p. 2
  • At ≲1-bit budgets, WikiText-2 perplexity falls from 27.56 to 22.96 (0.6B), 19.21 to 16.72 (1.7B), and 14.29 to 13.80 (4B), while zero-shot accuracy matches or improves.The reported comparison is against NanoQuant at the same bit budget; the contribution summary reports results from 0.6B to 14B.Source: §1 · p. 2
  • ShamAN-Q fits a dense two-sided reconstruction metric with KL-Shampoo curvature, changes continuous ADMM updates to Sylvester solves, and keeps NanoQuant’s projection and packed format unchanged.It also redistributes NanoQuant’s uniform rank budget across layers at bit parity.Source: §1 · p. 2

2 RELATED WORK

Prior sub-1-bit approaches include training with quantization-aware objectives, partial binarization, and PTQ through low-rank binary factors. Related curvature methods approximate layer sensitivity with input statistics or Kronecker-structured factors.

2 RELATED WORK

  • BitNet and OneBit use quantization-aware objectives, PB-LLM keeps salient weights at higher precision, and NanoQuant factorizes each weight matrix into two low-rank binary matrices and two scale vectors.NanoQuant’s factorization breaks the one-bit-per-weight lower bound for in-place binarization with full-precision scales.Source: §2 · pp. 2–3
  • GPTQ weights reconstruction by input second moments, while K-FAC and Shampoo approximate layer Fisher information with Kronecker products of input and output statistics.NKP fits a Kronecker product in Frobenius norm, whereas KL-Shampoo uses a matrix-normal maximum-likelihood fit; ShamAN-Q adopts the latter for its token weighting.Source: §2 · p. 3
  • NanoQuant’s LB-ADMM uses SVID-projected consensus variables to separate low-rank reconstruction from binary structure; ShamAN-Q retains that decoupling and changes only the continuous subproblems to two-sided Sylvester equations.Both methods retain Euclidean consensus penalties.Source: §2 · p. 3

3 BACKGROUND

NanoQuant represents linear layers with low-rank binary factors and scale vectors, then optimizes them through a blockwise pipeline. Kronecker curvature statistics approximate the layer Fisher with tractable input- and output-side matrices, with several fitting and preconditioning choices.

3.1 NANOQUANT

  • A rank-r NanoQuant layer uses two binary factors and FP16 scale vectors, costing (r + 16)(m + n) bits and (r + 16)(m + n)/mn bpw.The signs require r(m + n) bits, and the scales require 16(m + n) bits.Source: §3.1 · p. 3
  • NanoQuant calibrates a full-precision teacher, reconstructs decoder blocks in order, tunes latent factors and scales, then distills only scale vectors against the teacher.During block reconstruction, full-precision weights are tuned to reproduce teacher outputs from the already-quantized prefix’s activations.Source: §3.1 · p. 3
  • NanoQuant shrinks clipped input-activation and output-gradient second moments into diagonal preconditioners, then minimizes a weighted Frobenius reconstruction loss.Its LB-ADMM separates continuous factors from SVID-constrained consensus variables; the U update is a right-sided r × r normal equation.Source: §3.1 · pp. 3–4

3.2 KRONECKER CURVATURE STATISTICS: SHAMPOO AND KL-SHAMPOO

  • The empirical layer Fisher is approximated as A ⊗ G, where A and G act on input and output coordinates; the resulting error metric is vec(E)⊤(A ⊗ G) vec(E).The statistics can be applied directly without forming the full mn × mn Kronecker product.Source: §3.2 · p. 4
  • Shampoo accumulates two matrix statistics to precondition gradients; its approximation’s square takes one power-iteration step toward the optimal Kronecker approximation of the empirical Fisher.The passage relates this approximation to using exponent 1/4 per statistic as an inverse square root of approximate curvature.Source: §3.2 · p. 4
  • NKP fits the Kronecker product in Frobenius norm, whereas KL-Shampoo fits it by matrix-normal maximum likelihood with token weights given by inverse quadratic forms.The KL-Shampoo fit uses a trace constraint to remove scale ambiguity between A and G.Source: §3.2 · p. 4

4 SHAMAN-Q

ShamAN-Q retains NanoQuant’s representation but replaces its diagonal reconstruction geometry with dense curvature at multiple stages of fitting. It combines curvature-aware reconstruction with rank redistribution at the same total bit budget.

4 SHAMAN-Q

  • ShamAN-Q keeps NanoQuant’s representation and changes low-rank fitting geometry through calibration, rank allocation, continuous updates, and curvature refresh.SVID, scale extraction, latent tuning, packing, and scale distillation remain unchanged.Source: §4 · p. 5, §4.3 · p. 8

4.1 DENSE KRONECKER CURVATURE FROM GLOBAL CALIBRATION

  • For each linear layer, three alternating KL-Shampoo passes fit a Kronecker curvature approximation from calibration data, avoiding materialization of the empirical Fisher matrix.KL weighting limits the influence of massive-activation tokens compared with NKP weighting, whose activation terms are quartic in magnitude.Source: §4.1 · p. 6
  • Shrinkage changes anisotropy while preserving each statistic’s mean eigenvalue, with no additional damping term.The diagonal of the shrunken statistic is NanoQuant’s shrunken d2, and robust token clipping is retained.Source: §4.1 · p. 6

4.2 NON-UNIFORM RANK ALLOCATION AT BIT PARITY

  • At bit parity, ShamAN-Q redistributes NanoQuant’s uniform rank budget across layers using depth and layer-type priors or measured reconstruction-loss sensitivity.The measured allocator probes layers at three ranks, fits a power law, and distributes 32-rank increments by predicted loss decrease per bit; the depth prior compensates for its tendency to undervalue late blocks.Source: §4.2 · p. 6

4.3 BLOCK RECONSTRUCTION PIPELINE

  • Before factorizing each layer, ShamAN-Q remeasures its input statistic and periodically refreshes statistics for unfactorized layers on the partially quantized model.The refreshed statistics are shrunk, then transported by congruence to preserve the Mahalanobis error; diagonal preconditioners are floored to remain invertible.Source: §4.3 · pp. 6–7
  • Because one-bit reconstruction errors are large relative to the Fisher’s accurate-curvature neighborhood, ShamAN-Q applies trace-preserving spectral tempering before ADMM.Tempering halves the log-condition number and reduces trust in dominant curvature directions; the square root is described as regret-optimal under uncertain curvature.Source: §4.3 · p. 7
  • Dense-curvature LB-ADMM replaces the Frobenius reconstruction term with a transported Mahalanobis metric while leaving consensus penalties Euclidean.Its continuous U/V updates solve Mahalanobis Sylvester equations; SVID, dual updates, scale extraction, and sign extraction remain as in NanoQuant.Source: §4.3 · pp. 7–8
  • Table 1 covers WikiText-2 perplexity and zero-shot mean across reasoning tasks for Qwen3-Base models at 1.0, 0.8, and 0.55 bpw budgets.It compares models at the same total bits as NanoQuant and includes reported and reproduced values.Source: Table 1 · p. 8

5 EXPERIMENTS

At identical bit budgets, ShamAN-Q improves WikiText-2 perplexity over NanoQuant, with the size of the gain varying across model sizes and allocator choices. Experiments also assess setup, quantization cost, and cumulative contributions of the method’s modifications.

5.1 EXPERIMENTAL SETUP

  • The evaluation quantizes Qwen3-Base models from 0.6B to 14B using WikiText-2 calibration and reports WikiText-2 perplexity and mean zero-shot accuracy; all results are single runs.ShamAN-Q and NanoQuant use the same total number of bits, with reported budgets spanning approximately 0.55, 0.8, and 1 bpw.Source: §5.1 · p. 8
  • At 0.8 bpw, ShamAN-Q matches or beats NanoQuant at 1.0 bpw on 0.6B and 1.7B, while its perplexity is lower at every model size.At 1.7B and 4B, the 0.8-bpw margin widens from 13% to 26% and from 3% to 11%, respectively; at 8B it narrows to 0.4%.Source: §5.1 · p. 8

5.2 RANK ALLOCATION AT BIT PARITY

  • At 0.6B, the measured sensitivity-by-depth allocator is 2.5 perplexity below uniform ranks; at 4B, it gives 13.80 PPL and 0.463 zero-shot mean versus 13.55 and 0.450 for the hand table.Both 4B allocator results are single runs.Source: §5.2 · p. 9, Table 2 · p. 9

5.3 QUANTIZATION COST

  • ShamAN-Q retains NanoQuant’s deployment format, leaving inference kernels, memory footprint, and decoding throughput unchanged while incurring additional cost during quantization.The supplied passage identifies Table 3 as an end-to-end wall-clock comparison but gives no timing values.Source: §5.3 · p. 9, Table 3 · p. 9

5.4 ABLATION STUDIES

  • The cumulative 1.7B ablation finds tempering the largest individual gain, while the fresh input statistic and periodic refresh consistently help at block boundaries but have an aggregate effect within single-run noise.A dense NKP fit without tempering adds nothing over the diagonal baseline, whereas KL-Shampoo improves both untempered and tempered models; p = 1 and p = 1/4 are worse than p = 1/2.Source: §5.4 · p. 9, Table 4 · p. 9

6 CONCLUSION

ShamAN-Q replaces NanoQuant’s diagonal reconstruction geometry with a dense, KL-Shampoo-style Kronecker fit of the empirical Fisher matrix. At matched bit budgets, it lowers perplexity and increases zero-shot reasoning performance on Qwen3-Base models from 0.6B to 14B.

6 CONCLUSION

  • ShamAN-Q replaces NanoQuant’s diagonal reconstruction geometry with a dense metric estimated by a KL-Shampoo-style Kronecker fit of the empirical Fisher Information matrix.The continuous ADMM updates become Sylvester equation solutions, while the packed format remains unchanged.Source: §6 · p. 9
  • At the same total number of bits, ShamAN-Q lowers WikiText-2 perplexity and increases zero-shot reasoning performance on ≲1-bpw Qwen3-Base models from 0.6B to 14B.The method is intended to retain the speed of NanoQuant’s customized inference kernel through its unchanged packed format.Source: §6 · p. 9
  • Future work includes parallelizing computation, approximating dense curvature statistics, tuning the depth prior and calibration budget, and investigating scaling behavior.The depth prior and calibration budget were tuned at 0.6B.Source: §6 · p. 9

AI USE STATEMENT

The authors report using generative AI tools for implementation and several supporting tasks, while excluding AI involvement in proposing hypotheses, research methodology, experiments, or results interpretation. They state that AI-assisted work was reviewed, and that humans reviewed code tests and revised generated figures and writing.

AI USE STATEMENT

  • The authors used generative AI tools for implementing methods, but not for proposing or refining hypotheses, designing or providing feedback on methodology or experiments, or interpreting results.Other listed tasks, including synthetic data generation and proof assistance, are marked not applicable.Source: AI USE STATEMENT · p. 10
  • AI tools were also used to create scientific figures, suggest experimental parameters, brainstorm, identify some literature, and edit.The authors state that they reviewed all AI-assisted work.Source: AI USE STATEMENT · p. 10
  • LLM-generated code was tested with human-reviewed tests, and AI-generated TikZ figures and writing were manually revised or reviewed by humans.The authors take responsibility for the final content, including AI-assisted text, claims, and artifacts.Source: AI USE STATEMENT · p. 10

ETHICS STATEMENT

The authors intend the work to reduce energy consumption required to run AI models. They caution that lower energy costs per use could increase aggregate consumption through wider access and adoption, though this outcome is not guaranteed.

ETHICS STATEMENT

  • The work is intended to reduce energy consumption required to run AI models, amid concerns about resource scarcity, limited production capacity, and environmental impact.Source: ETHICS STATEMENT · p. 10
  • The authors caution that lower energy costs for use could increase aggregate energy consumption if access and adoption expand sufficiently.Source: ETHICS STATEMENT · p. 10
  • The authors describe this possible aggregate increase as a Jevon’s paradox effect and state that it is not guaranteed.Source: ETHICS STATEMENT · p. 10

A IMPLEMENTATION DETAILS

The implementation follows NanoQuant's protocol except where ShamAN-Q changes it, adding KL-Shampoo passes, curvature refreshes, a fresh input statistic, and rank allocation. The paper also reports hardware and engineering measures for managing the dense recipe's cost.

A IMPLEMENTATION DETAILS

  • Runs use NanoQuant’s protocol with 128 WikiText-2 calibration sequences of 2048 tokens, shrinkage γ = 0.2, and K = 400 LB-ADMM iterations.The 8B and 14B models use 512 sequences for curvature statistics; their refresh periods are ∆ = 9 and ∆ = 10, respectively.Source: A IMPLEMENTATION DETAILS · p. 12
  • ShamAN-Q adds three KL-Shampoo calibration passes, tempering exponent p = 1/2, periodic refreshes, a fresh input statistic, and a rank allocator.Refresh periods are ∆ = 7 for 0.6B and 1.7B, and ∆ = 9 for 4B and 8B.Source: A IMPLEMENTATION DETAILS · p. 12
  • At 4B, grouping calibration statistics by layer budget reduced calibration from about 5 hours to about 35 minutes.Caching fresh input measurements and eigendecompositions by shared-input group saves three of seven measurements and three large eigendecompositions per block.Source: A IMPLEMENTATION DETAILS · p. 12
  • Although reported 1.7B and 4B results used fp64, n × n eigendecompositions can run in fp32 with eigenvalue clamping before sign projection; at 1.7B, perplexity was unchanged (16.66 vs. 16.72).Source: A IMPLEMENTATION DETAILS · p. 12

B RANK ALLOCATION DETAILS

Rank allocation preserves NanoQuant’s total bit budget while distributing ranks across layers, and the accompanying algorithms describe curvature preparation, layer reconstruction, and unchanged projection and tuning stages.

B RANK ALLOCATION DETAILS

  • Both allocators start from uniform ranks and preserve total Nbits, with rank ceilings set by each layer’s dimensions.The appendix describes the measured and hand-table allocation procedures as completing §4.2.Source: B RANK ALLOCATION DETAILS · p. 12
  • The hand allocator sets per-layer budgets using a depth prior and layer-type weights, then rounds ranks in 32-rank steps to preserve the total bit budget.The type weights range from 0.85 for q to 1.15 for down, based on median relative block-loss jumps when layer types were binarized in 4B logs.Source: B RANK ALLOCATION DETAILS · p. 12
  • The measured allocator probes each layer at three candidate ranks and greedily assigns 32-rank steps by predicted loss decrease per bit until the total bit budget is reached.Its curves rank layer types within a block well (Spearman 0.8–0.96 against logged block-loss jumps) but under-value late blocks, which the depth prior compensates for.Source: B RANK ALLOCATION DETAILS · pp. 12–13

C ALGORITHMS

  • Metric preparation forms transformed weights, transports and tempers the shrunken statistics without refitting, then caches the resulting metric state and eigensystems.The transformed target is M := G fW A, using the tempered transported statistics.Source: C ALGORITHMS · p. 13
  • The continuous KRON-LB-ADMM updates use two systems solved by Eq. 16, while SVID projection, scale extraction, and final return to original weight coordinates follow NanoQuant’s procedures.The Euclidean consensus penalties mean the SVID projection and its inner iteration count do not depend on A or G.Source: C ALGORITHMS · pp. 13–14
  • The pipeline calibrates and initially allocates ranks on the full-precision teacher, refreshes statistics for unfactorized layers periodically, and measures fresh input statistics before layer factorization.TUNEFP is repeated between layer factorizations; latent tuning and scale distillation retain NanoQuant’s procedures.Source: C ALGORITHMS · pp. 14–15

D COMPLEXITY ANALYSIS

The complexity accounting separates dense setup from repeated reconstruction costs and reports that the per-iteration rank-sized eigendecomposition dominates at 4B.

D COMPLEXITY ANALYSIS

  • Table 5 separates per-layer calibration, preparation, and one outer ADMM sweep for an m × n weight, rank r, and N calibration tokens.Dense setup includes forming M and eigendecomposing the two transported statistics.Source: D COMPLEXITY ANALYSIS · p. 15, Table 5 · p. 16
  • Each layer performs two continuous solves per outer iteration, with two SVID projections at initialization and two more per iteration.The runtime accounting distinguishes setup, one pair of continuous solves, one pair of proxy projections, and diagnostic evaluations.Source: D COMPLEXITY ANALYSIS · p. 16
  • At 4B, removing every n^2r product changes per-layer ADMM time by only 5–8%, because the repeated r × r eigendecomposition dominates.At allocated ranks around 1500–2500, each layer performs 800 r × r eigendecompositions over K = 400 iterations.Source: D COMPLEXITY ANALYSIS · p. 16

E DESIGN ALTERNATIVES NOT ADOPTED

The evaluated alternatives were not adopted; several speed shortcuts yielded limited savings, while reducing TUNEFP rounds incurred a measurable perplexity cost.

E DESIGN ALTERNATIVES NOT ADOPTED

  • Reducing TUNEFP to one round per shared-input group cost +0.28 block-3 perplexity on the 0.6B four-block screen, outside single-run noise.The control was 14.15 with about 0.2 single-run noise; the authors conclude every round matters, including between layers reading the same activation.Source: E DESIGN ALTERNATIVES NOT ADOPTED · p. 16
  • Carrying an explicit per-rank middle scale vector through tuning and distillation was 3.5–5 perplexity worse at 0.6B, while mean-one rebalancing was within noise.The explicit vector’s learning-rate sweep was flat to negative.Source: E DESIGN ALTERNATIVES NOT ADOPTED · p. 16
  • At 4B, the low-rank-plus-identity Sylvester step saved only 5–8% of ADMM time because the r × r eigendecomposition dominated, and it was not used.Its projection cost +0.10 to +0.34 block-3 perplexity at k = 64 per side and was within noise at k = 256.Source: E DESIGN ALTERNATIVES NOT ADOPTED · p. 16
Loading 2609.38521v1…