Variable Bit-width Quantization: Learning Per-Group Precision for "Bigger-but-Smaller" Language Models

Abstract: Low-bit quantization shrinks LLMs but treats precision as a single global hyper-parameter: every weight uses the same bit-width. We introduce Variable Bit-width Quantization (VBQ), a training-time method in which each contiguous group of 64 weights learns its own resolution from {1,2,4,8} bits via a Gumbel-Softmax relaxation, trained jointly by an alternating optimization that gives the precision logits a clean, task-aligned signal. VBQ discovers a consistent, strongly heterogeneous allocation within individual projection types, not merely across layers, impossible to express with per-layer methods: 69% of groups collapse to 1 bit, the LM head averages 1.09 bits, while the first MLP block keeps ~2.5 bits. This pattern is stable enough to freeze into a fixed recipe and reuse without further search. The recipe yields a "bigger-but-smaller" regime: a 131M model at 1.82 mean bits reaches perplexity 4.2 on TinyStories, beating a 55M FP16 model (PPL 4.4) at 3.8x less storage, and lets a 1.46B model on FineWeb-Edu match a 593M FP16 control at ~3.7x less storage with 2.5x more parameters. As quality-per-byte, VBQ is 3.9-8.4x more efficient than FP16. The recipe maps directly to packed low-bit storage, so it also accelerates inference: with custom fused dequantize-and-multiply kernels, memory-bandwidth-bound autoregressive decode is faster at equal output, and the speedup grows with scale (parity at 131M, 1.9x at 1.0B, 4.7x at 9B on Apple silicon). A distributional analysis (KL divergence and argmax-flip rate) reveals a striking mechanism: deeper layers progressively self-heal the quantization error injected by early layers. The win is a from-scratch, train-time phenomenon; scaling the search economically beyond 1.5B parameters remains open. VBQ reframes precision as a learnable, non-uniform resource and shows that spending a fixed bit budget unevenly beats spending it uniformly.

• The paper introduces VBQ, a train-time quantization method that learns per-group precision using Gumbel-Softmax and alternating optimization.
• It demonstrates that adaptive bit allocation across transformer weights yields up to 8.4× efficiency improvements over FP16, with a 'bigger-but-smaller' regime.
• Empirical results show enhanced inference speed, storage reduction, and a self-healing error mechanism that preserves quality in deeper model layers.

Variable Bit-width Quantization for LLMs: Learning Heterogeneous Precision

Introduction

Variable Bit-width Quantization (VBQ) is introduced as a train-time quantization-aware technique that enables transformer-based LMs to allocate numerical precision — specifically bit-width — at a highly granular, per-group level. Departing from the conventional approach of uniform precision assignment across all weights or using fixed heuristics, VBQ leverages Gumbel-Softmax relaxation and alternating optimization to allow each group of 64 weights to learn its optimal precision from a discrete set $\{1,2,4,8\}$ bits. This results in models that are simultaneously larger (in parameter count), significantly smaller (in storage footprint), and considerably faster in inference, especially at scale.

The VBQ methodology fundamentally reframes precision during model training as a learnable resource that can be unevenly distributed, rather than a static hyper-parameter. Strong empirical evidence demonstrates that VBQ achieves a heterogeneous precision hierarchy impossible with per-layer or heuristic mixed-precision strategies, consistently allocating precision where it is most task-aligned.

Methodology

VBQ partitions each weight matrix into contiguous groups (64 weights per group). Each group is associated with a categorical precision variable, relaxed using Gumbel-Softmax for differentiability, and trained via alternating optimization — where weight updates and precision updates are interleaved and gradients are isolated to prevent cross-entropy overwhelming the bit penalty.

The training objective trades task loss (cross-entropy) against mean bit-width via a Lagrangian penalty, $$\mathcal{L} = \mathcal{L}_{\text{CE}} + \lambda\,\bar{b}$$, with $\lambda$ explicit chosen to target a mean precision (typically $\sim 1.8$ bits), below the empirically demonstrated information floor ($3.5$–$4$ bits per weight).

Precision allocation quickly converges to a strongly committed distribution—characterized by low entropy—and can be distilled into a transferable fixed recipe, eliminating the need for Gumbel machinery in production runs. The recipe specifies higher precision for the first MLP block and high-frequency vocabulary rows in the LM head, with lower precision elsewhere.

Heterogeneous Precision Allocation

VBQ discovers a substantially heterogeneous precision allocation both within and across projection types in transformers. Approximately $69\%$ of all weight groups select $1$-bit quantization, the LM head averages $1.09$ bits — with frequent tokens receiving higher precision — while the first MLP block retains $>2$ bits Figure 1.

Figure 1: VBQ learns a heterogeneous precision allocation; most groups select 1 bit, the LM head collapses to near-binary, and the MLP first block remains at higher precision.

This allocation is not only stable throughout training but semantically interpretable, with function words and punctuation dominating the high-precision assignments in the LM head. Importantly, the pattern emerges from optimization and is not manually set, highlighting the effectiveness of learning precision at per-group granularity.

Quality/Storage Pareto Frontier and "Bigger-but-Smaller" Regime

VBQ facilitates a "bigger-but-smaller" regime, where models with higher parameter count at lower average bit-width outperform smaller, full-precision baselines in validation perplexity, while consuming a fraction of the storage Figure 2. For example, a 131M-parameter VBQ model at $$\mathcal{L} = \mathcal{L}_{\text{CE}} + \lambda\,\bar{b}$$0 mean bits achieves PPL $$\mathcal{L} = \mathcal{L}_{\text{CE}} + \lambda\,\bar{b}$$1 on TinyStories, surpassing a 55M FP16 model (PPL $$\mathcal{L} = \mathcal{L}_{\text{CE}} + \lambda\,\bar{b}$$2) with $$\mathcal{L} = \mathcal{L}_{\text{CE}} + \lambda\,\bar{b}$$3 less storage.

Figure 2: Quality vs. storage frontier; VBQ models consistently achieve lower perplexity for a given storage budget compared to FP16 baselines.

Measured as quality-per-byte $$\mathcal{L} = \mathcal{L}_{\text{CE}} + \lambda\,\bar{b}$$4, VBQ models exhibit $$\mathcal{L} = \mathcal{L}_{\text{CE}} + \lambda\,\bar{b}$$5–$$\mathcal{L} = \mathcal{L}_{\text{CE}} + \lambda\,\bar{b}$$6 higher efficiency than FP16 counterparts Figure 3. The recipe generalizes well, requiring only a token-frequency table for the LM head and a projection-type map for deployment.

Figure 3: Quality-per-byte efficiency; VBQ models outperform dense FP16 by multiple times across scaling ladder.

Distributional Analysis: Depth Self-Healing of Quantization Error

Distributional audits (KL divergence and argmax-flip rate) against FP16 references reveal that quantization error injected by early, low-bit-width layers (e.g., the initial 4-bit MLP block) is progressively mitigated by deeper, subsequent 2-bit blocks Figure 4. This self-healing mechanism enables co-adaptation between blocks and heads, essentially allowing the network to learn an error-correcting paradigm through depth rather than uniform precision.

Figure 4: Cumulative KL divergence and argmax-flip rate across Transformer depth; quantization error is injected early but reduced in later blocks, demonstrating depth self-healing.

Notably, increasing block-level precision provides diminishing returns, as the co-adaptation absorbs available bits, making network width a more effective lever than depth alone for quality improvement.

Inference Efficiency: Scaling Wins

VBQ topology directly maps to packed low-bit storage, enabling custom fused dequantize-and-multiply kernels to operate end-to-end without materializing FP16 weights. Autoregressive decode is memory-bandwidth bound, so the storage reduction translates into inference speedup, compounding with model size Figure 5. The speedup is negligible at small scales (131M parameters), but grows substantially to $$\mathcal{L} = \mathcal{L}_{\text{CE}} + \lambda\,\bar{b}$$7 at 9B parameters.

Figure 5: Inference speedup grows with model size; VBQ models at large scale decode several times faster compared to dense FP16, owing to reduced weight traffic.

This demonstrates the unique deployment potential of VBQ: with more parameters, smaller storage, and faster inference, it becomes particularly attractive for serving large-scale LLMs where wall-clock performance and memory bandwidth are limiting factors.

Limitations and Open Problems

VBQ is a train-time technique; it does not serve as a post-training compressor for arbitrary pretrained checkpoints. Economical scaling of the search beyond $$\mathcal{L} = \mathcal{L}_{\text{CE}} + \lambda\,\bar{b}$$8B parameters and recipe refinement for larger models remain as open problems. Attempts to add low-bit correction blocks post hoc to frozen bases did not yield further gains, reinforcing the necessity of co-adaptation during training. The method currently focuses on next-token prediction with English corpora; evaluation on downstream tasks and further scaling are outstanding research directions.

Conclusion

VBQ advances quantization-aware training by treating bit-level numerical precision as a structured, learnable resource, refined via Gumbel-Softmax and alternating optimization. This approach discovers and leverages a highly heterogeneous, per-group precision hierarchy, yielding models that are not only more efficient in terms of storage and inference but also capable of scaling to larger parameter counts under severe memory constraints. The empirical results underscore the superiority of an uneven, learned precision allocation over uniform or heuristic strategies, especially in the "bigger-but-smaller" regime where parameter count offsets aggressive quantization.

The self-healing error-correction observed across network depth provides a theoretical underpinning for robust quality preservation under aggressive precision reduction. Practically, the VBQ recipe is transferable and enables deployment-friendly inference—packed weights accelerate large-scale LLM decode in memory-bound settings. Future exploration should address scaling to multi-billion parameter models, extension beyond next-token prediction, and adaptation to new architectures and tasks, supporting the broader paradigm shift in model training: precision as a structural input, not merely post-training tuning.