Where Did This Sentence Come From? Tracing Provenance in LLM Reasoning Distillation

1. 연구 배경 및 문제 제기 (Problem)

• 배경: 대형 언어 모델 (LLM) 의 추론 능력을 소규모 모델로 전이하는 '추론 증류 (Reasoning Distillation)' 기법이 주목받고 있습니다. 이는 일반적으로 대형 '교사 모델 (Teacher Model)'이 생성한 추론 경로를 '학생 모델 (Student Model)'이 학습하도록 하는 방식으로 수행됩니다.
• 문제: 기존 연구들은 증류된 모델의 성능 향상 자체에 집중했으나, 학생 모델이 실제로 교사 모델의 추론 논리를 계승했는지, 아니면 단순히 기존에 가지고 있던 출력 패턴으로 회귀 (regress) 했는지에 대한 심층적인 분석이 부족했습니다.
• 핵심 질문: 테스트 시간 (Test-time) 에 새로운 문맥에서 학생 모델이 일관되게 교사 모델의 행동을 따를 수 있는지, 아니면 원래의 출력 분포로 돌아갈 것인지에 대한 불확실성이 존재합니다. 이는 증류 모델의 일반화 능력에 대한 의문을 제기합니다.

2. 제안된 방법론: Reasoning Distillation Provenance Tracing (방법론)

저자들은 증류된 모델의 각 행동 (문장) 이 어디에서 기원했는지를 추적하는 새로운 프레임워크인 **'Reasoning Distillation Provenance Tracing'**을 제안합니다.

• 핵심 아이디어: 증류된 모델이 생성한 각 문장 (Action) 에 대해, 동일한 문맥에서 교사 모델 ($M_T$), 원본 학생 모델 ($M_S$), 증류된 모델 ($M_D$) 이 부여한 예측 확률 (Predictive Probabilities) 을 비교합니다.
• 행동 분류 (4 가지 카테고리):
  1. Teacher-originated (Teacher Sentence): 교사 모델이 학생 모델보다 훨씬 높은 확률을 부여한 행동. (증류를 통해 전이된 지식)
  2. Student-originated (Student Sentence): 학생 모델이 교사 모델보다 훨씬 높은 확률을 부여한 행동. (학생의 기존 패턴)
  3. Shared (Shared Sentence): 세 모델 모두 유사한 확률을 부여한 행동. (증류 전부터 존재하던 공통 패턴)
  4. Boosted (Boosted Sentence): 증류 전에는 세 모델 모두 유사했으나, 증류 후 학생 모델의 확률이 크게 향상된 행동. (증류를 통해 강화된 학생 내부 패턴)
• 분석 프로세스:
  1. 테스트 데이터에서 증류된 모델의 추론 경로 (Trajectory) 를 샘플링합니다.
  2. 해당 경로를 교사, 원본 학생, 증류된 모델에 다시 입력하여 각 토큰 (또는 문장 단위) 의 확률을 계산합니다.
  3. 확률 차이 ($\Delta$) 를 기반으로 위 4 가지 카테고리로 분류하여 각 행동의 기원을 규명합니다.

3. 주요 분석 결과 (Key Findings)

Open-source 증류 모델 (DeepSeek-Distill-Qwen-7B, DeepSeek-R1-0528-Qwen3-8B, LIMO-v2 등) 을 AIME24 와 GPQA-D 벤치마크에서 분석한 결과 다음과 같은 현상을 발견했습니다.

• 교사 기원 행동의 중요성: 증류된 모델은 새로운 테스트 환경에서도 Teacher Sentence를 생성할 수 있으며, 이는 정답 (Correct Answer) 과 높은 상관관계를 보입니다. 특히 추론 초기 단계에서 Teacher Sentence 의 비율이 높을수록 성능이 좋습니다.
• 학생 내부 패턴의 활성화: 증류는 학생 모델 내부에 잠재되어 있던 추론 지식 (Boosted Sentence) 을 활성화시킵니다. 그러나 모든 활성화가 유익한 것은 아니며, 모델 크기나 추론 단계에 따라 부정적인 영향을 줄 수도 있습니다.
• 일반화 메커니즘: 증류의 성능 향상은 단순히 학생 모델의 기존 패턴을 강화하는 것이 아니라, 교사 모델의 고유한 추론 패턴 (Teacher Sentence) 을 성공적으로 전이하고 유지하기 때문에 발생함을 입증했습니다.

4. 제안된 전략: Teacher-Guided Data Selection (학습 가이드)

분석 결과를 바탕으로, 증류 학습의 효율성을 높이기 위한 교사 주도 데이터 선택 (Teacher-Guided Data Selection) 전략을 제안합니다.

• 기존 방법의 한계: 기존 연구 (예: GRAPE) 는 학생 모델의 원래 분포와 가장 잘 맞는 데이터를 선택하는 휴리스틱 방식을 사용했습니다.
• 제안 방법: 교사 모델과 학생 모델 간의 **발산 (Divergence)**을 기준으로 데이터를 선택합니다.
  • 교사 모델이 학생 모델보다 확률이 훨씬 높은 Teacher Sentence가 포함된 샘플을 우선적으로 선택합니다.
  • 이는 증류 과정에서 학생 모델이 교사 모델의 추론 능력을 더 많이 흡수하도록 유도합니다.
• 비용 효율성: 추가적인 생성 과정 없이, 기존 데이터에 대해 모델의 로그 확률 (Logits) 을 한 번만 계산하여 필터링할 수 있어 계산 비용이 낮습니다.

5. 실험 결과 (Results)

다양한 교사 모델 (DeepSeek-R1-671B, QwQ-32B, GPT-OSS-120B) 과 학생 모델 (Qwen 시리즈) 조합으로 실험을 수행했습니다.

• 성능 향상: 제안한 데이터 선택 전략은 Random 선택 (Vanilla) 및 기존 방법 (GRAPE) 보다 AIME24, AIME25, MATH500, OlympiadBench 등 주요 추론 벤치마크에서 평균 1.7% ~ 2.5% 의 정확도 향상을 보였습니다.
• 영역 일반화: 수학 데이터로 학습된 모델이 과학 (OpenScienceReasoning) 데이터로 학습되었을 때도, Teacher Sentence 비율이 높은 데이터를 선택하는 전략이 수학 및 과학 영역 모두에서 우수한 일반화 성능을 보여주었습니다.
• 분석의 유효성: 증류 후 모델의 추론 경로를 다시 분석했을 때, 제안된 방법을 사용한 모델은 Teacher Sentence 생성 확률이 증가하고 Student Sentence 확률이 감소하여, 교사 모델의 행동 패턴을 더 잘 따르고 있음을 확인했습니다.

6. 의의 및 기여 (Significance & Contributions)

1. 새로운 분석 프레임워크: 증류된 모델의 출력 기원을 미세하게 분류 (Teacher/Student/Shared/Boosted) 하여, 증류가 실제로 추론 능력을 전이했는지 아니면 기존 패턴을 강화했는지를 정량적으로 분석할 수 있는 Provenance Tracing을 최초로 제안했습니다.
2. 이론적 통찰: 증류의 일반화 성능 향상이 '교사 모델의 추론 패턴 (Teacher Sentence)'의 성공적인 전이와 밀접하게 연관되어 있음을 실험적으로 증명했습니다.
3. 실용적 기여: 추론 증류의 효율성을 높이기 위해, 학생 모델의 분포가 아닌 교사 - 학생 간 차이 (Divergence) 를 기반으로 한 데이터 선택 전략을 제시하여, 더 적은 데이터로도 높은 성능을 달성할 수 있는 방법을 제시했습니다.

이 논문은 증류된 모델의 '블랙박스'를 열어 그 내부 작동 원리를 해석하고, 이를 바탕으로 더 효과적인 학습 전략을 수립할 수 있는 토대를 마련했다는 점에서 의의가 큽니다.