SQuAT: Spherical Lattice Vector Quantization-Aware Training

📄 arXiv: 2610.31412v1 📥 PDF

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

分类: cs.LG, cs.CL

发布日期: 2026-10-22

备注: 25 pages, 7 figures

💡 一句话要点

提出SQuAT以提升大语言模型的量化训练效果

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 向量量化 量化感知训练 大语言模型 深度学习 优化算法

📋 核心要点

  1. 现有的量化感知训练方法对大语言模型的量化效果理解不足，导致性能提升有限。
  2. SQuAT通过将权重分解为方向和幅度，并使用独立的STE进行量化，优化了量化过程。
  3. 在多个大型模型上，SQuAT显著提高了量化训练的准确性，训练时间减少了12.8倍。

📝 摘要（中文）

向量量化（VQ），尤其是基于晶格的方法，为低比特大语言模型（LLMs）提供了良好的率失真权衡，但在量化感知训练（VQ-QAT）下的表现尚不明确。本文提出SQuAT，一种形状增益VQ-QAT方法，采用独立的直通估计器（STE）来处理形状和增益。SQuAT将权重分解为方向和幅度：归一化方向使用基于$E_8$的球形码本进行量化，而增益则使用低比特标量量化器独立量化。研究表明，与应用于完整量化器的标准STE不同，SQuAT的独立STE为方向引入切空间更新，形成隐式的球面（黎曼）优化。通过开发分布式精确最近邻分配方案，SQuAT将VQ-QAT训练时间缩短了12.8倍，并扩展到100亿个标记。实验结果显示，SQuAT在多个模型上均优于标准VQ-QAT STE，达到了最先进的VQ-QAT准确率。

🔬 方法详解

问题定义：本文旨在解决现有量化感知训练方法在大语言模型中的不足，特别是在量化效果和训练效率方面的挑战。现有方法往往忽视了权重的形状和增益结构，导致性能下降。

核心思路：SQuAT的核心思路是将权重分解为方向和幅度，分别进行量化，从而更好地利用量化结构。通过独立的STE处理方向和增益，SQuAT实现了更有效的切空间更新，进而优化了量化过程。

技术框架：SQuAT的整体架构包括权重分解模块、独立STE模块和分布式最近邻分配方案。权重首先被分解为方向和幅度，然后分别进行量化，最后通过分布式方案进行训练，以提高效率。

关键创新：SQuAT的主要创新在于引入了独立的STE来处理权重的形状和增益，形成了一种隐式的球面优化方法。这一设计与传统的STE方法显著不同，后者通常对量化器的结构缺乏考虑。

关键设计：在参数设置上，SQuAT使用基于$E_8$的球形码本进行方向量化，增益则使用低比特标量量化器。损失函数设计上，SQuAT确保了方向和增益的独立性，以优化训练效果。

📊 实验亮点

在多个大型模型（如LLaMA-3.2-1B/3B和Qwen3-4B）上，SQuAT在相同的量化器和初始化条件下，均显著优于标准VQ-QAT STE，且训练时间减少了12.8倍，达到了最先进的VQ-QAT准确率。

🎯 应用场景

SQuAT的研究成果在大语言模型的训练和部署中具有广泛的应用潜力，尤其是在需要高效量化和低比特表示的场景中，如移动设备和边缘计算。其优化的训练过程和高准确性将推动自然语言处理和生成任务的进步。

📄 摘要（原文）

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.