SQuAT: Spherical Lattice Vector Quantization-Aware Training

1 INTRODUCTION

• VQ-QAT remains underexplored at billion-parameter and multi-billion-token scales, and prior work has not used shape-gain decomposition to align QAT updates with spherical geometry.The resulting question is how to derive a QAT rule whose gradient estimation and parameter updates fit spherical VQ geometry.Source: §1 · p. 2
• SQUAT decomposes weight blocks into direction and gain, quantizes them independently, and applies separate STEs through the E8-derived spherical direction and low-bit scalar gain quantizers.Unlike an STE on the reconstructed weight, this factored STE induces tangent-space updates for directions and decouples magnitude optimization into the gain parameter.Source: §1 · p. 2
• SQUAT's normalized spherical codebook enables fractional bitwidths and Pareto-optimal rates that scalar methods cannot access.The paper identifies E8-derived spherical codebooks, including its E8S construction, as part of the shape quantizer.Source: §1 · pp. 2–3
• Distributed exact assignment lets workers quantize disjoint weight blocks and exchange only discrete indices, reducing VQ-QAT training time by 12.8×.The implementation supports billion-parameter models and multi-billion-token training.Source: §1 · pp. 2–3
• SQUAT improves over standard VQ-QAT under matched quantizers and initialization, achieving state-of-the-art VQ-QAT accuracy.The evaluation covers instruction-tuned LLaMA-3.2-1B/3B and Qwen3-4B, plus base Qwen3-1.7B, with budgets from 0.25B to 10B tokens.Source: §1 · p. 3
• At 2.25 bpw, SQUAT beats INT-QAT on WikiText2 perplexity across models and training budgets, with MMLU gains up to 7.9% at 1B token training.The Figure 2 comparison is for LLaMA-3.2-1B; the broader WikiText2 and MMLU claims are reported across models and budgets.Source: Figure 2 · p. 2, §1 · p. 3

2.1 SCALAR QUANTIZATION

• Scalar quantization maps each weight independently to low-precision levels, with QAT using a straight-through estimator to propagate gradients.The scalar baseline described here is Stretched Elastic Quantization (SEQ).Source: §2.1 · p. 3

2.2 VECTOR AND LATTICE QUANTIZATION

• Vector quantization maps weight blocks to codewords and can improve rate–distortion trade-offs over scalar quantization, while lattice structure avoids explicit exponentially growing codebooks.VQ-QAT still requires repeated nearest-neighbor assignment during training.Source: §2.2 · p. 3

2.3 SHAPE-GAIN (POLAR) DECOMPOSITION

• Shape–gain quantization decomposes weights into direction and magnitude, quantizes them independently, and separates angular error from gain error.Prior work cited in the passage suggests most representational budget should go to direction, while gains can use few bits.Source: §2.3 · pp. 3–4

2.4 LATTICE CODEBOOKS FROM E8

• E8-derived codebooks use one or more finite lattice shells; normalized codewords form spherical codes for shape–gain quantization.The paper studies E8P12 and introduces E8S, which selects low-norm E8 points and normalizes them.Source: §2.4 · p. 4

3.1 STANDARD VQ-QAT: EUCLIDEAN SHADOW WEIGHTS AND STANDARD STE

• A standard joint STE treats the whole quantizer as the identity in backpropagation, so shape–gain factorization alone does not change the resulting update.In the standard formulation, Euclidean shadow vectors are updated directly by the upstream gradient.Source: §3.1 · p. 5

3.2 SPHERICAL VQ-QAT: SPHERICAL PARAMETERIZATION AND STE

• SQUAT interprets separate shape–gain STEs with spherical parameterization as a tangent-space direction step plus retraction, coupled with a Euclidean gain update.The paper states the one-step SGD interpretation is exact and remains useful for common optimizers.Source: §3 · p. 4, §3.2 · p. 5
• Alternative parameterization and STE placements produce different effective updates, clarifying that both choices shape the optimization dynamics.The appendix variants use spherical variables with a joint STE or Euclidean weights with a factored STE.Source: §3.2 · p. 6

3.3 PRACTICAL QUANTIZERS

• SQUAT uses normalized E8P12 or E8S codebooks for directions and a chi-quantile gain quantizer whose levels add no learned codebook parameters.The gain quantizer requires no extra bits beyond those used to index its gain levels.Source: §3.3 · p. 6

3.4 SCALABILITY AND IMPLEMENTATION

• 12.8× speedup reduces E8S VQ-QAT training from 62.5 s/it to 4.9 s/it while preserving exact nearest-neighbor assignment.The implementation partitions weight blocks across workers, synchronizes discrete indices, fuses kernels, and caches quantized weights across gradient-accumulation steps.Source: §3.4 · p. 6

4 EXPERIMENTAL SETUP

• Experiments use instruction-tuned LLaMA-3.2-1B/3B and Qwen3-4B, plus base Qwen3-1.7B, training for 0.25B–10B tokens.QAT uses DCLM-Edu, JSD distillation from a frozen FP16 teacher, batch size 256, and context length 1024.Source: §4 · pp. 6–7
• Table 1 compares Standard VQ-QAT and SQUAT using the same E8S shape–gain quantizer; Figure 3 depicts progressively added training optimizations and speed comparisons.The Figure 3 caption reports INT2 at 4.4 s/it and BF16 at 4.3 s/it.Source: Figure 3 · p. 6, Table 1 · p. 7
• The baselines include scalar INT-QAT using SEQ and VQ or PTQ methods including QTIP and PV-Tuning where available.Source: §4 · p. 7
• Evaluation reports WikiText2 perplexity, common-sense reasoning and MMLU accuracy, plus GSM8K, IFEval, and LongBench results.The common-sense reasoning score averages six tasks.Source: §4 · p. 7

Standard VQ-QAT vs Spherical VQ-QAT (SQUAT)

• Spherical VQ-QAT improves WikiText2 perplexity and GSM8K accuracy across tested models and budgets, with GSM8K gains reaching 4.9%.CSR and MMLU improve mainly at the 1B-token budget, except LLaMA-3.2-3B CSR.Source: §5 · p. 7

Scaling across model family and QAT budget

• At 2.25 bpw and 10B tokens, WikiText2 perplexity falls from 17.90 to 13.74 on LLaMA-3.2-1B, from 15.00 to 10.83 on LLaMA-3.2-3B, and from 12.52 to 9.90 on Qwen3-4B after QAT.On Qwen3-4B, MMLU rises from 54.2 after PTQ to 69.5 after QAT.Source: §5 · p. 7, Table 2 · p. 8
• At 2.25 bpw and 10B tokens, SQUAT matches or surpasses INT-QAT on nearly all evaluated models and metrics; Qwen3-4B CSR is the exception.At 2.0 bpw, the methods are broadly comparable; SQUAT tends to do better below 5B tokens, while INT-QAT often gains with more training.Source: §5 · p. 7
• Compared with QTIP, QUIP#, and PV-tuning, SQUAT sometimes has higher perplexity but substantially stronger downstream accuracy; on Qwen3-1.7B, CSR is 61.2 versus LC-QAT’s 59.8 after 1B tokens.The LC-QAT comparison is a 1.4% advantage for SQUAT.Source: §5 · p. 8

Chi-based gain quantization

• At matched bitwidths, fixed chi-based gain scaling outperforms 1-bit uniform integer gain quantization, underscoring distribution-aware design at extreme quantization levels.Source: §5 · p. 8

E8S vs. E8P12

• Under matched 1B-token training on LLaMA-3.2-1B, E8S outperforms E8P12 both with and without shape-gain; at 2.25 bpw, SQUAT consistently improves results for both codebooks.Source: §5 · p. 8

Pareto fronts and gain bits allocation

• The best gain allocation shifts with target rate: 0.125 gain bpw below 2.0 total bpw, 0.25 up to 2.75 total bpw, and 0.375 at 3.0 total bpw.E8S supports fractional bitwidths for finer-grained compression-accuracy trade-offs than INT-QAT.Source: §5 · p. 8

Post-training quantization (PTQ)

• PTQ methods progressed from scalar round-to-nearest to Hessian-based reconstruction and approaches addressing outliers through activation-aware scaling and sparsity.QuIP uses Hadamard rotations; YAQA improves Hessian estimation.Source: §6 · pp. 8–9

Vector and lattice quantization

• VQ can exploit parameter correlations, but naive codebooks and nearest-neighbor search scale poorly; lattice, trellis, and pyramid structures address these costs.Examples include AQLM’s sums of codewords, GPTVQ, QuIP#, and QTIP.Source: §6 · p. 9
• PVQ and PCDVQ use shape-gain decomposition, while prior QAT for large vector or lattice codebooks remains limited and existing shape-gain methods primarily target PTQ.PVQ projects vectors onto a hypersphere using an implicit lattice; PCDVQ learns separate shape and gain codebooks.Source: §6 · p. 9

Quantization-aware training (QAT)

• Scalar QAT research includes clipping, step-size learning, stochastic rounding, low-rank adaptations, and long-horizon methods, while PV-Tuning and LC-QAT represent limited VQ-QAT work.PV-Tuning does not study longer token training; LC-QAT supports up to 10B QAT tokens.Source: §6 · p. 9

Discussion

• Shape-gain parameterization with spherical E8 codebooks improves over VQ-PTQ and standard joint-STE VQ-QAT under the same quantizer and initialization.Source: §7 · p. 9
• SQUAT outperforms INT-QAT at 2.25 bpw across almost all models, metrics, and QAT budgets, while remaining competitive at 2.0 bpw.Source: §7 · p. 9

Future work

• The work reports the first exact nearest-neighbor VQ-QAT demonstration at billion-parameter and multi-billion-token scale, enabled by distributed assignment.Source: §7 · p. 10

Limitations

• Evaluation is limited to E8P12 and E8S, models up to 4B parameters, and budgets up to 10B tokens; larger models, broader architectures, and higher-bit hierarchical schemes remain future work.Source: §7 · p. 10
• E8P12 has a fixed 2.0 bpw rate, whereas flexible-rate E8S currently lacks an efficient dequantizer.Fast dequantization for spherical lattice codebooks is presented as a promising direction.Source: §7 · p. 10

Conclusion

• SQUAT decomposes each weight block into normalized direction and scalar gain, quantizes shape with an E8-derived spherical codebook, and quantizes gains separately.Source: §8 · p. 10
• Separate STEs induce spherical optimization of quantized directions, and ablations show improvement over a standard VQ STE under the same quantizer and PTQ initialization.The ablations also show the importance of chi-based gain quantization, gain bit-rate selection, and shadow-weight updates.Source: §8 · p. 10
• Distributed exact nearest-neighbor assignment reduces training time from 62.5 s/it to 4.9 s/it; SQUAT improves over VQ-PTQ and joint-STE VQ-QAT while remaining competitive with INT-QAT.Source: §8 · p. 10

• LLMs assisted with implementing parts of the code and editing and polishing the manuscript wording.Source: AI USE STATEMENT · p. 10
• The authors reviewed, tested, and verified all LLM-generated code.Source: AI USE STATEMENT · p. 10
• The authors produced the experimental results, reported numbers, and tables themselves and take full responsibility for the paper.Source: AI USE STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

• The authors describe the method and report experimental details across the paper, and plan to release code and training scripts upon acceptance.Source: REPRODUCIBILITY STATEMENT · p. 10

A CONSTRUCTING AN E8 CODEBOOK AT A TARGET BIT BUDGET

• E8 codebooks meet a target bit budget by accumulating lattice shells and, when needed, selecting codewords from the final partial shell using spherical farthest-point sampling.For shape–gain codebooks, normalized directions make integer-scaled lattice vectors redundant; E8P12 has two duplicate direction pairs, leaving 65,534 unique directions from 65,536.Source: A CONSTRUCTING AN E8 CODEBOOK AT A… · pp. 15–16

B ALTERNATIVE RELATED CHOICES OF PARAMETERIZATION AND STE

• The appendix considers alternatives to the two primary combinations of parameterization and STE, including spherical weights with a joint STE and Euclidean weights with a factored STE.Source: B ALTERNATIVE RELATED CHOICES OF… · p. 16, §B.1 · p. 16, §B.2 · p. 17

B.1 WHAT IF WE USE SPHERICAL SHADOW WEIGHTS, BUT KEEP THE JOINT STE?

• With spherical shadow weights but a joint STE, the direction update is tangential, but depends on unquantized rather than forward-pass quantized values.Spherical VQ-QAT instead uses quantized values in its surrogate gradients, aligning updates with the effective model used during training.Source: §B.1 · p. 17

B.2 WHAT IF WE USE SPHERICALLY-FACTORED STE, BUT KEEP THE EUCLIDEAN SHADOW WEIGHTS?

• With Euclidean shadow weights, a factored STE decomposes updates into tangent and radial components but lacks sphere retraction, so it is not a clean Riemannian step.Source: §B.2 · p. 18

B.3 IMPLEMENTATION TRICK: PLACING STE SUCH THAT ‘STANDARD VQ-QAT’ UPDATES WITH SHAPE-GAIN PARAMETERS (u, γ)

• An outer joint STE can reproduce the Standard VQ-QAT update while storing parameters in a factored form, because the backward pass ignores the shape–gain decomposition.Source: §B.3 · p. 18

• PTQ initialization follows QuIP# with GPTQ-style column-wise reconstruction and Hessian-based error propagation; calibration uses 128 WikiText2 sequences of length 2048.Shape–gain uses gain groups of size 8; reported bit budgets omit auxiliary storage below 0.01 bpw.Source: C ADDITIONAL EXPERIMENTAL DETAILS · p. 19
• One-billion-token QAT runs use GKD with AdamW for 3,815 steps on DCLM-edu, with sequence length 1024 and effective batch size 256 across 8 GPUs.Source: C ADDITIONAL EXPERIMENTAL DETAILS · p. 19
• For the STE ablation, Table 7 compares variants at 2.25 bpw across LLaMA-3.2-1B/3B and Qwen3-4B using the same E8S shape–gain quantizer and PTQ initialization.Source: Table 7 · p. 20

E.1 RESULTS ON GENERATIVE TASKS

• Spherical VQ-QAT outperforms Standard VQ-QAT across IFEval and LongBench generative tasks.IFEval reports the mean of four accuracy variants; LongBench-E averages three category scores. Both are evaluated using greedy decoding.Source: D EVALUATION PROTOCOL · p. 20, §E.1 · p. 21

E.2 STATISTICAL STABILITY ACROSS RANDOM SEEDS

• Across three matched seeds, SQUAT beats Standard VQ-QAT on every metric at both 1B and 5B tokens, with mean PPL and CSR differences exceeding standard deviations.Table 8 reports per-seed results for LLaMA-3.2-1B-Instruct at 2.25 bpw.Source: §E.2 · p. 21, Table 8 · p. 21

E.3 EXPLICIT GAIN PARAMETERIZATION AND SHADOW-WEIGHT UPDATES

• After 1B tokens, explicit gain yields 13.85 vs. 14.01 perplexity and higher CSR and MMLU than implicit gain; freezing shape weights or removing shape–gain degrades results.Source: §E.3 · p. 21, Table 9 · p. 22

E.4 ISOLATING THE EFFECT OF THE STE CHOICE FROM THE SPHERICAL PARAMETERIZATION

• SQUAT improves over Standard VQ-QAT across all three models and metrics except LLaMA-3.2-3B CSR; its spherical parameterization and factored STE together produce a tangent update followed by retraction.The factored STE alone with Euclidean shadow weights shows no consistent improvement over Standard VQ-QAT.Source: §E.4 · p. 22

E.5 COMPARISON WITH LC-QAT

• 61.2 vs. 59.8 CSR: SQUAT exceeds LC-QAT on Qwen3-1.7B after 1B training tokens.LC-QAT reports no PPL or MMLU for that checkpoint, preventing direct comparison on those metrics.Source: §E.5 · p. 22, Table 11 · p. 23

E.6 HADAMARD ROTATIONS

• Hadamard rotations improve performance, while SQUAT’s shape–gain benefit persists with or without them; unfused rotations add latency and may be unsupported on some deployment targets.Source: §E.6 · p. 23, Table 12 · p. 24, Table 13 · p. 24

E.7 PARETO PLOT OF E8S-BASED SQUAT VS. SEQ.

• Below 2.0 bpw, 0.125 bpw of gain is preferred; at higher total bitwidth, gain quantization matters less, and E8S permits finer fractional bitwidths than scalar quantization.Source: §E.7 · p. 23, Figure 6 · p. 25

F CHI-BASED VS GROUPED INTEGER GAIN QUANTIZATION

• At 1–2 gain bits, chi-quantile gains outperform uniform integer gains; as precision rises, both approach FP8 performance, so chi-quantile is the low-bit default.The evaluation protocol defines WikiText2 perplexity, CSR, MMLU, GSM8K, IFEval, and LongBench-E metrics and decoding settings.Source: D EVALUATION PROTOCOL · p. 20, F CHI-BASED VS GROUPED INTEGER GAIN… · p. 24, Figure 7 · p. 25