■ LLM은 complex reasoning tasks에서 뛰어난 성능을 보이며, LLM의 이러한 reasoning capabilities을 더 작은 모델로 distillation하는 방식이 등장했다.
■ 그러나 저자들은 흥미로운 현상을 발견했고, 이를 "Small Model Learnability Gap"이라고 부른다: 3B 이하의 소규모 모델들이 긴 CoT reasoning이나 더 큰 모델로부터의 distillation에서 일관되게 이점을 얻지 못한다.
■ 오히려 작은 모델들은 자신들의 내재적 학습 능력에 더 잘 맞는 짧고 단순한 reasoning chain으로 fine-tuning될 때 더 좋은 성능을 보인다.
■ 그래서 논문에서는 Mix Distillation을 제안한다. 이는 long CoT와 short CoT 예시를 결합하거나, 큰 모델과 작은 모델의 reasoning을 함께 사용하는 방법으로, 저자들은 이 방법이 reasoning complexity의 균형을 맞추는 단순하지만 효과적인 전략이라고 소개한다.
■ 실험을 통해 Mix Distillation이 둘 중 하나의 데이터만으로 학습하는 경우보다 작은 모델의 reasoning 성능을 유의미하게 향상시킨다는 것을 보여준다.
■ 이러한 발견은 strong model로부터 직접적으로 distillation하는 방식의 한계를 보여주며, reasoning 능력을 효과적으로 전달하기 위해서는 reasoning complexity를 student model에 맞게 조정하는 것이 중요함을 강조한다.
[2502.12143] Small Models Struggle to Learn from Strong Reasoners
Small Models Struggle to Learn from Strong Reasoners
Large language models (LLMs) excel in complex reasoning tasks, and distilling their reasoning capabilities into smaller models has shown promise. However, we uncover an interesting phenomenon, which we term the Small Model Learnability Gap: small models ($
arxiv.org
1. Introduction
■ LLM은 mathematical problem-solving, logical inference, structured decision-making과 같은 complex reasoning tasks에서 뛰어난 성능을 보여 왔다.
■ LLM의 complex reasoning 능력을 향상시키는 데 중요한 발전 중 하나는 chain-of-thought prompting이다. CoT는 복잡한 문제를 intermediate reasoning steps로 분해하여 성능과 interpretability(즉, 모델이 왜 그런 답을 냈는지에 대한 과정)을 모두 향상시킨다.
■ 그러나 LLM의 높은 계산 비용은 자원이 제한된 장치에서의 배포를 어렵게 만들고, 이에 따라 더 낮은 비용으로 비슷한 능력을 제공하는 작은 모델 개발이 필요해졌다.
■ 이를 위해 흔히 사용되는 전략은 distillation이다. strong teacher model이 생성한 CoT sequence를 사용해 weaker student model을 파인튜닝하는 것이다. 이렇게 더 좋은 모델로부터 CoT sequence를 distillation하면 small model의 complex reasoning capabilities이 향상될 것이라고 기대할 수 있다.
■ 그러나 저자들은 "Small Model Learnability Gap"이라고 명명한 현상을 발견했다: small model은 long CoT reasoning이나 large model로부터의 distillation처럼 strong teacher가 제공하는 complex reasoning sequences로부터 일관된 이점을 얻지 못한다는 것이다.
■ 저자들의 실험에서, small model이 길고 복잡한 reasoning trace에 노출되면, 제한된 능력 때문에 multi-step logic을 내재화하는 데 어려움을 겪는다는 것을 관찰했다.
■ 오히려 small model은 자신의 내재적 학습 능력에 더 잘 맞는 짧고 단순한 reasoning chain으로 fine-tuning될 때 더 좋은 성능을 보인다.
■ 이는 small model이 지나치게 정교한 reasoning trace를 처리하거나, stronger teacher가 만든 것에 대한 distribution shift에 적응하는 데 어려움을 겪으며, 결국 효과적으로 generalize하는 능력이 제한된다는 것을 시사한다.
■ 이러한 문제를 해결하기 위한 방법으로, 저자들은 Mix Distillation을 제안한다. 이는 서로 다른 유형의 reasoning traces을 섞어서 reasoning complexity를 조율하는 간단하면서도 효과적인 접근법이다.
■ 구체적으로, 이 방법은 두 가지 configuration으로 구성된다.
- (1) Mix-Long: long CoT와 short CoT examples의 조합으로, small model이 자세한 reasoning step과 간결한 reasoning step 모두에 노출되도록 한다.
- (2) Mix-Large: larger model과 small model의 responses을 함께 섞는 방식으로, small model이 자신의 capacity에 더 적합한 reasoning chain으로부터 학습할 수 있게 한다.
■ 저자들의 실험은 Mix Distillation이 standard distillation과 비교해 small model의 reasoning 성능을 일관되게 향상시킨다는 것을 보여준다.
- Qwen2.5-3B-Instruct는 long CoT data만 직접 학습한 경우와 비교했을 때, Mix-Long을 사용하면 MATH와 AMC에서 8점 이상 향상된다.
- Qwen2.5-3B-Instruct는 large teacher CoT data만 학습한 경우와 비교했을 때, Mix-Large를 사용하면 MATH, AIME, AMC에서 7점 이상 향상된다.
■ 이러한 발견은 strong model로부터 직접적으로 distillation하는 방식의 근본적인 한계를 보여주며, 효과적인 knowledge transfer를 위해 reasoning complexity를 조정하는 것이 중요함을 강조한다.
2. Preliminaries
2.1 Notation
■ input sequence(예: 프롬프트)를 \( x = (x_1, x_2, \dots, x_n) \)로, 이에 대응되는 output sequence를 \( y = (y_1, y_2, \dots, y_m) \)으로 나타낸다.
■ parameter \( \theta \)로 parameterized된 LLM을 고려하며, 이 모델은 conditional distribution \( \pi_\theta(y_t | x, y_{1:t-1}) \)를 따라 next token을 예측한다.
■ 생성된 output 안에서 chain-of-thought를 담고 있는 token들의 부분집합을 \( \text{CoT}(y) \subseteq y \)로 표기한다. 이것은 reasoning trace 또는 explanation sequence 역할을 한다.
■ 논문 전체에 걸쳐, "short CoT"는 solution에 도달하기 위한 간결한 reasoning path를 가리키는 말로 사용된다.
■ 그리고 단순히 길이가 더 길 뿐만 아니라 더 복잡한 reflective thought까지 보여주는 확장된 reasoning sequence를 "long CoT"라고 부른다.
■ 더 큰 teacher model이 생성한 reasoning trace를 "large teacher CoT", 더 작은 teacher model이 생성한 reasoning step을 "small teacher CoT"라고 부른다.
2.2 Supervised Fine-Tuning (SFT)
■ supervised fine-tuning (SFT)은 dataset \( \mathcal{D} = \{(x^i, y^i)\}_{i=1}^N \)에서 LLM의 reasoning capability를 향상시키기 위해 널리 사용된다.
■ 여기서 \( y^i \)는 short CoT, long CoT, strong model CoT 또는 weak model CoT sequence가 될 수 있다.
■ SFT 과정은 instruction dataset \( \mathcal{D} \)에 대한 negative log-likelihood loss를 최소화함으로써 language model의 parameter \( \theta \)를 업데이트한다.
3. Small Model Learnability Gap
■ 이 섹션에서는 서로 다른 CoT data를 사용하여 student models을 파인튜닝한다. 그런 다음 파인튜닝된 모델들의 성능을 바탕으로 "small model learnability gap"을 밝혀낸다.
3.1 Experiment Setup
Datasets
■ MATH(hendrycks의 MATH) 의 7,500개 prompt set을 사용한다. 이 dataset은 advanced calculus, geometry, linear algebra 등을 포함한 7가지 수학 topics로 구성된다.
Student models
■ 서로 다른 크기를 가진 Qwen과 Llama model family의 student model 10개를 고려한다.
■ 이 모델들에는 Qwen2.5-0.5B, Qwen2.5-1.5B, Qwen2.5-3B, Qwen2.5-7B, Qwen2.5-14B, Qwen2.5-32B의 Instruct 버전과 Llama3.2-1B, Llama3.2-3B, Llama3.1-8B, Llama3.3-70B의 Instruct 버전이 포함된다.
Teacher models
■ long CoT와 short CoT를 비교하기 위해, long CoT sequence 생성에는 QwQ-32B-Preview를, short CoT 생성에는 Qwen2.5-32B-Instruct를 response generator로 사용한다.
■ 각 model family 안에서, 더 큰 scale의 model을 large teacher로, 더 작은 scale의 model을 small teacher로 지정한다.
- Qwen의 경우 large teacher는 Qwen2.5-72B-Instruct, small teacher는 Qwen2.5-3B-Instruct (Table 4)
Evaluation Benchmarks
■ fine-tuning된 student model의 reasoning capability를 MATH, GSM8K, AMC 2023, AIME 2024, OlympiadBench의 English math subset에서 평가한다.
■ 이 벤치마크들은 초등 수학부터 고급 경시대회 문제에 이르기까지 광범위한 난이도를 포괄한다.
■ student model의 성능을 5개 벤치마크의 평균 점수로 정의한다.
■ 별도 언급이 없는 한, 모든 fine-tuning된 model은 zero-shot setting에서 greedy decoding으로 평가하며, maximum generation token 수는 16k로 설정한다.
■ 다음과 같은 성능 점수를 정의한다.
- \( P_{Long} \): long CoT data로 fine-tuning한 student model의 성능.
- \( P_{Short} \): short CoT data로 fine-tuning한 student model의 성능.
- \( P_{Large} \): large teacher가 만든 CoT로 fine-tuning한 student model의 성능.
- \( P_{Small} \): small teacher가 만든 CoT로 fine-tuning한 student model의 성능.
Training Setup
■ teacher models은 rejection sampling으로 responses을 생성한다. 기본적으로 teacher models은 greedy decoding을 사용한다.
■ math problem instructions과 teacher models이 생성한 ( math problem instructions에 대한) solutions을 결합하여 student model을 fine-tuning하기 위한 problem-solution pair를 구성한다.
■ 저자들은 LLaMA-Factory framework를 사용해 models을 학습시켰다.
■ 14B보다 작은 student model에는 full-parameter SFT를 사용하고, cosine learning-rate schedule과 최대 learning rate \( 10^{-5} \)를 적용해 2 epochs 동안 fine-tuning한다.
■ 14B보다 큰 student model에는 LoRA fine-tuning을 사용하고, learning rate \( 10^{-4} \)로 2 epochs 동안 fine-tuning한다.
3.2 Long CoT Gap
■ 이 섹션에서는 long CoT data와 short CoT data로 fine-tuning된 student model들의 reasoning capability를 평가한다.
■ 다음과 같이 정의된 "long CoT gap" \( \Delta_{Long} \)을 사용하여 long CoT data와 short CoT data 사이의 성능 차이를 정량화한다.
■ Fig 2는 서로 다른 크기의 student models에 걸친 Long CoT Gap을 보여준다. 모델 크기가 커질수록 gap의 방향이 바뀌는 것을 볼 수 있다.
■ Table 8은 MATH, GSM8K, AIME, AMC, OlympiadBench에서의 score이다.
■ long CoT가 larger model에서 더 효과적이며, 대부분의 math benchmark에서 일관된 성능 향상으로 이어진다.
■ 그러나 smaller models에서 long CoT data는 효과적이지 않으며, short CoT와 비교하면 훨씬 적은 개선을 보인다.
■ MATH와 AMC benchmark에서 Qwen2.5-1.5B-Instruct는 long CoT data로 fine-tuning되었을 때, short CoT data로 fine-tuning된 경우보다 10 point 이상 낮은 성능을 보인다.
■ 이는 smaller model이 long CoT paradigm을 효과적으로 학습하고 활용하지 못할 수 있음을 보여준다.
■ 이 섹션의 결론은, small student models은 short CoT에서 더 큰 이을 얻는 경향이 있고, large student models long CoT에서 더 큰 이을 얻는다.
3.3 Large Teacher CoT Gap
■ small models이 large teacher와 small teacher로부터 얼마나 효과적으로 학습하는지를 확인한다. 이를 위해 large teacher CoT gap을 다음과 같이 정의한다.
■ Fig 3은 모든 student model에서 발생한 large teacher CoT gap \( \Delta_{Large} \)을 비교한 것이다.
■ Table 7은 MATH, GSM8K, AIME, AMC, OlympiadBench에서의 benchmark score이며, Table 9는 Llama3.1-70B 대 Llama3.1-8B, Gemma2-27B 대 Gemma2-9B 등을 포함한 다른 teacher model 쌍의 추가 실험 결과이다.
■ 더 큰 student model large teacher CoT로부터 효과적으로 학습한다. 예를 들어 Qwen2.5-7B-Instruct와 Qwen2.5-32B-Instruct student model은 large teacher CoT로 학습했을 때 평균적으로 5 point 이상 향상된다. 특히 Qwen2.5-32B-Instruct는 AMC benchmark에서 15 point가 넘는 향상을 달성한다.
■ 그러나 더 작은 model은 Qwen2.5-72B-Instruct 같은 large teacher model로부터 효과적으로 학습하지 못한다.
■ 오히려 Qwen2.5-3B-Instruct 같은 small teacher model이 작은 student model에게는 더 좋은 teacher가 될 수 있다. 예를 들어 Qwen2.5-0.5B-Instruct의 경우, large teacher CoT로 학습했을 때 AMC benchmark 성능이 10 point 이상 하락한다.
■ 이 실험에서 large teacher와 small teacher 모두 short CoT data를 생성해 student model을 fine-tuning했다. 두 CoT의 평균 길이에는 유의미한 차이가 없다. 72B teacher가 생성한 CoT의 평균 token length는 432.98이고, 3B teacher가 생성한 CoT의 평균 token length는 440.70이다.
■ 이러한 길이 통제는 large teacher CoT gap 결과에서 CoT length를 confounding variable로 제거하는 데 도움이 된다.
■ 기존 연구들도 stronger models이 반드시 더 좋은 teacher인 것은 아니라는 점을 보여주었으며, response generator나 teacher-side factor를 강조했다. 이 논문은 그 현상을 student model size에
■ 이 섹션의 결론은, small student models은 small teachers로부터 더 잘 학습하는 경향이 있고, large student models은 large teachers로부터 더 큰 이점을 얻는다.
3.4 Analysis of Small Model Learnability Gap
Domain knowledge affects learnability gap
■ math expert models은 model size가 작더라도 general model보다 long CoT data와 large teacher CoT data 모두에서 더 작은 learnability gap을 보인다. (Fig 4)
■ Qwen2.5-Math-1.5B-Instruct와 Qwen2.5-1.5B-Instruct를 비교한 결과로, small math expert model의 long CoT gap이 general small models의 long CoT gap보다 유의미하게 더 작다.
■ 또한 large teacher CoT로 fine-tuning했을 때 Qwen2.5-Math-1.5B의 성능 향상은 Qwen2.5-1.5B의 성능 향상보다 더 크다. 이는 math expert model이 large teacher CoT로부터 더 큰 이득을 얻는다는 점을 시사한다.
■ 이러한 결과를 토대로, 저자들은 Small Model Learnability Gap을 일으키는 핵심 요인 중 하나가 작은 student model의 제한된 in-domain knowledge(이 실험의 경우 수학적 배경지식)라고 추측한다.
■ small model의 제한된(즉, 얕은) domain knowledge가 strong reasoning teacher로부터 학습하는 것을 방해할 수 있다는 것이다.
Distribution Mismatch between student and teacher models
■ Small Model Learnability Gap의 한 가지 기여 요인은 small model이 지나치게 정교한 reasoning traces을 처리하거나 larger teacher에 의해 생긴 distribution shifts에 적응하는 데 어려움을 겪으며, 이것이 궁극적으로 일반화하는 능력을 제한한다는 점이다. (Appendix B.3)
Base models exhibit a more significant learnability gap
■ small base model이 Instruct model보다 뚜렷하게 더 큰 learnability gap을 보인다. (Fig 5)
■ 이는 작은 base model이 long CoT data나 large teacher CoT로부터 효과적으로 학습하는 것이 더 어렵다는 점을 시사한다.
Speaking styles shift
■ 저자들은 Lin et al. (2023)의 방법을 사용하여, long CoT와 large teacher CoT data로 fine-tuning하기 전후 각 token의 rank shift를 평가한다. 이 방법을 통해 저자들은 fine-tuning 과정이 유도한 token distribution shift를 비교할 수 있다.
■ 가장 큰 rank shift를 보인 token들을 most shifted token으로 표시했으며, 그 결과 이러한 token들이 주로 wait, But, Let 같은 표현적 및 문체적 요소와 관련되어 있음을 보여준다. (Appendix C)
■ Fig 8은 most shifted token을 계산하는 과정을 보여준다. 원래 student model이 거의 선택하지 않던 token(But, So, Let)이 fine-tuning 후 강하게 선호되기 시작한다.
■ 이 실험에서 Qwen2.5-3B-Instruct를 student model로 선택했다. Fig 9와 10은 long CoT data로 fine-tuning한 뒤의 most shifted token 결과이다. large teacher CoT data로 fine-tuning한 뒤의 most shifted token 결과는 Figure 11에 볼 수 있다.
■ long CoT와 large teacher CoT는 작은 student의 token distribution을 주로 speaking style과 관련된 token 쪽으로 이동시킨다. Fig 9, 10에서 볼 수 있는 대표적인 표현은 Wait, But, Let, Hmm이고 Fig 11에서 볼 수 있는 대표 표현은 First, Next, For, Now, Thus이다.
■ 즉, long CoT와 large teacher CoT는 small student에게 단순히 더 좋은 reasoning만 전달한 것이 아니라,
■ 특히 small student가 바뀐 부분이 wait, But, Let 같은 표현에 많이 집중되었다는 점은, small model이 complex reasoning의 핵심 논리보다 표면적인 reasoning style을 많이 흡수했을 가능성을 시사한다.
4. Mix Distillation: Bridge Small Model Learnability Gap
■ Small Model Learnability Gap을 해소하기 위한 방법으로 "Mix Distillation"을 제안한다.
4.1 Mix Distillation
■ 이 접근법은 small model이 배우기 쉬운 data와 더 어려운 data를 섞어서, 두 data의 장점을 모두 활용한다.
■ 저자들의 insight는 small model이 자신의 inherent distribution과 더 가까운 data, 예를 들어 short CoT나 small teacher CoT에서 더 잘 작동하는 경향이 있는 반면, distribution shift가 큰 data에서는 어려움을 겪는다는 것이다.
■ long CoT와 large teacher CoT를 섞 data의 token distribution은 small model의 inherent distribution에 더 가까워질 수 있고, 그 결과 small model이 어려운 dataset으로부터 더 효과적으로 학습할 수 있게 될 수 있다.
■ 저자들은 long CoT data에 \( \alpha \)의 가중치를, short CoT data에 \( 1 - \alpha \)의 가중치를 부여하여 결합하는 Mix-Long과 large teacher CoT에 \( \alpha \)의 가중치를, small teacher CoT에 \( 1 - \alpha \)의 가중치를 부여하여 혼합하는 Mix-Large를 제안한다.
■ 이는 small student가 short CoT 또는 small teacher CoT처럼 자신에게 맞는 reasoning distribution을 기반으로 학습하면서도, long CoT 또는 large teacher CoT가 가진 더 풍부한 reasoning pattern을 일부 흡수하도록 만드는 전략이다.
4.2 Experiment Results
■ Qwen2.5-3B-Instruct를 student model로 사용하고 MATH(7.5k)를 training set으로 사용한다.
■ 그리고 baseline으로 다양한 teacher-generated responses을 사용한다. 여기에는 QwQ-32B (long CoT), Qwen2.5-32B (short CoT), Qwen2.5-72B (large teacher CoT), Qwen2.5-3B (small teacher CoT)가 포함된다.
■ 또 다른 long CoT baseline을 만들기 위해 DeepSeek-R1-32B도 teacher로 사용한다.
■ Mix-Long과 Mix-Large 두 설정 모두에서 \( \alpha = 0.2 \)를 사용한다. 즉, samll model이 잘 배울 수 있는 short/small-teacher data를 주축으로 두고, 어려운 data를 20%만 넣었다.
■ 실험 결과, Mix-Long과 Mix-Large는 대부분의 evaluation metric에서 baseline을 능가한다.
■ Mix Distillation은 single dataset으로 학습하는 것과 비교했을 때 small student model의 성능을 향상시킬 수 있다.
■ Mix-Long은 long CoT data만으로 직접 학습한 경우와 비교했을 때, Qwen2.5-3B-Instruct의 MATH와 AMC 성능을 8 point 이상 향상시킨다.
■ Mix-Large는 large teacher CoT data만으로 학습한 경우와 비교했을 때, Qwen2.5-3B-Instruct의 MATH, AIME, AMC에서 7 point 이상의 향상을 보인다.
■ 이 결과는 small student model이 Mix Distillation으로 생성된 dataset으로부터 더 쉽게 학습할 수 있음을 시사한다.
■ long CoT와 short CoT를 섞거나, large teacher CoT와 small teacher CoT를 섞으면, small student model은 두 data source 중 하나만 단독으로 학습했을 때보다 더 좋은 성능을 낼 수 있다.
■ Fig 6은 long CoT 또는 large teacher CoT data의 mix weight \( \alpha \)를 바꿨을 때의 평균 성능을 보여준다.
■ Qwen2.5-3B-Instruct를 student model로 사용할 때, Mix-Long과 Mix-Large 모두에서 \( \alpha = 0.2 \)가 benchmark 평균 성능을 가장 높게 만든다.
■ 아래의 Figure 7은 하나의 수학 문제에 대한 example이다.
■ long CoT와 short CoT를 섞은 뒤, small student model의 output은 branching process 같은 long CoT의 특성은 유지하면서도, 더 짧은 길이를 유지하고 지나치게 복잡한 thinking은 피한다.
■ long CoT data로 fine-tuning된 small student는 반복적인 사고에 압도되어 멈추지 못하는 반면, short CoT data로 fine-tuning된 모델은 잘못된 답을 생성한다. 반대로 Mix-Long은 정답을 낸다.
■ long CoT, short CoT, Mix-Long으로 학습한 모델이 생성한 response의 평균 token length는 각각 3384.7, 575.7, 1248.9다.
■ 이러한 결과는 long CoT data와 short CoT data를 섞어 균형 잡힌 CoT 길이를 만들어 사용하면, 결과적으로 small student model의 reasoning capabilities을 향상시킬 수 있음을 보여준.
