BitNet 증류

초록

본 논문에서는 특정 하위 작업을 위해 기성의 완전 정밀도 대형 언어 모델(LLM, 예: Qwen)을 1.58비트 정밀도(즉, 삼항 가중치 {-1, 0, 1})로 미세 조정하는 경량 파이프라인인 BitNet Distillation(BitDistill)을 제안한다. 이를 통해 최소한의 계산 비용으로 강력한 작업별 성능을 달성한다. 구체적으로, BitDistill은 BitNet에서 소개된 SubLN 모듈, MiniLM을 기반으로 한 다중 헤드 어텐션 디스틸레이션, 그리고 특정 작업에서 미세 조정된 완전 정밀도와 1.58비트 LLM 간의 성능 격차 확장성 문제를 완화하기 위한 중요한 준비 단계로 작용하는 지속적 사전 학습이라는 세 가지 핵심 기술을 통합한다. 실험 결과, BitDistill은 모델 크기에 걸쳐 완전 정밀도 대응 모델과 비슷한 성능을 달성하면서도 최대 10배의 메모리 절약과 CPU에서 2.65배 더 빠른 추론을 가능하게 한다. 코드는 https://github.com/microsoft/BitNet에서 확인할 수 있다.

English

In this paper, we present BitNet Distillation (BitDistill), a lightweight pipeline that fine-tunes off-the-shelf full-precision LLMs (e.g., Qwen) into 1.58-bit precision (i.e., ternary weights {-1, 0, 1}) for specific downstream tasks, achieving strong task-specific performance with minimal computational cost. Specifically, BitDistill incorporates three key techniques: the SubLN module, as introduced in BitNet; multi-head attention distillation, based on MiniLM; and continual pre-training, which serves as a crucial warm-up step to mitigate the scalability issue of the performance gap between finetuned full-precision and 1.58-bit LLMs on specific tasks. Experimental results show that BitDistill achieves performance comparable to the full-precision counterpart models across model size, while enabling up to 10x memory savings and 2.65x faster inference on CPUs. Code is available at https://github.com/microsoft/BitNet.