Compact Language Models via Pruning and Knowledge Distillation

Large language models (LLMs) targeting different deployment scales and sizes are currently produced by training each variant from scratch; this is extremely compute-intensive. In this paper, we investigate if pruning an existing LLM and then re-training it with a fraction (<3%) of the original training data can be a suitable alternative to repeated, full retraining. To this end, we develop a set of practical and effective compression best practices for LLMs that combine depth, width, attention and MLP pruning with knowledge distillation-based retraining; we arrive at these best practices through a detailed empirical exploration of pruning strategies for each axis, methods to combine axes, distillation strategies, and search techniques for arriving at optimal compressed architectures. We use this guide to compress the Nemotron-4 family of LLMs by a factor of 2-4x, and compare their performance to similarly-sized models on a variety of language modeling tasks. Deriving 8B and 4B models from an already pretrained 15B model using our approach requires up to 40x fewer training tokens per model compared to training from scratch; this results in compute cost savings of 1.8x for training the full model family (15B, 8B, and 4B). Minitron models exhibit up to a 16% improvement in MMLU scores compared to training from scratch, perform comparably to other community models such as Mistral 7B, Gemma 7B and Llama-3 8B, and outperform state-of-the-art compression techniques from the literature. We have open-sourced Minitron model weights on Huggingface, with corresponding supplementary material including example code available on GitHub.

한 줄 요약

Nemotron-4 15B 모델을 2-4× 압축한 Minitron 모델은 40× 적은 훈련 데이터로 기존 모델과 유사하거나 더 높은 성능을 보인다.

핵심 기여도

Nemotron-4 15B 모델을 2-4× 압축하여 Minitron 8B 및 4B 모델을 개발.
40× 적은 훈련 토큰으로 Minitron 모델을 재훈련하여 Nemotron-3 8B 및 Llama-2 7B 대비 더 높은 성능 달성.
Minitron 8B는 MMLU에서 16% 개선된 성능, Minitron 4B는 Gemma2 대비 우수한 성능 보임.
학습 비용 1.8× 절감하며, Minitron 모델 가중치는 Huggingface에 오픈소스로 공개.

핵심 아이디어

기존 대형 언어 모델(LLM)을 매번 처음부터 훈련하는 방식은 계산 비용이 매우 높기 때문에, 이 연구는 기존 LLM을 가지치기(pruning)하고 재훈련하는 방식을 제안한다. 구체적으로, 네 가지 축(깊이, 너비, 어텐션, MLP)을 조합한 구조적 가지치기와 지식 증류(knowledge distillation) 기반의 재훈련을 통해 모델 압축을 수행한다. 연구팀은 각 가지치기 축의 전략, 축 간 조합 방법, 재훈련 전략, 최적 압축 아키텍처 탐색 기법을 체계적으로 실험하고, 이를 바탕으로 실용적인 압축 최고 기준(best practices)을 제시한다. 예를 들어, MLP 및 어텐션 헤드를 가지치키고 재훈련하는 경우, 초기에는 단일 축만 가지치키는 것이 더 효과적이지만, 재훈련 후에는 여러 축을 조합한 가지치기가 더 나은 결과를 보인다는 통찰을 얻었다.

기술적 접근법

**가지치기 축**: MLP, 어텐션 헤드, 임베딩 채널, 모델 깊이.
**재훈련 전략**: 지식 증류(distillation)를 사용하여 기존 15B 모델의 지식을 8B 및 4B 모델로 전달.
**훈련 데이터**: Nemotron-4의 8T 토큰 데이터 중 1.8B 토큰(400 스텝)만 사용.
**하이퍼파라미터**: 코사인 학습률 감소 스케줄(2⁻⁴ → 4.5⁻⁷), 동일 최적화 설정 사용.
**프레임워크**: NVIDIA Megatron-LM을 사용하여 가지치기 및 재훈련 알고리즘 구현.

주요 결과

**MMLU 성능**: Minitron 8B는 63.8% (기존 8B 모델 대비 +16%)
**Winogrande**: 79.0% (Llama-2 7B 대비 +5%)
**GSM8K**: 51.3% (Nemotron-3 8B 대비 +26.8%)
**XLSum**: 31.2% (Gemma 7B 대비 +14.2%)
**MBPP**: 35.2% (Phi-2 2.7B 대비 +8.2%)
**전체 훈련 토큰 절감**: 40× (15B → 8B, 4B)
**계산 비용 절감**: 1.8× (15B, 8B, 4B 모델 가족 전체)

의의 및 한계

이 연구는 LLM 압축 분야에서 구조적 가지치기와 재훈련을 결합한 첫 번째 체계적인 실험을 수행한 것으로, 다양한 축의 조합 효과와 최적 재훈련 전략을 제시한다. Minitron 모델은 기존 8B 및 4B 규모의 커뮤니티 모델과 비교해 유사하거나 더 높은 성능을 보이며, 기존 압축 기법 대비 우수한 결과를 보인다. 그러나 연구는 Nemotron-4 15B 모델에만 적용되었으며, 다른 아키텍처나 데이터셋에서 동일한 효과가 나타날지는 추가 실험이 필요하다. 또한, 재훈련 단계에서 사용된 데이터는 3% 미만이지만, 이는 여전히 상당한 규모의 데이터를 필요로 하므로, 더 데이터 효율적인 방법이 필요할 수 있다.

실용적 활용

Minitron 모델은 저비용으로 다양한 규모의 LLM을 생성할 수 있어, 클라우드, 모바일, 임베디드 시스템 등에서의 배포에 적합하다. 연구자나 엔지니어는 기존 대형 모델을 기반으로 빠르게 소형 모델을 개발할 수 있으며, 특히 토큰 수 제한이 있는 환경에서 유용하다. 오픈소스로 공개된 Minitron 모델은 커뮤니티에서 즉시 활용할 수 있어, 연구 및 산업적 적용이 용이하다.