不用STE，用学习的雅可比代理实现更稳定的低比特大模型量化训练。
JacQuant: STE-Free Quantization-Aware Training via Learned Jacobian Surrogates
- 用数据驱动的轻量级雅可比代理替代STE，提升梯度稳定性。
- 2比特以下量化下，准确率显著高于传统STE方法。
- 无需修改量化器，兼容主流量化方案，适合超低比特大模型部署。
量化感知训练（QAT）广泛应用，但通常依赖直通估计器（STE），其通过强制传递梯度来处理不可微的量化操作，导致训练在分箱边界附近不稳定，且与低精度模型实际行为偏差较大。本文提出JacQuant，一种基于学习的局部敏感性代理的QAT框架，该代理为轻量级（对角或分块对角），可在标准方差缩减优化器中稳定并加速训练。理论证明，在代码保持训练阶段，非凸目标下可保证收敛，满足PL条件时达到线性收敛速率；并通过简单校准论证了学习到的敏感性与端到端输出保真度的关系。在≤2比特的大语言模型基准测试中，JacQuant consistently 达到更高准确率，运行时分析显示，合理分组规模下额外开销可忽略。方法即插即用，无需修改前向量化器，实验范围限定于超低比特大模型量化训练。
原文摘要 · Abstract (English)
Quantization-aware training (QAT) is widely deployed but typically relies on the Straight-Through Estimator (STE), which passes gradients through non-differentiable quantizers by fiat. This often makes training brittle near bin boundaries and weakly aligned with the actual behavior of the low-precision model. We introduce JacQuant, a QAT framework that learns a lightweight surrogate of the model's local sensitivity to parameter changes and uses it to stabilize and accelerate training within standard variance-reduced optimizers. The surrogate is inexpensive (diagonal or block-diagonal), data-driven, and compatible with common weight and activation quantizers. On code-preserving training phases, we prove convergence for non-convex objectives and obtain linear rates under a PL condition, and we relate the learned sensitivity to end-to-end output fidelity via a simple calibration argument. Across LLM benchmarks at $\leq 2$ bits, JacQuant consistently reaches higher accuracy than STE-based QAT, and the runtime analyses on various models show that the added cost remains negligible under practical group sizes. The method is drop-in and requires no changes to the forward quantizers; our empirical claims are scoped to ultra-low-bit LLM QAT.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。