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

Jun 9, 2026

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

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.

Haoyu Wang, Xingyu Yu, Haiyan Zhao +2

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

Ask PDF

Jun 9, 2026

UniSVQ: 2-bit Unified Scalar-Vector Quantization

Post-training quantization at the 2-bit level enables low-cost deployment and inference acceleration for large language models (LLMs). Scalar quantization (SQ) and vector quantization (VQ) are two primary quantization methods, however, the former suffers from significant performance degradation, and the latter incurs computational and storage overhead. We propose UniSVQ, a unified 2-bit quantization framework that bridges scalar and vector quantization by parameterizing codewords as an affine transform of integer lattices. This structure preserves compatibility with optimized integer kernels while retaining much of VQ's flexibility. We further introduce a data-driven block-wise fine-tuning strategy to directly minimize quantization reconstruction error. Extensive experiments across multiple LLM families and zero-shot benchmarks demonstrate that UniSVQ consistently outperforms state-of-the-art SQ methods and achieves performance comparable to advanced VQ methods, while providing higher inference throughput. Codes are publicly available at https://github.com/AI9Stars/UniSVQ.

Haoyu Wang, Haiyan Zhao, Xingyu Yu +4

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

LLM Quantization Vector Quantization Post-Training Quantization Low-Bit Quantization Neural Network Quantization

Ask PDF

Oct 20, 2026

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

Low-bit quantization reduces the memory footprint of large language models, but preserving accuracy with efficient training remains challenging. Scalar quantization encodes weights independently, while explicit-codebook vector quantization introduces costly search during training. We propose E8-QAT, a 2-bit quantization-aware training framework that uses the

E_8

lattice as an implicit eight-dimensional codebook. It jointly quantizes weight groups through arithmetic projection without storing or searching an explicit codebook during the training forward pass. A 2-bit core encoding and sparse extension preserve the selected lattice points exactly, adding less than 0.1% to core index storage. On Qwen3.6-35B-A3B, E8-QAT matches the performance of the full-precision model in terms of the average score across the four tasks. E8-QAT improves over the state-of-the-art baseline by 2.03% on IFEval and by 1.51% on LiveCodeBench, under the same data, schedule, and bit width. The 2-bit representation reduces peak memory by 80% and supports a

4\times

larger batch. At each representation's largest feasible batch, E8-QAT reaches

1.97\times

the decoding throughput of BF16.

Xinyu Lei, Hong Huang, Jiebin Zhang +6

Peking University · City University of Hong Kong · Tencent

LLM Quantization Quantization-Aware Training Low-Bit Quantization Vector Quantization

Ask PDF

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