PaperlayerView exact revision on arXiv

Source-linked AI summary

BitDistiller: Unleashing the Potential of Sub-4-Bit LLMs via Self-Distillation

Dayou Du, Yijia Zhang, Shijie Cao, Jiaqi Guo, Ting Cao, Xiaowen Chu, Ningyi Xu

arXiv:2402.10631v1cs.CL

TL;DR

Sub-4-bit quantization can degrade LLM weight fidelity and performance, complicating efficient deployment. BitDistiller combines QAT with self-distillation and reports better performance than existing PTQ and QAT methods across language-understanding and reasoning tasks.

  • Problem

    Sub-4-bit quantization degrades weight fidelity and model performance, particularly for smaller models and complex reasoning, amid LLM deployment memory and compute demands.

    Source: §1 · p. 1

  • Method

    BitDistiller combines QAT with self-distillation, using asymmetric quantization and clipping alongside a confidence-aware KL objective.

    Source: §1 · p. 2, §3 · p. 3

  • Results

    BitDistiller significantly outperforms existing PTQ and QAT methods on sub-4-bit language-understanding and complex-reasoning tasks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    BitDistiller achieves notable improvements in 3- and 2-bit settings while requiring less data and fewer training resources.

    Source: §5 · p. 9

    The empirical findings leave unexplained why a 7B teacher outperforms a 13B teacher when distilling a 2-bit 7B student.

    Source: Limitations · p. 9

Abstract

from arXiv · show

The upscaling of Large Language Models (LLMs) has yielded impressive advances in natural language processing, yet it also poses significant deployment challenges. Weight quantization has emerged as a widely embraced solution to reduce memory and computational demands. This paper introduces BitDistiller, a framework that synergizes Quantization-Aware Training (QAT) with Knowledge Distillation (KD) to boost the performance of LLMs at ultra-low precisions (sub-4-bit). Specifically, BitDistiller first incorporates a tailored asymmetric quantization and clipping technique to maximally preserve the fidelity of quantized weights, and then proposes a novel Confidence-Aware Kullback-Leibler Divergence (CAKLD) objective, which is employed in a self-distillation manner to enable faster convergence and superior model performance. Empirical evaluations demonstrate that BitDistiller significantly surpasses existing methods in both 3-bit and 2-bit configurations on general language understanding and complex reasoning benchmarks. Notably, BitDistiller is shown to be more cost-effective, demanding fewer data and training resources. The code is available at https://github.com/DD-DuDa/BitDistiller.

1 Introduction

LLM growth intensifies deployment memory and compute demands, while sub-4-bit quantization risks quality loss. BitDistiller combines QAT and self-distillation, and is reported to outperform existing methods across understanding and reasoning tasks.

  • Quantization reduces model size with minimal performance loss, but below 4 bits weight fidelity and performance suffer, especially for smaller models and complex reasoning.Source: §1 · p. 1
  • BitDistiller significantly outperforms existing PTQ and QAT methods on sub-4-bit language-understanding and complex-reasoning tasks; its code-generation scaling also beats existing QAT methods at 3 and 2 bits.Source: Figure 1 · p. 1, §1 · p. 2
  • BitDistiller combines QAT with self-distillation, using asymmetric quantization and clipping to preserve weight fidelity and a CAKLD objective to improve low-bit learning.The full-precision model acts as teacher for the low-precision student.Source: §1 · p. 2

2 Background and Related Work

Related work contrasts PTQ and QAT for low-bit quantization, then reviews KD approaches and QAT-based distillation. BitDistiller is positioned as combining QAT and KD for extremely low-bit models.

2.1 Weight Quantization for LLMs

  • PTQ avoids retraining but may lose performance at extremely low precision, while QAT learns low-bit representations but still requires extensive training and data.Finer-grained and asymmetric quantization are among the approaches discussed for improving accuracy.Source: §2.1 · p. 2

2.2 Knowledge Distillation for LLMs

  • Prior KD work explores reverse KLD and generalized JSD, while TSLD and LLM-QAT apply distillation with QAT; BitDistiller distinguishes itself by targeting performance and cost-efficiency at extremely low bit levels.Source: §2.2 · pp. 2–3

3 Methodology

BitDistiller combines quantization-aware training with self-distillation, first preparing low-bit weights through asymmetric quantization and clipping, then training them against a full-precision teacher with CAKLD. Its method adapts the distillation trade-off to the teacher’s confidence.

3 Methodology

  • BitDistiller combines QAT with self-distillation, using a full-precision model as teacher and its quantized variant as student.The framework is depicted in Figure 2.Source: §3 · p. 3, Figure 2 · p. 3

3.1 Asymmetric Quantization and Clipping

  • For 2-bit quantization BitDistiller uses INT, while it uses NF formats above 2 bits; asymmetric quantization and clipping are intended to preserve weight fidelity.The paper attributes the weaker 2-bit performance of FP/NF formats to their limited four-value capacity.Source: §3.1 · p. 3
  • Asymmetric clipping is applied only before QAT, using layer-specific α and β chosen to minimize output difference after quantization.This initialization aims to improve final quantized-model accuracy without the cost of iterative clipping optimization.Source: §3.1 · p. 4
  • On a LLaMA-2-7B weight channel quantized at 3 bits with group size 128, asymmetric quantization and clipping preserve higher fidelity than symmetric quantization.Figure 3 compares the original distribution with the two quantized distributions.Source: §3.1 · p. 4, Figure 3 · p. 4

3.2 Self Distillation with CAKLD

  • CAKLD blends reverse and forward KL using a coefficient estimated from averaged token probability, trading mode-seeking against mode-covering according to full-precision-model confidence.It favors mode-seeking when the teacher is confident and mode-covering when it is not.Source: §3.2 · pp. 4–5

4 Experiments

BitDistiller is evaluated against PTQ and QAT baselines on general-language and reasoning tasks, with experiments spanning 3-bit and 2-bit quantization. It leads on the reported language and reasoning benchmarks, while ablations and comparisons examine its quantization strategy, distillation choices, and efficiency.

4.1 Experimental Settings

  • Experiments benchmark LLaMA-2 and domain-specific models on language modeling, QA, coding, and mathematical reasoning, using group-wise 3-bit or 2-bit quantization.The default group size is 128, except for 3B models, which use 64; evaluation compares PTQ and QAT baselines.Source: §4.1 · p. 5, §4 · p. 5

4.2 Evaluation on Language Modeling Tasks; 4.3 Evaluation on Reasoning Tasks

  • BitDistiller leads competing PTQ and QAT methods on general-language tasks and reasoning benchmarks at both 3-bit and 2-bit quantization.For 2-bit quantization, it raises average QA accuracy by +3.54% over LLM-QAT and +12.43% over the leading PTQ method; on complex mathematical reasoning, it achieves 61.33% accuracy, outperforming LLM-QAT by 24.69%.Source: Table 1 · p. 6, §4.2 · p. 6, §4.3 · p. 6

4.4 Ablation Studies

  • On LLaMA-2-7B, asymmetric quantization reduces 2-bit post-training perplexity from 3.4e2 to 16.94, and asymmetric clipping at initialization yields further gains after training.The ablation evaluates perplexity and MMLU (5-shot).Source: §4.4 · p. 6, Table 3 · p. 7
  • Teacher-generated outputs with CAKLD yield lower cross-entropy loss and better HumanEval Pass@1 than fixed-dataset or student-generated outputs, while CAKLD outperforms other tested objectives.The paper also reports that JSD has weak convergence ability for QAT in practice.Source: §4.4 · p. 7, Figure 5 · p. 7, Figure 6 · p. 8

4.5 Analysis and Discussion

  • At 2 bits, BitDistiller surpasses QuIP across the reported benchmarks, while QuIP# performs better on mathematical reasoning and BitDistiller remains stronger in language modeling and programming.The authors suggest that PTQ with incoherence processing and vector quantization could serve as an initialization for BitDistiller.Source: §4.5 · p. 8, Table 4 · p. 8
  • With other settings unchanged, CAKLD converges more rapidly and achieves better overall performance than TSLD.The comparison replaces CAKLD with TSLD in the BitDistiller framework.Source: §4.5 · p. 8, Figure 7 · p. 8
  • Using LLaMA-2-13B rather than LLaMA-2-7B as the teacher did not improve accuracy, and the effects of teacher-student size and architecture require further investigation.The authors suggest that matching student and teacher architectures may improve weight alignment and probability-distribution matching.Source: §4.5 · p. 8, Table 5 · p. 8
  • BitDistiller quantizes WizardCoder-7B in approximately 3 hours on one A100-80G GPU, compared with hundreds of GPU hours for LLM-QAT.For a direct comparison, the authors evaluate LLM-QAT's GPU hours on a single GPU, although its original setup uses 64 GPUs.Source: §4.5 · p. 8, Table 6 · p. 8

5 Conclusion

BitDistiller combines QAT with self-distillation and CAKLD to improve sub-4-bit LLM performance. It reports gains across 3/2-bit language and reasoning tasks with fewer data and training resources.

5 Conclusion

  • BitDistiller outperforms existing PTQ and QAT methods, with notable improvements in 3/2-bit settings across diverse language and reasoning tasks.It is also reported to require fewer data and training resources.Source: §5 · pp. 8–9

Limitations

The authors identify unresolved questions about their empirical findings and two directions for future work: 1-bit quantization and adapting BitDistiller to vector quantization.

Limitations

  • The empirical findings do not explain why a 7B teacher outperforms a 13B teacher when distilling a 2-bit 7B student.The authors suggest matching model architectures may be a reason, but say this is not understood.Source: Limitations · p. 9
  • Extending BitDistiller to 1-bit binary quantization is future work, although the authors describe it as more challenging.They note binary weights could enable inference computation using additions without multiplications.Source: Limitations · p. 9
  • BitDistiller currently applies exclusively to scalar quantization; adapting it to vector quantization is future work.The authors describe vector quantization as a potentially beneficial extension.Source: Limitations · p. 9

A.1 Details of PTQ and QAT Configuration

The PTQ comparison examines calibration-dataset distributions, while the QAT comparison uses asymmetric uniform quantization for fairness. Domain-specific calibration is reported to improve task-specific performance.

A.1 Details of PTQ and QAT Configuration

  • Domain-specific calibration data significantly enhances task-specific PTQ performance.Figure 8 compares PTQ methods using various calibration datasets; domain-specific datasets are used for reasoning tasks.Source: §A.1 · pp. 11–12, Figure 8 · p. 11
  • For a fair QAT comparison, the authors replicate LLM-QAT with asymmetric uniform quantization rather than its symmetric quantization setup.They report that symmetric quantization in LLM-QAT degrades performance with grouped quantization.Source: §A.1 · p. 12

A.2 Implementation Details and Analysis of Confidence-Aware KLD

The authors estimate CAKLD’s coefficient using teacher-model token probabilities from ten training batches. Their confidence analysis motivates balancing different distillation behaviors for text generation and reasoning.

A.2 Implementation Details and Analysis of Confidence-Aware KLD

  • Ten training batches are run through the teacher without parameter updates to estimate γ from the average token probability.The teacher logits are used to compute that probability.Source: §A.2 · p. 12
  • Teacher confidence varies by task: text generation has variable confidence, while reasoning next-token predictions are highly confident.The analysis is presented in Figure 9.Source: §A.2 · p. 12, Figure 9 · p. 12
  • CAKLD balances learning multiple teacher modes for text generation with focusing on the predominant mode for reasoning.The authors motivate mean-seeking KL for generation and predominant-mode learning for reasoning.Source: §A.2 · p. 12

A.3 Training Datasets Examples

Training data varies by task: general language uses Alpaca and WikiText-2, code tasks use Evol-Instruct-Code, and mathematical reasoning uses MetaMathQA. Self-distillation uses only a small portion of the involved datasets.

A.3 Training Datasets Examples

  • General language training mixes Alpaca and WikiText-2 token sequences at a 2:1 ratio.For WikiText-2, the first 128 corpus tokens prompt teacher generation at temperature 0.7.Source: §A.3 · p. 12
  • Code understanding and generation use Evol-Instruct-Code, while mathematical reasoning uses MetaMathQA.Examples of the training data are shown in Table 8.Source: §A.3 · p. 12
  • The self-distillation process uses only a small portion of the involved datasets.Source: §A.3 · p. 12

A.5 Integration with AWQ For Quantization Strategies

The section evaluates asymmetric clipping alongside AWQ for QAT initialization. Asymmetric clipping performs robustly at sub-4-bit precision, while AWQ provides no additional gain after training.

A.5 Integration with AWQ For Quantization Strategies

  • At 2-bit quantization, INT-Asym and AWQ cannot complete the task, whereas Clip-Asym succeeds and improves WikiText-2 perplexity.Table 7 reports quantization-initialization results for QAT with AWQ on WikiText-2 perplexity.Source: §A.5 · p. 12, Table 7 · p. 12
  • AWQ before QAT initially improves performance, but yields no additional gain after training; straightforward clipping is sufficient to initialize QAT.Source: §A.5 · pp. 12–13
  • On LLaMA-2-13B general language tasks, BitDistiller achieves leading performance against established PTQ and QAT methods at both 3-bit and 2-bit quantization.Table 9 reports the comparison.Source: Table 9 · p. 14
Loading 2402.10631v1…