BCJR-QAT: A Differentiable Relaxation of Trellis-Coded Weight Quantization
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 consu...
| Authors | Venugopalan Iyengar |
|---|---|
| Primary category | cs.LG |
| Published | 11 May 2026 |
| Last updated at source | 11 May 2026 |
Silytics shows public metadata from arXiv and links to the original. Information may be incomplete or out of date. Last refreshed 17 October 2026. Licenses and terms of the original apply to the work itself. Not financial, legal or professional advice.