Reasoning Quality Emerges Early: Data Curation for Reasoning Models

Motivation

• Known: 소량의 고품질 long reasoning trace로 SFT를 수행하면 대규모 RL 없이도 LLM의 reasoning 능력을 효과적으로 이끌어낼 수 있으며, 이때 데이터의 diversity와 difficulty가 데이터 품질의 핵심 요소로 알려져 있다.
• Gap: 기존 방법들은 diversity와 difficulty 기반 필터링을 위해 DeepSeek-R1이나 GPT-4o-mini 같은 강력한 reasoning model/LLM judge를 반복적으로 호출해야 하므로 curation 비용이 매우 크고, response-length 등 대리 지표에 의존할 경우 실제 난이도가 아닌 trace 길이나 문체적 특성에 좌우되어 데이터 품질이 suboptimal해지는 문제가 있다.
• Why: reasoning trace 전체(최대 91k 토큰)를 처리하지 않고 초반 토큰만으로 difficulty와 diversity를 판별할 수 있다면, 고비용의 LLM 기반 필터링 없이도 SFT 데이터 curation 비용을 크게 절감하면서 성능은 유지하거나 향상시킬 수 있어, reasoning model 학습 파이프라인의 효율성과 접근성을 크게 높일 수 있다.
• Approach: 사전학습 모델을 무작위로 perturbation한 checkpoint에서 reasoning trace의 첫 100 토큰에 대한 loss를 측정해 difficulty를 판별하고, fine-tuning 궤적을 따라 외삽한 소수의 perturbed checkpoint들에서 첫 1k 토큰의 loss 패턴이 유사한 예제들은 학습 중 유사한 gradient를 유도함을 이론적으로 증명하여 diversity 기반 샘플링에 활용한다.

Achievement

Figure 3. The correlation of different heuristics with difficulty.

1. 저비용 difficulty 판별: reasoning trace 전체가 아닌 첫 100 토큰만으로, 무작위 perturbed pretrained checkpoint에서의 loss를 이용해 어려운 문제를 신뢰성 있게 식별할 수 있음을 실증하였다.
2. 이론적으로 정당화된 diversity 샘플링: fine-tuning 방향으로 외삽한 소수의 noisy checkpoint에서 첫 1k 토큰의 loss가 유사한 예제들이 학습 전 과정에서 유사한 gradient를 유도함을 증명(provable)하여 gradient 유사도를 loss 유사도로 근사할 수 있는 이론적 근거를 제시하였다.
3. 성능 및 효율성 검증: Qwen2.5-7B, Llama3.1-8B를 M23K 의료 reasoning 및 OpenThoughts-Math 데이터셋으로 fine-tuning한 결과, 기존 baseline 대비 최대 1.7% 성능 향상과 91% 더 높은 토큰 효율성을 동시에 달성하였다.

How

Figure 1. Overview of TEMP: Token-Efficient Model Perturbation Reasoning Data Selection. First, using only the first 100

• Problem Understanding Phase(첫 100 토큰)에서의 loss를 측정: 모델이 문제를 이해하고 핵심 포인트를 파악하는 초기 단계의 loss가 difficulty와 강하게 상관됨을 관찰(Fig 2).
• 무작위로 perturbed된 pretrained model checkpoint에서 이 초기 loss를 측정하여 difficulty filtering을 수행.
• Initial Reasoning Steps(첫 1k 토큰)에 대해, fine-tuning 방향으로 외삽한 소수의 directionally perturbed checkpoint들에서 loss를 측정.
• SFT 과정에서 파라미터 변화가 상대적으로 작다는 점(wide/flat basin, Fig 4, 5)에 착안해, 이 loss 값들이 유사한 예제들은 유사한 gradient를 유도함을 증명.
• 이 loss 기반 클러스터링을 통해 diversity sampling을 수행하여 difficult하고 diverse한 SFT 데이터셋을 최종적으로 curation(Early Termination으로 불필요한 계산 절약).
• Qwen2.5-7B, Llama3.1-8B에 대해 M23K, OpenThoughts-Math 데이터셋으로 실험하여 기존 LLM 기반 필터링 baseline과 비교.

Originality

• 기존 연구들이 difficulty/diversity 판별에 강력한 LLM judge나 전체 reasoning trace를 요구했던 것과 달리, 첫 100~1k 토큰만으로 이를 판별할 수 있다는 새로운 경험적 관찰(problem understanding phase)을 제시함.
• loss 유사도와 gradient 유사도 사이의 관계를 SFT의 flat/wide basin 특성과 연결하여 이론적으로 증명한 점이 기존 gradient/embedding 기반 diversity 방법론과 차별됨.
• randomly perturbed checkpoint와 fine-tuning 방향으로 외삽한 directionally perturbed checkpoint라는 두 종류의 perturbation을 각각 difficulty와 diversity 판별에 결합적으로 활용하는 파이프라인(TEMP) 설계가 독창적임.

Limitation & Further Study

• 제안 방법의 이론적 보장은 SFT 과정에서 파라미터 변화가 작다는 가정(flat/wide basin)에 기반하므로, 파라미터 변화가 큰 RL 기반 post-training이나 매우 긴 SFT 학습 스케줄에는 적용성이 제한적일 수 있음.
• perturbed checkpoint의 개수, perturbation 크기, 외삽 방향 등 하이퍼파라미터에 대한 민감도 분석과 다양한 모델 크기·아키텍처로의 일반화 검증이 추가로 필요함.
• 실험이 M23K 의료 reasoning과 OpenThoughts-Math라는 두 도메인에 한정되어 있어, 코딩·과학 등 더 폭넓은 reasoning 도메인에서의 검증이 후속 연구로 요구됨.
• 1.7%라는 성능 향상 폭이 크지 않아, 다양한 baseline과 통계적 유의성 검증이 추가로 뒷받침되면 주장이 더 견고해질 것임.

Evaluation

총평: reasoning trace의 초기 토큰만으로 difficulty와 diversity를 저비용으로 판별할 수 있다는 참신하고 실용적인 통찰을 이론적 증명과 실증 실험으로 뒷받침한 우수한 연구이나, 이론의 가정 범위와 도메인 다양성 측면에서 추가 검증이 필요하다.