SQuAT: Spherical Lattice Vector Quantization-Aware Training

Sebastian Loeschcke; Ivan Koryakovskiy; Tycho F. A. van der Ouderaa; Riccardo Del Chiaro; Mart van Baalen; Paul N. Whatmough; Markus Nagel

cs.LGOct 22, 2026

SQuAT: Spherical Lattice Vector Quantization-Aware Training

Authors: Sebastian Loeschcke, Ivan Koryakovskiy, Tycho F. A. van der Ouderaa, Riccardo Del Chiaro, Mart van Baalen, Paul N. Whatmough, Markus Nagel

Organizations: University of Copenhagen · Qualcomm AI Research

Abstract

Vector quantization (VQ), in particular lattice-based methods, offer strong rate-distortion trade-offs for low-bit large language models (LLMs), yet their behavior under quantization-aware training (VQ-QAT) remains poorly understood. We scale lattice VQ-QAT to billion-parameter LLMs and multi-billion-token training regimes, and introduce SQuAT, a shape-gain VQ-QAT approach trained with a separate straight-through estimator (STE) for the shape and gain. SQuAT decomposes weights into direction and magnitude: normalized directions are quantized using an $E_8$ -derived spherical codebook, while gains are quantized independently with low-bit scalar quantizers. We show that, unlike a standard STE applied to the full quantizer, which ignores the shape-gain structure, these separate STEs induce tangent-space updates for directions, yielding an implicit form of spherical (Riemannian) optimization. To enable scalability, we develop a distributed exact nearest-neighbor assignment scheme, in which workers quantize disjoint weight blocks and communicate only discrete indices, reducing VQ-QAT training time by 12.8 $\times$ and allowing us to scale it to 10B tokens. Across instruction-tuned LLaMA-3.2-1B/3B, Qwen3-4B, and base Qwen3-1.7B, SQuAT consistently outperforms standard VQ-QAT STE under identical quantizers and initialization, and achieves SOTA VQ-QAT accuracy.

Figures & tables

Figure 1: Overview of SQuAT . Left: a weight block is decomposed as

w=\gamma s

, where

s\in\mathbb{S}^{d-1}

is the shape and

\gamma\in\mathbb{R}^{+}

is the gain. The two components are quantized independently and reconstructed as

\hat{w}=\hat{\gamma}\hat{s}

. Middle: our spherical factored STE applies separate STEs to the gain and shape quantizers, mapping the gradient

g=\partial\mathcal{L}/\partial\hat{w}

to a scalar gain gradient

\nabla_{\gamma}\mathcal{L}\approx\hat{s}^{\top}g

and a tangent shape signal

\hat{\gamma}P_{s}g

T_{s}\mathbb{S}^{d-1}

. Right: the tangent shape signal induces the update

u^{+}=u-\eta\,\hat{\gamma}(\|u\|)^{-1}P_{s}g

on the shape parameter, followed by normalization (retraction back to the unit sphere),

s^{+}=u^{+}/(\|u^{+}\|)

Figure 2: SQuAT improves WikiText2 perplexity over integer scalar (INT-QAT) at 2.25 bpw on LLaMA-3.2-1B.

Figure 3: SQuAT reduces training time via progressively added fused kernels, index caching, and sharded assignment, approaching INT2 (4.4 s/it) and BF16 (4.3 s/it) performance, and allows us to scale VQ-QAT to billions of tokens.

Method	WikiText2 PPL $\downarrow$			CSR, % $\uparrow$			MMLU, % $\uparrow$			GSM8K, % $\uparrow$
Method	L1B	L3B	Q4B	L1B	L3B	Q4B	L1B	L3B	Q4B	L1B	L3B	Q4B
FP16	12.15	10.49	9.41	61.0	67.5	72.9	45.4	59.6	72.6	36.8	65.8	85.4
Tok = 1B
Standard	14.22	11.04	10.17	58.4	66.3	70.2	39.3	55.0	67.1	14.7	47.0	76.3
Spherical	13.85	10.96	10.03	58.8	65.8	71.0	40.3	55.8	67.4	18.0	51.9	77.0
Tok = 5B

Table 1: Standard vs Spherical VQ-QAT ( SQuAT ) results obtained with the same E8S shape-gain quantizer. Standard VQ-QAT applies STE to the reconstructed block, while Spherical VQ-QAT applies separate STEs to the shape and gain quantizers. App. E.1 extends the table to IFEval and LongBench generative tasks.

Method	Tok	bpw	WikiText2 PPL $\downarrow$			CSR, % $\uparrow$			MMLU, % $\uparrow$
Method	Tok	bpw	L1B	L3B	Q4B	L1B	L3B	Q4B	L1B	L3B	Q4B
FP16	–	16.0	12.15	10.49	9.41	61.0	67.5	72.9	45.4	59.6	72.6
QuIP# (fine-tuned)	52M	2.0	16.67	11.13	–	52.0	61.0	–	26.5	42.1	–
QTIP	52M	2.0	14.47	10.12	8.87	54.3	62.7	69.5	33.0	49.4	62.7
PV-tuning	1B a	2.0	–	–	–	–	–	–	28.9	48.4	–
INT-QAT	–	2.0	2e6	2e5	6e6	36.5	35.5	36.8	25.5	24.7	26.9

Table 2: QAT scaling across instruction-tuned LLaMA and Qwen models. We report WikiText2 perplexity, CSR, and MMLU as a function of QAT token budget. Missing entries indicate results that were not run or not reported under the same protocol.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

	Joint / Standard STE	Factored (Shape-Gain) STE
	$\displaystyle\begin{array}[]{c}\frac{\partial\hat{w}}{\partial w}\approx I\end{array}$	$\displaystyle\begin{array}[]{c}\frac{\partial\hat{s}}{\partial s}\approx I,\frac{\partial\hat{\gamma}}{\partial\gamma}\approx 1\end{array}$
Euclidean parameterization ( $w$ )	$\displaystyle\begin{array}[]{c}w^{+}=w-\eta\,g\\[3.99994pt] \text{Sec.~\ref{sec:optA}}\end{array}$	$\displaystyle\begin{array}[]{c}w^{+}=w-\eta\!\left(\frac{\hat{\gamma}}{\gamma}P_{s}g+(\hat{s}^{\top}g)\,s\right)\\[3.99994pt] \text{Sec.~\ref{sec:optE_factored}}\end{array}$
Spherical parameterization ( $u,\gamma$ )	$\displaystyle\begin{array}[]{c}u^{+}=u-\eta\frac{\gamma}{\\|u\\|}P_{s}g\\ \gamma^{+}=\gamma-\eta(s^{\top}g)\\[3.00003pt] w^{+}=\gamma^{+}\,\frac{u^{+}}{\\|u^{+}\\|}\\[3.99994pt] \text{App.~\ref{sec:optBprime}}\end{array}$	$\displaystyle\begin{array}[]{c}u^{+}=u-\eta\frac{\hat{\gamma}}{\\|u\\|}P_{s}g\\ \gamma^{+}=\gamma-\eta(\hat{s}^{\top}g)\\[3.00003pt] w^{+}=\gamma^{+}\,\frac{u^{+}}{\\|u^{+}\\|}\\[3.99994pt] \text{Sec.~\ref{sec:optB}}\end{array}$

Appendix

Table 4: Effect of STE choice and parameterization on the effective weight update

w^{+}

bpw	Scheduler	Hadamard	WikiText2 perplexity $\downarrow$		CSR, % $\uparrow$		MMLU, % $\uparrow$
bpw	Scheduler	Hadamard	Signed	SEQ	Signed	SEQ	Signed	SEQ
FP16			12.15		61.0		45.4
2.0	Cosine		18.07	16.97	54.5	55.7	26.0	28.5
	WSD		16.91	16.25	54.7	56.3	28.1	30.8
	WSD	✓	15.91	15.07	56.3	57.3	28.3	31.7
2.25	Cosine		16.71	16.32	55.1	55.8	31.8	30.9

Appendix

Table 5: The influence of learning rate scheduler, Hadamard transformation, and quantizer on the performance of the quantized LLama-3.2-1B model after 1B token INT-QAT.

Method	Granularity	Tokens	Weight/Shape LR	Scale/Gain LR	Model
SQuAT (E8S-SG)	all bpv variants	1B, 5B, 10B	$3\times 10^{-5}$	$3\times 10^{-2}$	L1B, L3B,
SQuAT (E8S-SG)	all bpv variants	1B, 5B, 10B	$3\times 10^{-5}$	$3\times 10^{-2}$	Q1B7, Q4B
Standard VQ-QAT	all bpv variants	1B, 5B	$3\times 10^{-5}$	—	L1B, L3B
Standard VQ-QAT	all bpv variants	1B, 5B	$1\times 10^{-4}$	—	Q4B
E8S	all bpv variants	1B	$1\times 10^{-4}$	$1\times 10^{-2}$	L1B
SQuAT (E8P12-SG)	all bpv variants	1B	$3\times 10^{-5}$	$3\times 10^{-2}$	L1B

Appendix

Table 6: Learning rates selected for experiments.

Method	Tok	IFEval, % $\uparrow$			LongBench, % $\uparrow$
Method	Tok	L1B	L3B	Q4B	L1B	L3B	Q4B
FP16	–	57.0	73.9	86.3	28.4	23.8	22.4
Standard VQ-QAT	1B	37.8	63.0	79.8	21.1	19.4	17.9
Spherical VQ-QAT ( SQuAT )	1B	41.4	62.5	82.3	23.6	21.6	18.1
Standard VQ-QAT	5B	40.8	59.5	78.9	24.5	18.4	17.2
Spherical VQ-QAT ( SQuAT )	5B	43.6	64.2	82.4	23.4	20.7	17.6

Appendix

Table 7: Ablation of the STE used in SQuAT at 2.25 bpw across LLaMA-3.2-1B/3B and Qwen3-4B on generative tasks. Both variants use the same E8S shape-gain quantizer and PTQ initialization.

		Standard VQ-QAT			Spherical VQ-QAT (SQUAT)
		PPL	CSR	MMLU	PPL	CSR	MMLU
1B training tokens
	Seed 0	14.22	58.43	39.28	13.85	58.77	40.26
	Seed 1	14.29	58.18	36.54	13.90	58.68	39.50
	Seed 2	14.13	58.27	38.53	13.78	58.65	40.73
	Mean	14.21	58.29	38.12	13.84	58.70	40.16

Appendix

Table 8: Per-seed results for Naive vs. Spherical gradient on LLaMA-3.2-1B-Instruct at 2.25 bpw. PPL: WikiText2 perplexity (

\downarrow

). CSR: common-sense reasoning (%,

\uparrow

). MMLU (%,

\uparrow

). For 1B token training, ranges do not overlap for any metric, confirming that the Spherical improvement exceeds run-to-run noise.

Setting	PPL / CSR / MMLU
LLaMA-3.2-1B-Instruct
Explicit gain parameter ( SQuAT , Sec. 3.2 )	13.85 / 58.8 / 40.3
No separate gain parameter ( Sec. B.2 )	14.01 / 58.2 / 38.4
Without training shape weights	13.85 / 58.2 / 40.0
No shape-gain (E8S)	14.24 / 58.4 / 38.2
LLaMA-3.2-3B-Instruct

Appendix

Table 9: Ablation of explicit gain parameterization and shadow-weight updates under a fixed 2.25 bpw budget using 1B tokens and the E8S codebook.

Panel of Figure 5 — Figure 5: Scaling behavior on LLaMA-3.2-1B as a function of QAT training tokens. We compare INT-QAT and SQuAT across WikiText2 perplexity, CSR, and MMLU.

Method	Tok	WikiText2 PPL $\downarrow$			CSR, % $\uparrow$			MMLU, % $\uparrow$
Method	Tok	L1B	L3B	Q4B	L1B	L3B	Q4B	L1B	L3B	Q4B
FP16	–	12.15	10.49	9.41	61.0	67.5	72.9	45.4	59.6	72.6
Standard VQ-QAT (Sec. 3.1 )	1B	14.10	11.04	10.17	58.7	66.3	70.2	37.5	55.0	67.1
Euclidian $w$ + factored STE (Sec. B.2 )	1B	14.01	11.11	10.15	58.2	66.0	70.8	38.4	54.6	66.5
Spherical VQ-QAT ( SQuAT , Sec. 3.2 )	1B	13.85	10.96	10.03	58.8	65.8	71.0	40.3	55.8	67.4

Appendix

Table 10: Additional comparison of Euclidean parametrization and Standard vs Spherical VQ-QAT results. Both variants use the same E8S shape-gain quantizer and PTQ initialization. Standard VQ-QAT applies STE to the reconstructed block, while Spherical VQ-QAT applies separate STEs to the shape and gain quantizers. App. E.1 shows results on multiple generative tasks.

	PPL 2048 $\downarrow$	PPL 4096 $\downarrow$	AC	AE	BQ	HS	PQ	WG	CSR $\uparrow$	MMLU $\uparrow$
FP16	16.72	15.68	43.08	69.69	77.52	60.37	72.14	61.80	64.10	60.12
LC-QAT	–	–	39.51	64.35	70.43	55.01	71.65	57.54	59.75	–
SQuAT	17.36	16.19	39.76	65.87	75.66	54.06	69.97	61.72	61.17	52.94

Appendix

Table 11: Qwen3-1.7B quantization results at 2.0 total bpw. We use E8S-SG codebook in SQuAT method. PPL is WikiText2. CSR is the average over 6 common sense reasoning tasks. LC-QAT and SQuAT both use 1B training tokens. PPL and MMLU are not reported for the 1B-token LC-QAT. AE, AC, BQ, HS, WG, and PQ denote ARC-Easy, ARC-Challenge, BoolQ, HellaSwag, WinoGrande, and PIQA, respectively.

Setting	bpw	Hadamard	PPL / CSR / MMLU
SQuAT (grouped)	2.25	Yes	13.85 / 58.8 / 40.3
		No	14.89 / 58.2 / 36.5
		$\Delta$	1.04 / 0.6 / 3.8
No SG (grouped)	2.25	Yes	14.53 / 58.0 / 37.0
		No	15.26 / 57.1 / 33.2
		$\Delta$	0.73 / 0.9 / 3.8

Appendix

Table 12: Ablation of Hadamard rotations with and without shape-gain on LLaMA-3.2-1B (E8S-SG) after 1B QAT.

Setting	bpw	Hadamard	PPL / CSR / MMLU
SQuAT	2.25	Yes	13.94 / 58.7 / 38.9
	2.25	No	14.36 / 58.3 / 36.9
		$\Delta$	0.42 / 0.3 / 1.9
No SG	2.25	Yes	14.67 / 57.5 / 35.1
	2.25	No	15.44 / 57.2 / 33.8
		$\Delta$	0.77 / 0.3 / 1.3

Appendix

Table 13: Ablation of Hadamard rotations with and without shape-gain on LLaMA-3.2-1B (E8P12-SG) after 1B QAT.

Figure 6: Pareto fronts of LLaMA-3.2-1B quantized with SQuAT (E8S-SG) vs. SEQ after 1B-token QAT. Each frontier corresponds to a different gain bitwidth. The normalized E8S spherical codebook supports fractional bitwidths at finer granularity than scalar quantization, enabling more precise rate-distortion trade-offs.

Figure 7: Scale-quantization ablation with shape-gain. Chi-based gains perform best at very low bit-rates, while all methods approach FP8-level quality as gain precision increases.

LLM Quantization

Large language model quantization reduces numerical precision to lower memory use and accelerate training or inference.

See all 381 papers

LC-QAT: Data-Efficient 2-Bit QAT for LLMs via Linear-Constrained Vector Quantization
Jun 9, 2026Haoyu Wang, Xingyu Yu, Haiyan Zhao +2LLM Quantization Quantization-Aware Training
Department of Computer Science and Technology, Tsinghua University, Beijing, China · School of Software and Microelectronics, Peking University, Beijing, China · College of Computer Science and Technology, National University of Defense Technology, Hunan, China.
UniSVQ: 2-bit Unified Scalar-Vector Quantization
Jun 9, 2026Haoyu Wang, Haiyan Zhao, Xingyu Yu +4LLM Quantization Vector Quantization
Department of Computer Science and Technology, Tsinghua University, Beijing, China · School of Software and Microelectronics, Peking University, Beijing, China · School of Physics, Beihang University, Beijing, China
E8-QAT: Efficient 2-Bit Quantization-Aware Training with Implicit Lattice Codebooks
Oct 20, 2026Xinyu Lei, Hong Huang, Jiebin Zhang +6LLM Quantization Quantization-Aware Training
Peking University · City University of Hong Kong · Tencent

Jun 9, 2026J/K move · Enter open · S save

LC-QAT: Data-Efficient 2-Bit QAT for LLMs via Linear-Constrained Vector Quantization

Haoyu Wang, Xingyu Yu, Haiyan Zhao, Fengxiang Wang+1

Department of Computer Science and Technology, Tsinghua University, Beijing, China · School of Software and Microelectronics, Peking University, Beijing, China · College of Computer Science and Technology, National University of Defense Technology, Hunan, China.

LLM Quantization Quantization-Aware Training Vector Quantization Weight-Only Quantization Neural Network Quantization Data-Efficient Learning

Open paper AskPDF arXiv

Quantization-aware training (QAT) is essential for extremely low-bit large language models (LLMs). Current QAT methods are mainly based on scalar quantization (SQ), which enables efficient optimization but suffers from severe performance degradation at 2-bit precision. On the other hand, vector quantization (VQ) provides substantially higher representational capacity, but its discrete codebook lookup prevents end-to-end training. We propose LC-QAT, a 2-bit weight-only VQ-QAT framework that represents quantized weights via a learned affine mapping over discrete vectors, which yields a high-quality PTQ initialization and enables fully differentiable end-to-end optimization without explicit codebook lookup in the training forward pass. This strong post-training initialization makes LC-QAT highly data-efficient. Experiments across diverse LLMs demonstrate that LC-QAT consistently outperforms state-of-the-art QAT methods while using only 0.1%--10% of the training data. Our results establish LC-QAT as a practical and scalable solution for extreme low-bit model deployment. Codes are publicly available at https://github.com/AI9Stars/UniSVQ.

Abstract

LLM Quantization

Explore similar work

LC-QAT: Data-Efficient 2-Bit QAT for LLMs via Linear-Constrained Vector Quantization

UniSVQ: 2-bit Unified Scalar-Vector Quantization

E8-QAT: Efficient 2-Bit Quantization-Aware Training with Implicit Lattice Codebooks