BCJR-QAT: A Differentiable Relaxation of Trellis-Coded Weight Quantization

arXiv:2605.10655v1 Announce Type: new Abstract: 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 o 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.57imes 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.

Source: arXiv cs.LG | 2026-05-12