ThinkBit: Search Codewords, Recalibrate Scales, Select by Reasoning

AI summary

Overview

Research area: Model compression for large language models, specifically ultra-low-bit (2-bit) weight quantization of reasoning models.

Technical level: Advanced. The paper assumes familiarity with transformer internals, GPTQ-style post-training quantization, knowledge distillation, quantized-aware training, and techniques such as Kronecker-factored Hessian approximations, Gumbel-softmax relaxations and ReinMax gradient estimators.

Scope: The paper proposes ThinkBit, a five-stage strictly scalar two-bit quantization framework for Qwen3 models (1.7B, 8B, 14B), and argues that such models should be selected and evaluated on generated reasoning accuracy rather than perplexity or multiple-choice accuracy.

What This Paper Is About

Reasoning models emit long chains of thought, and because each generated token requires reading all weights from memory, inference speed is bounded by memory bandwidth; storing weights in two bits instead of bfloat16 can make inference up to 8× cheaper. The problem is that two-bit quantization usually destroys a model's ability to reason through hard problems, while the standard yardsticks (perplexity, multiple-choice accuracy) fail to reveal that damage. The paper builds a two-bit framework that keeps the deployment format as simple as possible — one integer codeword per weight times one bf16 scale per group of 128 weights, with no vector codebooks, rotations or zero-points — and shows it preserves reasoning accuracy per byte better than published two-bit alternatives.

Key Contributions

1. An argument for reasoning-based model selection. The authors show that removing the codeword search, the scale recalibration, both, or the bit-allocation stage leaves the 7-task QA average essentially flat at 53.7–54.8 (within a ±1.5-point margin), while sampled MATH-500 accuracy after QAT falls from 43.0 to the range 28.4–37.0. They conclude two-bit checkpoints must be chosen on reasoning tasks.

2. The ThinkBit framework itself, a five-stage pipeline: (1) Allocate — sensitivity-based assignment of a ternary, 2-bit or 3-bit tier per layer; (2) Search — block-by-block search for integer codewords starting from a layer-wise GPTQ solution; (3) Recalibrate — fit only the stored scales and RMSNorm gains with codewords frozen; (4) QAT — reassign codewords via straight-through latents; (5) Selection — keep whichever of the recalibrated or QAT checkpoint wins a paired reasoning test.

3. Better two-bit reasoning per byte. On Qwen3-1.7B at matched size (about 0.5 GB), the recalibrated checkpoint (Rec-h2) scores 46.7 on MATH-500 versus 30.0 for EfficientQAT trained on the same data and tokens. At 0.52 GB — 2.42 bits per parameter counting head and embeddings — the checkpoint is smaller than every published 2-bit GGUF of the model while scoring higher than all of them on MATH-500.

4. A precise accounting of what "2-bit" means. The paper documents that nominal two-bit releases actually store far more: Unsloth's 2-bit GGUFs store 2.82–3.62 bits per parameter for Qwen3-1.7B and 2.42–3.12 for Qwen3-14B, and PrismML's ternary Bonsai stores 2.15 bits per parameter for Qwen3-1.7B despite weights needing only log₂3 ≈ 1.585 bits. Comparing by nominal bit-width can hide size differences of up to 50%, so the authors compare whole-model size in bytes.

Main Findings

• Recalibration recovers most of what quantization loses, and it is where most of the gain lives. On Qwen3-1.7B (bf16 head and embedding), the codeword-search-only checkpoint (GSQ-hbf16) sits at 46.91 WikiText-2 perplexity, 43.8 on the 7-task mean and 14.8 on MATH-500. Recalibration (Rec-hbf16) brings this to 19.46 perplexity, 51.5 on the 7-task mean and 47.5 on MATH-500, against the bf16 model's 16.67, 58.9 and 91.2.

• QAT improves likelihood but not reasoning on the base mixture. The QAT checkpoint (QAT-hbf16) reaches 16.54 WikiText-2 perplexity and 54.6 on the 7-task mean — better than the recalibrated checkpoint on both — yet scores lower on MATH-500 (43.0 versus 47.5). This is the paper's central evidence that likelihood metrics and reasoning accuracy diverge.

• Standard benchmarks can mislead checkpoint selection. In the stage ablations summarized in Figure 1b, variants that drop codeword search, recalibration or both keep the 7-task mean in the 53.7–54.8 range while sampled MATH-500 drops from 43 to 35.0, 37.0 and 32.1; switching from mixed tiers to uniform 2-bit allocation gives 28.4. Meanwhile Rec-hbf16 is lower on the 7-task mean (51.5) but higher on reasoning (47.5).

• Codeword search matters a great deal at initialization. The framework reaches 47.5 on sampled MATH-500 before QAT and 43.0 after, against 18.5–24.1 and 32.6–37.8 from tuned GPTQ initializations.

• Reasoning-trace QAT lifts the small model considerably. On Qwen3-1.7B the score rises to 73.0 after QAT on reasoning traces, in a 0.52 GB checkpoint.

• The framework scales to 14B. At 3.74 GB a model scores 94.8 on MATH-500, against 97.5 for the 29.5 GB bf16 model.

• The bit allocator never picks 2 bits on its own. In the (bits, rounding-noise) plane the 2-bit point lies above the line joining the ternary and 3-bit points, so for every price λ the allocator prefers ternary or 3-bit. At a budget of 2.0 decoder bits per weight on Qwen3-1.7B the outcome is 75.3% ternary, 6.0% 2-bit and 18.7% 3-bit weights; only query and key projections are held at 2 bits or above, and an appendix shows this restriction is unnecessary on that model.

• Measured rounding noise matches the theoretical grids. Ternary, 2-bit and 3-bit tiers show noise of 0.190, 0.150 and 0.041 when measured with the actual scale search, versus exact-scale minima of 0.187, 0.146 and 0.038 — a difference the authors say changes none of their conclusions.

Methodology in Plain English

Every linear weight is stored as a small integer times one shared bf16 scale for each block of 128 consecutive input weights, so dequantization is a single multiplication. This costs 1.71 bits per weight for ternary, 2.125 for 2-bit and 3.125 for 3-bit, including the scale's share.

Stage 1 (Allocate). Using KronQ, the authors approximate each layer's Hessian by a Kronecker product of an input factor and an output-gradient factor, and reduce the sensitivity of a layer to the product of the mean squared input norm and the mean squared output-gradient norm. They multiply sensitivity by a per-tier rounding-noise constant and solve a multiple-choice knapsack problem by Lagrangian relaxation with bisection on the bit price, then spend leftover bits on the upgrades with the largest loss reduction per bit. Sensitive, large layers get more bits.

Stage 2 (Search). Each 2- and 3-bit layer is first quantized alone with GPTQ (with a damping term modulated by output-channel energy). Then a GSQ-style block-wise search treats each codeword as a distribution over candidate grid values and optimizes logits and scales jointly against the bf16 block's output, weighted by gradient second moments at each residual-stream channel. 3-bit and ternary layers use a Gumbel-softmax relaxation with annealed temperature; 2-bit layers use ReinMax, which samples a hard grid value in the forward pass and uses a second-order-corrected gradient estimate, so training and inference see the same kind of weight.

Stage 3 (Recalibrate). Codewords are frozen; each group scale is reparameterized as s = s₀e^r and each RMSNorm gain as ν = ν₀(1+ψ), with bounded offsets. Only r and ψ are trained — under 1% of the model's parameters — minimizing forward KL to the bf16 model's next-token distribution plus penalties that act as a trust region around the Stage-2 solution. The step with the lowest held-out KL is kept.

Stage 4 (QAT). Each weight gets an FP32 latent measured in codeword units; the forward pass rounds and clips it to the layer's grid, and the backward pass uses a straight-through estimator. Latents start partway toward the bf16 weight (α = 0.5), so no codeword changes at step 0 but each latent sits on the side where the bf16 weight lies, letting training move it across a rounding boundary when useful. The loss adds hidden-state distillation on selected blocks to the KL term.

Stage 5 (Selection). Two candidate checkpoints, recalibrated and QAT, are compared on 313 held-out level-3–5 MATH training problems disjoint from MATH-500, with four sampled answers per problem and strict grading (the trace must close and the final answer must be correct). A paired sign-flip test estimated from 20,000 random sign vectors decides: keep the higher-accuracy model if p ≤ 0.05, otherwise the smaller one, or the recalibrated one when sizes tie.

Setup. Models studied are Qwen3-1.7B, Qwen3-8B and Qwen3-14B with identical hyperparameters across all three. Training data mixes each model's own bf16 replies (code, math, tool calls, QA, MMLU-auxiliary), its own continuations of FineWeb-Edu and Wikipedia, and raw text from a range of corpora, with replies capped at 4,096 tokens. Stage budgets are 2M tokens (sensitivity), 57M (codeword search), 268M (recalibration), 537M (QAT), and 1.1–2.1B tokens for the longer QAT on a mixture where complete long reasoning traces from OpenMathReasoning, OpenCodeReasoning, DeepSeek-R1 solutions and code/math sources make up about 60%. The selected checkpoints at 1.7B run 2.1B tokens at learning rate 2·10⁻³, and at 8B/14B run 1.1B tokens at 1.5·10⁻³. Evaluation uses lm-eval with vLLM; sampled MATH-500 uses temperature 0.6, top-p 0.95, top-k 20 and up to 32k new tokens.

Why This Matters

The paper reframes low-bit quantization around reasoning rather than likelihood, and shows empirically that the two can move in opposite directions within the same framework. It also introduces a rigorous size-accounting convention (whole-model bytes and bits per parameter, counting head and embeddings) that makes competing "2-bit" claims comparable — a point with direct consequences for how compression results are reported.

Real-world applications:

• On-device and edge deployment of reasoning assistants, where a 0.52 GB checkpoint of a 1.7B model fits memory budgets that a bf16 model does not.
• Cost reduction in serving long chain-of-thought workloads, where per-token weight traffic rather than arithmetic dominates latency and 2-bit storage cuts that traffic by up to 8×.
• Consumer-GPU and laptop inference of larger reasoning models, illustrated by a 3.74 GB 14B-class model scoring 94.8 on MATH-500 versus 97.5 for the 29.5 GB bf16 original.
• Model-selection and release QA pipelines for compressed models, where the paper's paired reasoning test offers a cheap, statistically grounded alternative to perplexity sweeps.

Industry relevance is direct: scalar formats dequantize with a single multiplication, so they drop into existing runtimes without custom kernels — the paper's stated reason for insisting on a strictly scalar format. The finding that likelihood-based benchmark suites can keep QA scores flat while reasoning collapses by more than ten points is a practical warning for anyone shipping quantized reasoning models.

Future Directions

• Replace the hand-set tier rules. The allocator never selects 2 bits on its own and needed an explicit rule holding query and key projections at 2 bits or above; the appendix reportedly shows that rule is unnecessary on Qwen3-1.7B, leaving open how allocation should behave on other architectures.
• Close the recalibration-versus-QAT gap. Recalibration gives better reasoning (47.5 vs 43.0 on Qwen3-1.7B at bf16 head precision) while QAT gives better likelihood, and the paper resolves this only by running both and selecting. A single training procedure that achieves both would remove a full stage of compute.
• Push the scalar format further down. The paper establishes a two-bit scalar baseline; whether the same allocate–search–recalibrate–select loop holds at even lower budgets, where the rounding-noise geometry becomes more punishing, is not addressed in the reported results.
• Extend the reasoning-based evaluation protocol. The paired sign-flip selection test uses 313 MATH training problems and sampled strict scoring; how many problems and samples are needed to make such tests reliable across model families and non-math reasoning domains remains an open methodological question.

Target Audience

Researchers and engineers working on LLM quantization, compression and efficient inference, especially those deploying reasoning models under tight memory budgets. It is also useful for evaluation specialists interested in why likelihood benchmarks miss reasoning degradation, and for practitioners who need to compare competing "2-bit" releases on a fair size basis. Readers without a background in second-order quantization methods, distillation and gradient estimators for discrete variables will find the methodology sections demanding.

Authors’ abstract

Reasoning language models generate thousands of tokens per answer, and generation speed is limited by memory bandwidth, so storing weights in two bits instead of bfloat16 can potentially make inference 8$\times$ cheaper. Two-bit methods are usually judged by perplexity and multiple-choice question-answering accuracy, both of which say little about whether a model can still reason through more difficult problems. We present ThinkBit, a strictly scalar two-bit framework built to preserve reasoning, in which every weight is an integer codeword times a bfloat16 scale shared by a group of weights. We first use a sensitivity metric to assign each layer a ternary, 2-bit, or 3-bit tier to allocate bitwidths optimally among layers. This is followed by a layerwise Gumbel-softmax-based scalar quantization algorithm that searches for optimal codewords block by block. Since a layerwise scheme does not capture how quantized layers interact with each other, we recalibrate scales, with the codewords frozen, using a knowledge-distillation-based loss against the teacher model. We then perform a short quantization-aware training (QAT) stage that lets the model mimic the base model's distribution by reassigning codewords. Finally, a paired reasoning test on training problems selects the final checkpoint between the recalibration stage and the QAT stage. On Qwen3-1.7B at matched size, our recalibrated checkpoint scores 46.7 on MATH-500 against 30.0 for EfficientQAT trained on the same data, and 73.0 after QAT on reasoning traces; at 0.52 GB it is smaller than every published 2-bit GGUF of the model yet more accurate on MATH-500. We show that this framework scales and up to 14B, a 3.74 GB model scores 94.8 on MATH-500, against 97.5 for the 29.5 GB bfloat16 model.