🤖 machine learning

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

This paper introduces BCJR-QAT, a differentiable relaxation of trellis-coded weight quantization that replaces the non-differentiable Viterbi argmax with a temperature-controlled BCJR sum-product algorithm, enabling end-to-end quantization-aware training that empirically outperforms the state-of-the-art PTQ frontier on LLMs.

Original authors: Venugopalan Iyengar

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Venugopalan Iyengar

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a massive library of books (a Large Language Model) that you want to shrink down so it fits on a standard laptop or phone. To do this, you need to compress the "words" (weights) inside the library.

The paper introduces a new method called BCJR-QAT to compress these books even further—down to just 2 bits per word (which is like compressing a high-definition photo into a tiny, grainy thumbnail).

Here is the story of how they did it, using simple analogies:

1. The Problem: The "One-Way Street" Trap

Previously, the best way to shrink these models was a method called QTIP. Think of QTIP as a very smart librarian who looks at a page of text and says, "Okay, I'll replace this complex sentence with the closest simple phrase from my dictionary."

However, there was a limit. Once the librarian made that choice, they couldn't go back and change their mind to make the whole book sound better. It was like a "one-way street." If the librarian picked a slightly wrong phrase early on, the rest of the book suffered, and they couldn't fix it because the decision was "hard" and final.

To fix this, researchers usually use Training (QAT), where you let the librarian practice and adjust their choices. But here's the catch: The librarian's decision process involves a complex maze (a "trellis"). To find the best path through the maze, they use a rule called Viterbi, which is like a "pick the absolute best option" button.

The Problem: You can't teach a computer to learn through a "pick the best" button because that button is a mathematical dead end. If you push the button, the computer doesn't know how to nudge its choices to get better; it just snaps to the nearest option. It's like trying to steer a car by only allowing it to jump instantly to the next lane—you can't make smooth turns.

2. The Solution: The "Soft" Decision (BCJR)

The authors, Venugopalan Iyengar, invented a new way to let the librarian practice. Instead of forcing a hard "pick the best" decision immediately, they introduced a concept called Temperature ( $T$ ).

High Temperature (Hot): Imagine the librarian is very relaxed and indecisive. Instead of picking just one phrase, they consider many phrases at once, assigning a "probability" to each. It's like a foggy day where you can see several paths, not just one. This "soft" view is mathematically smooth, meaning the computer can easily calculate how to nudge the choices to improve the book.
Low Temperature (Cold): As the training finishes, the librarian gets "cooler" and more decisive. The fog clears, and they snap to the single best phrase (the hard decision).

This method, called BCJR, turns the "one-way street" into a smooth, drivable road. The computer can now learn exactly how to adjust the weights to make the whole book sound better, not just individual sentences.

3. The "Overshoot" Mistake

The authors discovered a surprising quirk in how this "temperature" works.

In traditional physics, when you try to find the best solution, you usually start very hot (very indecisive) to explore all possibilities, and then slowly cool down.

The Paper's Finding: For this specific type of compression, starting "too hot" is a disaster. If the librarian is too indecisive at the start, they wander off into a bad neighborhood (a "worse Voronoi basin") and get stuck there. By the time they cool down and try to find the best path, they are already lost in a bad area and can't get back.
The Fix: They found that starting at a moderate temperature (not too hot, not too cold) works best. It's like starting a hike with a clear map rather than wandering blindly in the fog. By skipping the "super hot" phase, they avoided getting lost and found a better solution.

4. The Results: A Win for Small Computers

The team tested this on two different types of models:

The "Proxy" Test (OLMoE): They tried to optimize the compression by just looking at how well each layer of the model reconstructed its own data (like checking if a photocopy looks like the original). Result: It failed. The model actually got worse than the old method. This taught them that just making a "better photocopy" doesn't mean the book reads better.
The "Real" Test (Llama-3.2): They optimized the model to actually read and understand text better (using a method called "distillation," where a smart teacher model guides the student). Result: It worked!
- On a specific layer of the model, their new method beat the old best method by a small but significant margin.
- When they applied this to multiple layers, the improvements added up in a "super-additive" way (the whole became greater than the sum of its parts).

5. The Engine Under the Hood

Doing this math is usually very slow and requires expensive supercomputers. The authors also built a specialized "engine" (a Triton kernel) that runs this complex math incredibly fast on a single consumer graphics card (like an RTX 4080). They made it 6.5 times faster than standard methods, proving that this advanced technique doesn't need a data center to run.

Summary

The paper presents a new way to shrink AI models to fit on consumer devices.

The Trick: They replaced a rigid, un-teachable decision process with a "soft," learnable one that gradually hardens.
The Insight: You shouldn't start the learning process too "hot" (too random); a moderate start prevents the model from getting lost.
The Outcome: When trained correctly to understand text (not just copy data), this method produces better results than previous state-of-the-art techniques, and it runs fast enough on a regular gaming PC.

The authors have released their code and the trained models so others can try it out.

Technical Summary: BCJR-QAT

Problem Statement

Achieving 2 bits per weight (bpw) or lower for Large Language Models (LLMs) on consumer hardware is a critical bottleneck. While Post-Training Quantization (PTQ) methods like QTIP have established a strong frontier for 2-bit quantization using Trellis-Coded Quantization (TCQ), they are limited by the "PTQ ceiling." Pushing beyond this requires Quantization-Aware Training (QAT). However, standard QAT techniques fail on trellis codes because the Viterbi decoding algorithm involves a non-differentiable argmax operation over exponentially many paths. Applying a Straight-Through Estimator (STE) to the trellis destroys the coupling between adjacent weight decisions, which is the primary advantage of trellis coding over scalar quantization.

Methodology: BCJR-QAT

The paper introduces BCJR-QAT, a differentiable relaxation of trellis-coded quantization that replaces the hard Viterbi argmax with a soft expectation derived from the BCJR (Bahl-Cocke-Jelinek-Raviv) forward-backward algorithm.

Core Mechanism

Statistical Mechanics Analogy: The trellis is treated as a factor graph equivalent to a 1D Ising-like spin chain. The Viterbi argmax is viewed as the zero-temperature limit ( $T \to 0$ ) of a Boltzmann distribution over paths.
Soft Codeword: At a finite temperature $T > 0$ , the method computes the exact marginal posterior probabilities of each state using the BCJR algorithm. The quantized output is the Boltzmann expectation of the emitted codeword:
$\hat{w}_t(T) = \langle c(s_t) \rangle_{p_T} = \sum_s c(s) p_T(s_t = s | w)$
This soft codeword is exactly differentiable with respect to the weight pre-image and trellis parameters.
Annealing: During training, the temperature $T$ is annealed from an initial value $T_0$ to a near-zero $T_{end}$ . As $T \to 0$ , the soft codeword converges to the hard Viterbi output, effectively "crystallizing" the continuous optimization into a discrete solution.

Key Technical Innovations

Differentiability: The BCJR forward-backward recursion (complexity $O(L \cdot S^2)$ ) provides exact gradients through the trellis, preserving the inter-state coupling that scalar STE methods lose.
Triton Kernel: The authors implemented a fused Triton kernel that combines gather, log-sum-exp, and addition operations into a single launch. This reduces the forward/backward pass from hundreds of kernel launches to one, achieving a 6.57× speedup over a reference PyTorch autograd implementation while maintaining FP32 numerical parity.
Schedule Correction (Skip-High-T): The paper diagnoses a critical failure mode in standard simulated annealing schedules (starting at $T_0=1.0$ ). At high temperatures, the soft codeword gradient is "smeared" and uninformative, causing the optimizer to drift into a worse Voronoi basin before the cooling phase can recover. The proposed fix is to skip the high- $T$ phase, initializing at a lower temperature (e.g., $T_0=0.3$ ) where gradients are already informative.

Experimental Results

The paper evaluates BCJR-QAT on two architectures: OLMoE-1B-7B (MoE) and Llama-3.2-1B (Dense).

1. OLMoE-1B-7B (Negative Result: Proxy Gap)

Setup: Per-layer greedy QAT optimizing for reconstruction MSE (a proxy objective).
Outcome: BCJR-QAT successfully reduced per-layer reconstruction MSE by 17–20% compared to QTIP-PTQ. However, this improvement did not translate to end-task performance. The models regressed by $\sim 1.3$ perplexity (PPL) on WikiText-2 compared to the QTIP-PTQ baseline.
Conclusion: This confirms a "proxy gap" where minimizing local reconstruction error does not guarantee better end-task performance at 2 bpw.

2. Llama-3.2-1B (Positive Result: End-to-End KL)

Setup: Single-layer QAT optimizing for end-to-end forward-KL distillation against an FP16 teacher.
Outcome:
- Schedule Overshoot: Using a naive high- $T$ schedule ( $T_0=1.0$ ) resulted in performance slightly worse than PTQ (+0.005 to +0.022 PPL).
- Skip-High-T Fix: Using the corrected schedule ( $T_0=0.3$ ) eliminated the overshoot. On Layer 4, the model achieved 10.1347 PPL, a −0.084 PPL improvement over the QTIP-PTQ baseline (10.2189).
- Super-Additivity: When combining trained layers (Layer 4 and Layer 8), the joint gain (−0.077 PPL) exceeded the sum of individual gains (−0.062 PPL), demonstrating a "cooperation surplus" where BCJR-trained codewords reinforce each other.

3. Theoretical Analysis

Drift-Budget Theory: The authors derive a condition for escaping the PTQ Voronoi basin: the cumulative weight drift ( $\eta \cdot N \cdot g_{max}$ ) must exceed the basin radius ( $r_{Voronoi} \approx 10^{-3}$ ). Experiments confirm that schedules failing this condition cannot escape the PTQ solution, while those exceeding it can find better solutions if the objective and schedule are correct.

Significance and Claims

The paper claims three primary contributions:

Algorithmic: BCJR-QAT is the first method to make trellis-coded quantization natively differentiable, enabling QAT on TCQ without losing the benefits of path coupling.
Engineering: The fused Triton kernel makes this computationally feasible on consumer hardware (single 12GB GPU), removing the inference-time deployment barrier for research.
Empirical Insight:
- It identifies that per-layer reconstruction MSE is a poor proxy for 2-bit QAT, whereas end-to-end KL distillation is effective.
- It challenges the conventional wisdom of simulated annealing for discrete relaxations, showing that starting "hot" ( $T=1.0$ ) is detrimental for BCJR-QAT because the high-temperature gradient is noise-dominated. The optimal strategy is to start at a temperature where the gradient is already informative.

The paper concludes that while single-layer results are promising, the full potential of BCJR-QAT likely lies in joint end-to-end training across all layers, a task made affordable by the released kernel but requiring cloud-scale hardware (H100/H200) for models larger than 1B parameters. The authors explicitly state they did not perform this full-scale joint training in this work.