LC-QAT: Data-Efficient 2-Bit QAT for LLMs via Linear-Constrained Vector Quantization

Abstract: Quantization-aware training (QAT) is essential for extremely low-bit LLMs. Current QAT methods are mainly based on scalar quantization (SQ), which enables efficient optimization but suffers from severe performance degradation at 2-bit precision. On the other hand, vector quantization (VQ) provides substantially higher representational capacity, but its discrete codebook lookup prevents end-to-end training. We propose LC-QAT, a 2-bit weight-only VQ-QAT framework that represents quantized weights via a learned affine mapping over discrete vectors, which yields a high-quality PTQ initialization and enables fully differentiable end-to-end optimization without explicit codebook lookup in the training forward pass. This strong post-training initialization makes LC-QAT highly data-efficient. Experiments across diverse LLMs demonstrate that LC-QAT consistently outperforms state-of-the-art QAT methods while using only 0.1%--10% of the training data. Our results establish LC-QAT as a practical and scalable solution for extreme low-bit model deployment.

• The paper introduces LC-QAT, which uses linearly constrained vector quantization with affine codebooks, LDLQ/Hadamard initialization, and differentiable gradient estimation to train 2-bit LLM weights end to end.
• LC-QAT achieves 72.18 average zero-shot accuracy on Qwen-3-8B and 73.42 on LLaMA-3-8B with only 4B training tokens, outperforming prior methods that use substantially more data.
• The lookup-free design supports optimized integer-weight inference, reaching 101.65 tokens per second on LLaMA-3-8B—1.68× faster than FP16—while showing that initialization quality is central to data-efficient 2-bit QAT.

Motivation and problem setting

At 2-bit weight-only precision, quantization-aware training (QAT) consistently outperforms post-training quantization (PTQ), but the dominant QAT paradigm is built on scalar quantization (SQ), which independently maps each weight to one of four values. SQ-based QAT is easy to optimize yet loses substantial information at ultra-low bit-widths, producing weak initializations that require very large training corpora to recover. Vector quantization (VQ), which represents groups of $d$ weights with codewords from a shared codebook ($K = 4^d$ entries at 2 bits), offers far higher representational capacity, but its discrete nearest-neighbor assignment and codebook lookup are non-differentiable, blocking end-to-end training. Prior VQ-QAT attempts, notably PV-Tuning, rely on coordinate descent or beam search with asynchronous parameter updates.

Linear-constrained codebook parameterization

LC-QAT's central idea is to replace unconstrained codebooks with a linearly constrained parameterization: each codeword is generated as $c = Az + B$, where $z \in \{0,1,2,3\}^d$ is a discrete quaternary vector and $A \in \mathbb{R}^{d\times d}$, $B \in \mathbb{R}^d$ are floating-point parameters shared per weight matrix. This reformulation converts discrete index selection into rounding and clamping followed by a linear projection, so gradients propagate through the dequantization without any index search. The authors distinguish this from lattice-based codebooks in Quip# and NestQuant, whose padding or nested codewords cannot be generated by a single shared affine map; those methods retain more flexibility but require lookup.

Initialization proceeds by extending a PTQ method (LDLQ from QuIP#, with Hadamard incoherence processing) to produce an initial quantized model; $A$ and $B$ are initialized as $A = sG$ and $B = -sG\mu\mathbf{1}$ with random orthogonal $K = 4^d$0, yielding a codebook with approximately zero mean, unit variance, and weak inter-dimensional correlation. Only 20 extra parameters per matrix are introduced at $K = 4^d$1.

The paper supports this design empirically: projecting the loss landscape of Qwen-3-1.7B onto two directions shows the LC-QAT initialization sits near the FP16 minimum with a saddle-point structure resembling the full-precision model, whereas SQ-based initialization lies far from any local minimum. On LLaMA-3 models, the LC-PTQ initialization preserves 84–87% of FP16 zero-shot QA accuracy while OSTQuant—a strong SQ method using the same Hadamard transform—retains only about 58–62%, with ARC-Challenge collapsing nearly to random guessing. Since both methods share the rotation, the gap isolates VQ versus SQ as the causal factor.

Lookup-free forward pass and differentiable gradients

During training, LC-QAT maintains floating-point proxy weights $K = 4^d$2 discretized into integer weights $K = 4^d$3, which are decoded group-wise through the affine map. The output computation is algebraically reorganized as $K = 4^d$4, so activations absorb the linear mapping and the integer weights remain compatible with optimized Int2-FP16 MatMul kernels designed for SQ models. With a custom fused CUDA kernel, LC-QAT reaches 101.65 tokens/s on LLaMA-3-8B (single A100, bs=1)—a 1.68× speedup over FP16 and 1.43× over AQLM—and throughput is essentially invariant to group size.

Two auxiliary components address training stability. First, because LC-QAT starts close to a local minimum, the Straight-Through Estimator's gradient noise is disproportionately harmful; the authors adopt a Differentiable Gradient Estimator (DGE) that smooths the discretization with a controllable sharpness parameter ($K = 4^d$5). Second, since $K = 4^d$6 originates from PTQ rather than from quantizing pretrained weights, proxy weights are initialized via an Xavier-aligned affine transformation of $K = 4^d$7. Ablations show both components are necessary: without preprocessing the loss stays flat, and STE produces large unrecoverable loss spikes that DGE eliminates.

Results

Model capacity. Under equal data budgets on Qwen-3-8B, LC-QAT achieves WikiText-2 perplexity of 10.23 and 72.18 average zero-shot accuracy (97% of FP16), exceeding PV-Tuning (10.65 PPL, 70.99) and all PTQ baselines including QTIP and YAQA. On the harder Qwen-3-1.7B, LC-QAT reaches 13.44 PPL and 61.43 average accuracy versus PV-Tuning's 17.43/58.05, surpassing even state-of-the-art trellis PTQ. Training curves show PV-Tuning saturates early—attributed to its asynchronous updates—while LC-QAT improves monotonically.

Data efficiency. On LLaMA-3-8B, LC-QAT trained on 4B tokens attains 73.42 average zero-shot accuracy (9.38 PPL), exceeding ParetoQ's 72.07 despite ParetoQ using 30B tokens—roughly 10% of the data budget. On LLaMA-3-3B it reaches 68.33 average (98% of FP16) versus ParetoQ's 66.82. In the instruction-tuned setting, LC-QAT applied to a 1.7B model averages 50.35 across OpenBookQA, IFEval, MMLU, GSM8K, MATH, and HumanEval, beating BitNet b1.58 2B4T (48.07) while using roughly 0.1% of BitNet's 4T-token pretraining budget. Total wall-clock cost is also favorable: 61 hours for LC-QAT (including 6 hours of PTQ initialization) versus 417 hours for ParetoQ on comparable hardware.

Ablations. Replacing the LDLQ initialization with GPTQ or random initialization collapses zero-shot accuracy from 67.95 to 37.53–39.51, confirming that high-quality PTQ initialization—not merely end-to-end training—is essential. Conversely, running PV-Tuning's optimizer from the same LC-PTQ start yields only 55.35 average versus LC-QAT's 59.75 on Qwen-3-1.7B, showing end-to-end synchronous optimization contributes beyond initialization quality alone. Scaling analyses show continued improvement from 1B to 10B tokens without saturation, consistent gains on a 14B model with only 58M tokens, and insensitivity to group size ($K = 4^d$8 vs. $K = 4^d$9).

Limitations and open questions

The authors acknowledge that experiments focus on data-efficient fine-tuning under limited budgets, leaving scalability to large-scale training regimes open. The additional linear projection in the forward pass has bounded asymptotic complexity, but its interaction with highly optimized GEMM kernels in production frameworks has not been systematically studied, and its effect on training efficiency under different deployment conditions remains unexplored. Two further caveats bear on interpretation: the data-scaling comparison against ParetoQ is asymmetric because ParetoQ's training data is not publicly available for a symmetric analysis, and the instruction-tuned comparison with BitNet contrasts fine-tuning of a pretrained dense model against training from scratch—an apples-to-oranges setting the paper presents as evidence of data efficiency rather than of equivalent training paradigms.

Conclusion

LC-QAT demonstrates that constraining VQ codebooks to the image of a shared affine map makes vector-quantized LLM weights fully trainable with standard backpropagation, eliminating codebook lookup from both training and inference. Combined with LDLQ/Hadamard PTQ initialization, DGE gradient estimation, and integer-weight preprocessing, the framework outperforms prior VQ-QAT under equal data and matches or exceeds state-of-the-art SQ-QAT with 0.1%–10% of the training data across Qwen-3, LLaMA-3, and 14B-scale models. The results position initialization quality, rather than training volume, as the dominant lever for practical 2-bit QAT, while leaving large-budget scaling and kernel-level systems integration as open questions.