用可微方法突破2比特大模型量化瓶颈，实现更优压缩效果。
BCJR-QAT: A Differentiable Relaxation of Trellis-Coded Weight Quantization
- 引入BCJR-QAT替代非可微的Viterbi算法，实现梯度传播。
- 在Llama-3.2-1B上以2比特每权实现-0.084的困惑度提升。
- 适合追求极致模型压缩与推理效率的研究者使用。
格码量化目前是大模型2比特后训练量化的前沿（QTIP），但要突破后训练量化极限需依赖量化感知训练（QAT）。然而，传统格码结构中的Viterbi argmax操作不可微，阻碍了QAT应用。本文提出BCJR-QAT，将argmax替换为温度T控制的BCJR前向-后向求和算法，生成玻尔兹曼加权路径的软码字，完全可微，且在T→0时退化为硬码字，数学形式等价于一维伊辛自旋链的转移矩阵计算。我们贡献：(i) 一个融合的Triton内核，使单张消费级显卡即可高效运行BCJR，速度提升6.57倍，精度接近fp32；(ii) 建立定量漂移预算理论，预测何时能跳出QTIP-PTQ的Voronoi区域，并在四个实验中验证；(iii) 在Llama-3.2-1B上，采用端到端前向KL蒸馏，在2比特每权条件下，通过合理调度（跳过高温阶段避免过冲），单层BCJR-QAT在WikiText-2上达到-0.084的困惑度降低，多层叠加呈现超加性增益。
原文摘要 · Abstract (English)
Trellis-coded quantization sets the current 2-bit post-training frontier for LLMs (QTIP), but pushing below the PTQ ceiling requires quantization-aware training, and QAT on a trellis is obstructed by the non-differentiable Viterbi argmax. We introduce BCJR-QAT, a relaxation that replaces the argmax with the BCJR forward-backward sum-product algorithm at temperature $T$, producing a soft codeword equal to the Boltzmann expectation over trellis paths, exactly differentiable, recovering the hard QTIP code as $T \to 0$, and mathematically identical to the transfer-matrix computation for a 1D Ising-like spin chain. We contribute (i) a fused Triton kernel making BCJR tractable on a single consumer GPU ($6.57\times$ speedup, fp32 parity); (ii) a quantitative drift-budget theory of when BCJR-QAT can escape the QTIP-PTQ Voronoi basin, verified across four experiments; and (iii) a positive empirical result on Llama-3.2-1B at 2 bpw under end-to-end forward-KL distillation: with the right schedule (skip the high-$T$ phase to avoid an overshoot we diagnose), single-layer BCJR-QAT beats QTIP-PTQ by $\mathbf{-0.084}$ PPL on WikiText-2, and multi-layer compounding is super-additive.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。