JacQuant:基于学习雅可比率近似的无梯度估计量感知训练
机器学习
2026-05-26 v1
摘要
感知训练 (QAT) 广泛部署,但通常依赖直通估计量 (STE),后者通过 fiat 将梯度通过不可导的量化器传递。这使得训练在二值阈值附近脆弱,且与低精度模型的实际行为较弱。我们引入 JacQuant,这是一种 QAT 框架,通过学习模型对参数变化的局部灵敏度的轻量化近似来稳定和加速训练。该近似成本低廉(对角或块对角),数据驱动,且与常见的权重和激活量化器兼容。在 code-preserving 训练阶段,我们证明了非凸目标的收敛性,并在 PL 条件下获得线性收敛率,我们将所学灵敏度通过一个简单的校准论证与端到端输出保真度联系。在 LLM 基准测试中使用 位时,JacQuant 在保持前向量化器不变的前提下,始终优于基于 STE 的 QAT 获得更高准确率;各种模型的运行时分析表明,所添加的开销在实际分组大小下仍可忽略不计。该方法即插即用,无需修改前向量化器;我们的实证主张限定于 ultra-low-bit LLM QAT。
引用
@article{arxiv.2605.25469,
title = {JacQuant: STE-Free Quantization-Aware Training via Learned Jacobian Surrogates},
author = {Kai Yi and Vignesh Vivekraja and Harshit Khaitan and Steven Li},
journal= {arXiv preprint arXiv:2605.25469},
year = {2026}
}