Distilling Step-by-Step! Outperforming Larger Language Models with Less Training Data and Smaller Model Sizes

https://arxiv.org/abs/2305.02301

Distilling Step-by-Step! Outperforming Larger Language Models with Less Training Data and Smaller Model Sizes

옛날 논문이라 이전의 지식들을 이야기 하긴 하지만 지금 사용하는 distill의 초창기겠네요

llm을 통해 step을 생성하고, 단순 정답만 학습하는게 아니라 그 step도 학습하여 더 높은 성능을 보여준다! 가 나옵니다. 

이게 distilling step-by-step이 llm보다 쌘 것이 의아하긴 하지만 뭐 범용모델이 성능은 그렇게 좋지 않다는 것을 보여주는 것 같습니다. 

 

CoT를 통해 추론 과정을 뽑고 작은 모델에 이걸 학습시킨다! 

근데 단순 Rationale + label 을 출력하도록 학습하는 것이 아닌 

Rationale, label을 출력하도록 따로 따로 로스를 더해 학습함 

 

Teacher는 540B PaLM이고, Student는 T5-Base 220M, T5-Large 770M, T5-XXL 11B가 사용되었습니다.

 

더 적은 데이터로 standard fine-tuning을 능가하는 모습을 보여줬다. 

단순 성능만을 올리는 것이 아니라 데이터 요구량 자체를 줄였습니다. 

라벨이 없는 경우에도 PaLM이 생성한 라벨과 출력물을 통해 학습할 수 있었다.- 여기서도 적은 데이터로 높은 성능을 가질 수 있었음 

작은 모델이 큰 모델의 CoT를 뛰어 넘는 것을 볼 수 있다. 

모델 사이즈가.... 흠 

지금은 모델 성능이 워낙 좋아서 이 때랑 비교가 안 될 것 같기도 하고요 

 

핵심 문제	대형 LLM은 few-shot / CoT 성능이 뛰어나지만, 모델 크기·GPU 메모리·추론 비용 때문에 실제 배포가 어렵다. 반면 작은 task-specific model은 배포는 쉽지만, 기존 fine-tuning이나 distillation 방식은 많은 학습 데이터가 필요하다.
핵심 아이디어	LLM을 단순히 pseudo-label 생성기로 쓰지 않고, 정답에 도달하는 rationale / reasoning step을 생성하는 teacher로 활용한다. 작은 모델은 label뿐 아니라 rationale도 함께 학습한다.
제안 방법	Distilling Step-by-Step: ① Few-shot CoT prompting으로 LLM에서 label과 rationale을 추출한다. ② 작은 T5 모델을 multi-task learning으로 학습한다. 하나의 task는 label prediction, 다른 task는 rationale generation이다.
학습 목적 함수	최종 loss는 L = L_label + λL_rationale이다. 즉, 작은 모델이 정답을 맞히는 능력과 reasoning 과정을 생성하는 능력을 동시에 학습하도록 한다.
중요한 설계 포인트	Rationale을 입력으로 넣지 않고 출력 supervision으로 사용한다. 따라서 추론 시에는 LLM이 필요 없고, 작은 모델이 label만 예측하면 된다.
비교 대상	Standard fine-tuning, standard task distillation, Few-shot CoT, PINTO tuning
Teacher model	주로 540B PaLM을 사용하고, ablation에서는 20B GPT-NeoX도 사용한다.
Student model	T5-Base 220M, T5-Large 770M, T5-XXL 11B
사용 데이터셋	e-SNLI, ANLI, CommonsenseQA(CQA), SVAMP. 각각 자연어 추론, adversarial NLI, 상식 질의응답, 수학 word problem을 평가한다.
주요 결과 1: 데이터 효율성	220M T5 기준, Distilling Step-by-Step은 standard fine-tuning보다 훨씬 적은 labeled data로 더 높은 성능을 낸다. e-SNLI에서는 전체 데이터의 12.5%만 사용해도 full-data fine-tuning을 능가한다.
주요 결과 2: Unlabeled distillation	라벨이 없는 경우에도 LLM이 생성한 pseudo-label과 rationale을 사용하면 standard task distillation보다 적은 unlabeled data로 더 좋은 성능을 달성한다.
주요 결과 3: 작은 모델로 LLM 능가	Distilling Step-by-Step으로 학습한 작은 T5 모델이 540B PaLM Few-shot CoT를 여러 데이터셋에서 능가한다. 예를 들어 e-SNLI에서는 220M T5가, ANLI와 SVAMP에서는 770M T5가 PaLM을 능가한다.
주요 결과 4: 최소 자원 분석	ANLI에서 770M T5가 80% 데이터만으로 540B PaLM Few-shot CoT를 능가한다. 반면 standard fine-tuning은 100% 데이터를 사용해도 PaLM 성능을 맞추기 어렵다.
Ablation 1: Teacher 크기	20B GPT-NeoX에서 추출한 rationale도 성능 향상을 제공하지만, 540B PaLM에서 추출한 rationale이 더 높은 성능을 낸다. 즉, teacher의 rationale quality가 중요하다.
Ablation 2: Multi-task의 중요성	Rationale과 label을 하나의 sequence로 붙여 예측하는 single-task 방식은 일부 데이터셋에서 fine-tuning보다 나쁘다. 반면 label prediction과 rationale generation을 분리한 multi-task 방식이 가장 안정적이다.
논문의 핵심 기여	LLM의 지식을 단순 label이 아니라 reasoning rationale 형태로 작은 모델에 압축했다. 이를 통해 데이터 효율성과 배포 효율성을 동시에 개선했다.
한계점	Few-shot CoT prompt 예시가 필요하고, teacher LLM의 rationale 품질에 성능이 의존한다. 또한 복잡한 reasoning/planning task에서는 LLM rationale 자체가 부정확할 수 있으며, teacher의 bias가 student로 전이될 수 있다.
최종 결론	이 논문은 “LLM의 답만 증류하지 말고, 답을 도출하는 reasoning 과정까지 증류하면 작은 모델도 적은 데이터로 대형 LLM을 능가할 수 있다”는 것을 실험적으로 보인 연구다. 특히 test-time에는 작은 모델만 사용하면 되므로 실용적인 LLM compression / task-specific deployment 방법으로 의미가 크다.