论文
JacQuant:通过学习雅可比代理进行无 STE 量化感知训练
JacQuant: STE-Free Quantization-Aware Training via Learned Jacobian Surrogates
摘要
量化感知训练 (QAT) 得到了广泛部署,但通常依赖于直通估计器 (STE),该估计器通过不可微分的量化器强制传递梯度。这通常会使箱边界附近的训练变得脆弱,并且与低精度模型的实际行为的一致性较差。我们引入了 JacQuant,这是一个 QAT 框架,它学习模型对参数变化的局部敏感性的轻量级代理,并使用它来稳定和加速标准方差减少优化器中的训练。该替代物价格低廉(对角线或块对角线),数据驱动,并且与常见的权重和激活量化器兼容。在代码保留训练阶段,我们证明了非凸目标的收敛性并在 PL 条件下获得线性速率,并且我们通过简单的校准参数将学习到的灵敏度与端到端输出保真度联系起来。在 $\leq 2$ 位的 LLM 基准测试中,JacQuant 始终达到比基于 STE 的 QAT 更高的精度,并且对各种模型的运行时分析表明,在实际组大小下,增加的成本仍然可以忽略不计。该方法是直接引入的,不需要改变前向量化器;我们的经验主张范围为超低位 LLM QAT。