论文
BCJR-QAT:网格编码权重量化的可微松弛
BCJR-QAT: A Differentiable Relaxation of Trellis-Coded Weight Quantization
摘要
网格编码量化为 LLM (QTIP) 设置当前 2 位 后训练 前沿,但低于 PTQ 上限需要量化感知训练,并且网格上的 QAT 受到不可微分维特比 argmax 的阻碍。我们引入了 BCJR-QAT,这是一种在温度 $T$ 下用 BCJR 前向-后向和积算法取代 argmax 的松弛方法,产生一个软码字,该软码字等于网格路径上的玻尔兹曼期望,完全可微,将硬 QTIP 代码恢复为 $T \to 0$,并且在数学上与 1D Ising 类自旋链的传递矩阵计算相同。我们贡献 (i) 一个融合的 Triton 内核,使 BCJR 在单个消费者 GPU 上易于处理(6.57 美元\times$ 加速,fp32 奇偶校验); (ii) BCJR-QAT 何时能够逃离 QTIP-PTQ Voronoi 盆地的定量漂移预算理论,已通过四次实验得到验证; (iii) Llama-3.2-1B 在端到端正向 KL 蒸馏 下以 2 bpw 获得积极的实证结果:采用正确的时间表(跳过高 $T$ 阶段以避免我们诊断的过冲),单层 BCJR-QAT 在 WikiText-2 上以 $\mathbf{-0.084}$ PPL 击败 QTIP-PTQ,并且多层复合为超添加剂。