[논문 리뷰] DistiLLM: Towards Streamlined Distillation for Large Language Models

Jongwoo Ko, Sungnyun Kim외 2명|arXiv (Cornell University)|2024. 02. 06.

Topic Modeling인용 수 4

한 줄 요약

DistiLLM는 자동회귀 대규모 언어 모델을 위한 간소화된 지식 정복 프레임워크를 제안하며, 안정적이고 수렴 가능한 학습을 위해 새로운 기울기 Kullback-Leibler 발산 손실을 도입하고, 학생 모델이 생성한 출력을 효율적으로 활용하기 위한 적응형 오프-정책 방법을 제안한다. 최근의 KD 방법보다 최대 4.3배 빠른 학습 속도를 기록하며 지시 따르기 및 요약 작업에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Knowledge distillation (KD) is widely used for compressing a teacher model to a smaller student model, reducing its inference cost and memory footprint while preserving model capabilities. However, current KD methods for auto-regressive sequence models (e.g., large language models) suffer from missing a standardized objective function. Moreover, the recent use of student-generated outputs to address training-inference mismatches has significantly escalated computational costs. To tackle these issues, we introduce DistiLLM, a more effective and efficient KD framework for auto-regressive language models. DistiLLM comprises two components: (1) a novel skew Kullback-Leibler divergence loss, where we unveil and leverage its theoretical properties, and (2) an adaptive off-policy approach designed to enhance the efficiency in utilizing student-generated outputs. Extensive experiments, including instruction-following tasks, demonstrate the effectiveness of DistiLLM in building high-performing student models while achieving up to 4.3$ imes$ speedup compared to recent KD methods.

연구 동기 및 목표

• 자동회귀 언어 모델을 위한 지식 정복에서 표준화되고 안정적인 목적 함수의 부족을 해결하기 위해.
• 오프-정책 정복에서 과도하게 생성된 학생 모델 출력(SGO)으로 인한 학습 비효율성을 줄이기 위해.
• 학생 모델 출력의 모드 붕괴와 분포 불일치를 최소화하여 정복 효과를 향상시키기 위해.
• 학습 시간과 계산 비용을 크게 줄여가며 고성능의 학생 모델을 달성하기 위해.
• 이론적으로 탄탄하고 효율적이며 효과적인 대규모 언어 모델 정복 프레임워크를 제공하기 위해.

제안 방법

• 학습 중 기울기 안정성과 근사 오차 감소를 위해 강력한 이론적 기반을 지닌 기울기 Kullback-Leibler(KLD) 발산 손실을 도입한다.
• 학습에 사용되는 학생 모델 출력 비율을 동적으로 조절하는 적응형 오프-정책 전략을 적용하여 샘플 효율성을 향상시킨다.
• 고정 데이터셋 감시와 온-정책 학생 모델 출력 시퀀스를 모두 활용하여 분포 일치성과 학습 효율성을 균형 잡는다.
• 교사와 학생의 출력 분포 간 비대칭 발산을 최소화하도록 정복 목적 함수를 최적화하여 모드 붕괴를 방지한다.
• 학습을 안정적으로 진행하기 위해 점진적으로 학생 모델 출력에 의존도를 높이는 커리큘럼 유사 학습 스케줄을 적용한다.
• 시퀀스 수준의 출력에 기울기 KLD 손실을 적용하여 다양한 자동회귀 작업에서 교사의 생성 행동과의 보다 나은 일치를 보장한다.

Figure 1 : Examples of SGOs from GPT-2 (student) and their corresponding validation loss by GPT-2 XL (teacher). Since the teacher model may be not familiar with the SGO, using $p(\mathbf{y}|\mathbf{x})$ as a target distribution can misguide the student model, as shown in Tab. 7 .

실험 결과

연구 질문

• RQ1 이론적으로 탄탄한 비대칭 기울기 KLD 손실이 자동회귀 LMs의 지식 정복에서 학습 안정성과 수렴성을 향상시킬 수 있는가?
• RQ2 높은 계산 비용 없이 학생 모델 출력을 효율적이고 적응적으로 정복에 통합할 수 있는가?
• RQ3 고정 데이터셋 감시와 적응형 오프-정책 SGO를 조합하면 온-정책 또는 고정 SGO 전략 대비 더 나은 일반화와 더 빠른 수렴을 이룰 수 있는가?
• RQ4 제안된 방법이 생성 작업에서 학생 모델 성능을 유지하거나 향상시키면서도 학습 시간을 얼마나 줄일 수 있는가?
• RQ5 기울기 KLD 손실이 표준 KLD 및 반전 KLD 접근 방식에 비해 다양한 지시 따르기 및 텍스트 생성 벤치마크에서 강건한가?

주요 결과

• DistiLLM는 지시 따르기 및 텍스트 요약 작업에서 기존의 KD 방법들을 능가하는 최신 기술 수준의 성능을 달성한다.
• 최근의 오프-정책 KD 방법들인 Gu et al. (2024) 및 Agarwal et al. (2024)에 비해 학습 시간을 최대 4.3배 단축시킨다.
• 기울기 KLD 손실은 더 빠른 수렴과 더 안정적인 학습을 가능하게 하여 기존 KLD에서 관찰된 기울기 폭발 위험을 줄인다.
• 적응형 오프-정책 모듈은 동적으로 학생 모델 출력 사용 비율을 조절함으로써 재료가 되거나 저품질인 시퀀스를 피함으로써 샘플 효율성을 향상시킨다.
• Self-Instruct 및 Vicuna 벤치마크에서, 다른 방법들이 실패하거나 환각된 출력을 생성한 경우에도 DistiLLM는 100%의 경우에서 간결하고 정확한 응답을 올바르게 생성한다.
• 제한된 SGO 사용 조건에서도 뛰어난 성능을 달성함으로써 고품질 정복이 끊임없이 비싼 학생 모델 생성이 필요하지 않음을 보여준다.

Figure 2 : ( Left ): Normalized runtime according to the maximum response length of SGOs with GPT-2 XL teacher and GPT-2 student. ( Right ): Normalized runtime for various sizes of teacher and student models with a response length of 256. FWD and BWD denote forward and backward propagation, respecti