대형 언어 모델을 위한 사전 훈련 증류: 디자인 공간 탐색

Pre-training Distillation for Large Language Models: A Design Space Exploration

October 21, 2024

저자: Hao Peng, Xin Lv, Yushi Bai, Zijun Yao, Jiajie Zhang, Lei Hou, Juanzi Li

cs.AI

초록

지식 증류(Knowledge distillation, KD)은 큰 교사 모델로부터 작은 학생 모델로 지식을 전달하는 것을 목표로 합니다. 이전 연구에서는 대형 언어 모델(Large Language Models, LLMs) 분야에서 KD를 적용할 때 주로 후훈련 단계에 초점을 맞추어 왔는데, 여기서 학생 LLM이 교사 모델이 생성한 지시와 해당 응답에서 직접 학습하는 방식이었습니다. 본 논문에서는 KD를 LLM의 사전훈련 단계로 확장하여 사전훈련 증류(Pre-training distillation, PD)이라고 명명하였습니다. 우리는 먼저 GLM-4-9B를 교사 LLM으로 사용하여 1.9B 파라미터 학생 LLM을 증류하는 초기 실험을 수행하여 PD의 효과를 검증했습니다. 증류의 주요 영향 요소를 고려하여 사전훈련 증류의 설계 공간을 로짓 처리, 손실 선택, 스케일링 법칙, 오프라인 또는 온라인 로짓 네 가지 측면에서 체계적으로 탐색했습니다. 우리는 사전훈련 증류의 설계 공간을 탐색하기 위해 포괄적인 실험을 수행하고 더 나은 구성 및 흥미로운 결론을 도출했습니다. 예를 들어, 일반적으로 큰 학생 LLM이 사전훈련 증류에서 더 많은 이점을 얻는 반면, 큰 교사 LLM이 반드시 더 나은 결과를 보장하지는 않는다는 사실을 발견했습니다. 우리의 설계 공간 탐색이 사전훈련 증류에 대한 미래 실천에 도움이 되기를 희망합니다.

English

Knowledge distillation (KD) aims to transfer knowledge from a large teacher model to a smaller student model. Previous work applying KD in the field of large language models (LLMs) typically focused on the post-training phase, where the student LLM learns directly from instructions and corresponding responses generated by the teacher model. In this paper, we extend KD to the pre-training phase of LLMs, named pre-training distillation (PD). We first conduct a preliminary experiment using GLM-4-9B as the teacher LLM to distill a 1.9B parameter student LLM, validating the effectiveness of PD. Considering the key impact factors of distillation, we systematically explore the design space of pre-training distillation across four aspects: logits processing, loss selection, scaling law, and offline or online logits. We conduct extensive experiments to explore the design space of pre-training distillation and find better configurations and interesting conclusions, such as larger student LLMs generally benefiting more from pre-training distillation, while a larger teacher LLM does not necessarily guarantee better results. We hope our exploration of the design space will inform future practices in pre-training distillation.