How does SQuAT make lattice quantization work for billion-parameter models?

Why standard VQ-QAT hits a wall

Vector quantization packs many weights into a single codebook entry. Lattice codebooks such as the one derived from the $E_8$ root system give excellent rate-distortion trade-offs. The problem appears when you try to train a quantized model end to end. A straight-through estimator applied to the whole quantizer treats the codebook as a flat lookup table. It ignores the fact that lattice points live on a sphere and that each weight has a natural split into direction and length. Gradient updates then push weights off the manifold where the codebook actually sits. The result is unstable training that has never been demonstrated beyond tiny models and token counts.

Standard VQ-QAT fails at scale.

Splitting shape from gain

SQuAT decomposes every weight vector into a normalized direction and a scalar gain. Directions are quantized against the spherical $E_8$ codebook. Gains go through a low-bit scalar quantizer. Each path gets its own straight-through estimator. The direction estimator produces updates that stay in the tangent space of the sphere. That is equivalent to running Riemannian optimization without ever computing an exponential map. The gain estimator works in ordinary Euclidean space. Because the two estimators are separate, the direction update never leaks into the gain and vice versa. The paper shows that a single joint STE cannot replicate this behavior. It also notes that the implicit spherical geometry emerges automatically from the separate estimators. No extra projection step is required.

Separate STEs change everything.

The direction update stays on the manifold. The gain update stays in Euclidean space. They never interfere. A single joint STE cannot replicate this behavior. The paper shows the implicit spherical geometry emerges automatically from the separate estimators. No extra projection step is required.

Distributed nearest-neighbor search that actually scales

Exact nearest-neighbor assignment is the bottleneck. The authors partition weight blocks across workers. Each worker quantizes its local blocks against the full codebook and sends back only the discrete indices. No floating-point vectors cross the network. Communication volume drops to a few bytes per parameter. Wall-clock time for the VQ-QAT step falls by a factor of 12.8. That reduction makes it feasible to run quantization-aware training on ten billion tokens. The scheme is exact. Approximate methods such as product quantization or IVF are not used. The paper does not report the absolute training time in hours. It also does not specify the cluster topology or the interconnect bandwidth.

What changes on real models

Experiments cover instruction-tuned LLaMA-3.2-1B and LLaMA-3.2-3B. They also cover Qwen3-4B and base Qwen3-1.7B. All runs use identical quantizers and initialization. SQuAT beats standard VQ-QAT STE on every configuration. The gap widens at lower bit widths. The method reaches state-of-the-art accuracy for VQ-QAT on these architectures. The paper does not compare against non-VQ quantization such as GPTQ or AWQ. It does not evaluate perplexity on held-out benchmarks beyond the training distribution. It does not measure inference latency or memory bandwidth on hardware. The code and checkpoints are not yet linked from the arXiv entry.