Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes

Each language version is independently generated for its own context, not a direct translation.

이 논문은 거대한 언어 모델 (LLM) 을 더 똑똑하게 만드는 '지식 전수 (Distillation)' 과정에서 발생한 흥미로운 문제와 그 해결책을 다룹니다.

비유하자면, 이 논문은 **"어린 학생 (Student) 이 스승 (Teacher) 의 가르침을 받으며 스스로 문제를 풀 때, 왜 가끔 엉뚱한 방향으로 헤매게 되는가?"**를 분석하고, **"더 안정적인 학습 방법"**을 제안합니다.

핵심 내용을 일상적인 언어와 비유로 설명해 드릴게요.

1. 배경: 왜 '온 - 정책 (On-Policy)' 학습이 필요한가?

과거에는 AI 를 가르칠 때, **스승이 미리 정해둔 정답지 (고정된 데이터)**만 보여주며 학습시켰습니다. 하지만 요즘처럼 복잡한 문제 (수학 추론, 에이전트 작업 등) 를 풀 때는 상황이 다릅니다.

상황: 학생이 문제를 풀다 보면, 스승이 경험해 본 적 없는 완전히 새로운 길로 가게 됩니다.
문제: 이때 스승이 "이건 틀렸어"라고만 말해주면, 학생은 길을 잃고 망설이게 됩니다.
해법 (OPD): 그래서 학생이 **스스로 만들어낸 답안 (Rollout)**을 스승에게 보여주고, 그 순간순간의 피드백을 받아 학습하는 '온 - 정책 (On-Policy)' 방식이 주목받게 되었습니다.

2. 문제점: "단 하나의 단어"에 모든 것을 걸다

기존의 온 - 정책 학습 방식은 아주 단순하고 효율적이었습니다. 하지만 세 가지 치명적인 약점이 있었습니다.

① 편향된 신호 (Imbalanced Signal)

비유: 시험지를 채점할 때, 스승이 "정답인 단어 하나만" 보고 점수를 매기는 것과 같습니다.
현실: 학생이 쓴 문장 100 개 중 99 개는 평범하고, 1 개만 스승이 좋아하는 단어가 나왔다면, 그 1 개 때문에 전체 학습 방향이 크게 흔들립니다. 마치 "한 번의 칭찬에 취해서, 나머지 99 개의 실수를 무시하고 그 단어만 반복하는 학생"과 같습니다.

② 스승의 눈이 멀다 (Unreliable Guidance)

비유: 학생이 스승이 가본 적 없는 낯선 길로 들어섰을 때, 스승은 "아, 이 길은 내가 잘 모르는 길이니까 그냥 '괜찮아'라고 말해줄 뿐"입니다.
현실: 학생이 엉뚱한 방향으로 (예: 같은 문장을 반복하거나, 의미 없는 말을 늘어놓는 것) 가도, 스승은 그 순간의 단어만 보고 "좋아, 이 단어는 내가 좋아해"라고 점수를 줍니다. 결과적으로 학생은 **엉뚱한 반복 (Repetition Loop)**이나 **의미 없는 말 (Gibberish)**을 계속하게 됩니다.

③ 언어 장벽 (Tokenizer Mismatch)

비유: 스승과 학생이 단어를 쪼개는 방식이 다릅니다.
- 학생: "사과"를 "사", "과"로 쪼갭니다.
- 스승: "사과"를 "사과" 하나로 쪼갭니다.
현실: 학생이 "사"라고 말했을 때, 스승은 "아, 내가 원하는 '사과'의 일부가 아니야"라고 오해해서 점수를 깎아줍니다. 실제로는 같은 뜻인데, **단어 쪼개기 방식 (토크나이저)**이 달라서 불필요한 오해를 사는 것입니다.

3. 해결책: "스승이 좋아하는 단어 10 개"를 함께 보자

저자들은 이 문제를 해결하기 위해 **"스승이 가장 좋아할 만한 상위 K 개 단어 (Top-K)"**를 함께 고려하는 새로운 방식을 제안했습니다.

기존 방식: 학생이 쓴 단 하나의 단어만 보고 "좋다/나쁘다"를 판단.
새로운 방식 (Local Support Matching):
1. 학생이 문장을 이어갈 때, 스승이 "이런 단어들을 쓸 수도 있겠다"라고 생각할 만한 상위 K 개 후보군을 먼저 뽑습니다.
2. 학생이 그 후보군 안에 있는 단어를 썼다면, 단 하나의 단어에 매달리지 않고 그 후보군 전체의 확률 분포를 비교합니다.
3. 결과: 학생이 엉뚱한 길로 빠지지 않도록, 스승이 "이쪽 방향은 괜찮아"라고 알려주는 여러 개의 안전지대를 만들어주는 것입니다.

4. 실험 결과: 더 안정적이고 똑똑해졌다

이 방식을 적용했을 때 두 가지 큰 변화가 있었습니다.

학습이 더 안정적: 학생이 엉뚱한 길 (반복, 의미 없는 말) 로 빠지지 않고, 스승의 지도를 더 정확하게 따라갑니다.
실제 성능 향상: 수학 문제 풀이와 복잡한 에이전트 작업 (예: 로봇이 집안일 하기) 에서 기존 방식보다 훨씬 좋은 점수를 받았습니다.

5. 한 줄 요약

"기존 방식은 학생이 쓴 '단 하나의 단어'만 보고 점수를 매겨서 엉뚱한 방향으로 헤매게 했지만, 새로운 방식은 '스승이 좋아할 만한 여러 후보'를 함께 보며 더 넓은 시야로 지도해줘서, 학생이 더 안정적이고 똑똑하게 성장하게 했다."

이 논문은 거대 AI 가 스스로 학습할 때 발생하는 **'오해'와 '불안정성'**을 해결하여, 더 똑똑하고 신뢰할 수 있는 AI 를 만드는 데 중요한 디딤돌이 된 연구입니다.

Each language version is independently generated for its own context, not a direct translation.

1. 문제 정의 (Problem Definition)

**온-폴리시 증류 (On-Policy Distillation, OPD)**는 대규모 언어 모델 (LLM) 의 후속 학습 (post-training) 에 매력적인 접근법으로, 고정된 교사 (Teacher) 데이터가 아닌 **학생 모델이 생성한 롤아웃 (rollouts)**에 대해 교사 모델의 피드백을 평가합니다. 이는 긴 시간 범위 (long-horizon) 의 추론이나 에이전트 학습에서 학생이 교사 데이터에 없는 접두사 (prefix) 에 도달할 때 특히 유용합니다.

그러나 현재 널리 사용되는 샘플링된 토큰 기반 (sampled-token) OPD 방식은 다음과 같은 이유로 긴 시퀀스 학습에서 취약하고 불안정합니다:

불균형한 신호: 분포 수준의 차이를 단일 토큰의 로그 비율 (log-ratio) 로 축소하여, 학습 신호가 극도로 불균형해집니다.
불신뢰한 교사 신호: 학생이 생성한 접두사가 교사에게 흔하지 않은 경우 (Out-of-Distribution), 교사는 그럴듯해 보이는 토큰에 높은 확률을 부여할 수 있어, 전체 궤적의 품질이 떨어지더라도 잘못된 학습 신호를 제공합니다.
토크나이저 불일치: 학생과 교사의 토크나이저가 다르거나 특수 토큰 처리가 다를 경우, 의미적으로 동일한 내용이더라도 토큰 매칭 실패로 인해 잘못된 패널티를 받습니다.

이러한 문제들은 학습이 불안정해지거나, 모델이 의미 없는 반복 (repetition loops) 이나 비합리적인 생성을 계속하도록 유도하는 '보상 해킹 (reward hacking)' 현상을 초래합니다.

2. 방법론 (Methodology)

저자들은 이 문제를 해결하기 위해 **교사 Top-K 로컬 서포트 매칭 (Teacher Top-K Local Support Matching)**을 제안합니다. 이는 단일 토큰 비교를 대체하여, 각 접두사 (prefix) 에서 교사가 유망하다고 판단하는 토큰 집합 (Support Set) 내에서 학생과 교사의 분포를 비교하는 방식입니다.

주요 구성 요소는 다음과 같습니다:

Truncated Reverse-KL Objective:
- 기존 샘플링된 단일 토큰 $y_t$ 에 대한 로그 비율 대신, 교사의 확률이 높은 Top-K 토큰 집합 $S(c_{i,t})$ 내에서 학생과 교사의 분포를 비교합니다.
- 목적 함수는 이 잘린 (truncated) 서포트 내에서 재규격화된 분포 $\hat{\pi}_\theta$ 와 $\hat{q}$ 에 대한 역 KL 발산 (Reverse-KL) 을 평균화합니다.
- 수식: $L_{LSM} = \mathbb{E} [\frac{1}{\sum |o_i|} \sum_{t} \sum_{v \in S} \hat{\pi}_\theta(v) \log \frac{\hat{\pi}_\theta(v)}{\hat{q}(v)}]$
실용적 안정화 기법 (Practical Stabilization Choices):
1. 서포트 재규격화 (Support-set Renormalization): 잘린 서포트 내에서만 확률 질량을 재규격화하여, 전체 어휘가 아닌 부분 집합 내에서의 비교가 가능하도록 합니다.
2. Top-p 롤아웃 샘플링: 극단적으로 낮은 확률의 토큰이 생성되는 것을 방지하여, 교사의 신호가 신뢰할 수 있는 영역에서 학습이 이루어지도록 합니다.
3. 특수 토큰 마스킹 (Special-token Masking): 토크나이저 불일치로 인한 잘못된 패널티를 줄이기 위해 특수 토큰을 마스킹합니다.

3. 주요 기여 (Key Contributions)

추정기 트레이드오프 분석:
- 토큰 수준 OPD 는 시퀀스 수준 역 KL 에 비해 편향 (bias) 이 있지만, 시퀀스 길이에 따른 분산 (variance) 증가가 훨씬 느리다는 것을 이론적으로 증명했습니다.
- 시퀀스 수준 (Sequence-level) 접근은 미래 보상에 대한 강한 커플링으로 인해 분산이 급격히 증가 ( $O(T^4)$ ) 하여 긴 시퀀스 학습에서 불안정해지는 반면, 토큰 수준은 분산이 $O(T^2)$ 로 더 안정적입니다.
- 제안된 방법은 이 두 극단 사이의 균형을 찾습니다.
실패 모드 식별:
- 샘플링된 토큰 OPD 의 세 가지 주요 실패 원인 (불균형한 신호, 학생 생성 접두사에서의 교사 신호 불신뢰, 토크나이저 불일치) 을 실증적으로 규명했습니다.
새로운 목적 함수 제안 및 검증:
- Top-K 로컬 서포트 매칭을 통해 단일 토큰 의존성을 제거하고 분포 수준의 비교를 도입했습니다.
- 수학 추론 (Math Reasoning) 과 에이전트 작업 (Agentic Tasks) 을 포함한 다양한 시나리오에서 기존 OPD 보다 안정적인 최적화와 더 나은 성능을 입증했습니다.

4. 실험 결과 (Results)

실험은 Qwen2.5-7B-Instruct 를 학생 모델로, OpenThinker3-7B 를 교사 모델로 사용하여 수행되었습니다.

단일 작업 수학 추론 (Single-task Math Reasoning):
- 제안된 방법 (Local Support Matching) 은 샘플링된 토큰 OPD 보다 평균 점수를 크게 향상시켰습니다 (예: AIME24 에서 10.0 → 23.3, Minerva 에서 32.4 → 34.9).
- 특수 토큰 마스킹만 적용한 베이스라인도 개선되었으나, 제안된 방법은 마스킹 여부와 상관없이 더 일관된 성능 향상을 보였습니다.
다중 작업 학습 (Multi-task Agentic-plus-Math Training):
- 수학 작업과 에이전트 작업 (ALFWorld) 을 번갈아 학습하는 환경에서, 제안된 방법은 수학 성능을 크게 향상시키면서도 에이전트 성능을 유지하거나 개선했습니다.
- 특히 긴 시퀀스 학습에서 토큰 수준 감독이 취약했던 수학 영역에서 가장 큰 개선을 보였습니다.
학습 동역학 (Training Dynamics):
- 더 낮은 분산: 제안된 방법은 그라디언트 노름 (Gradient Norm) 이 작고, 클리핑 (clipping) 비율이 낮아 최적화가 더 안정적입니다.
- 더 나은 정렬: 학생과 교사의 로그 확률 차이 (log-probability gap) 가 줄어들어, 모델이 교사의 의도된 분포에 더 잘 정렬됨을 보여줍니다.
- 안정적인 학습 곡선: 학습 초기부터 평가 성능까지 꾸준히 향상되는 경향을 보였습니다.

5. 의의 및 결론 (Significance)

이 논문은 긴 시퀀스 학습에서 온-폴리시 증류의 핵심적인 한계를 규명하고, 이를 해결하기 위한 실용적이고 효과적인 대안을 제시했습니다.

이론적 통찰: 시퀀스 수준의 완전한 역 KL 과 단일 토큰 기반 OPD 사이의 편향 - 분산 트레이드오프를 명확히 하고, 긴 시퀀스 학습에서는 분산 제어가 더 중요함을 강조했습니다.
실용적 해결책: 복잡한 알고리즘 변경 없이, Top-K 서포트 매칭과 Top-p 샘플링이라는 간단한 수정만으로도 학습의 안정성과 성능을 획기적으로 개선할 수 있음을 보였습니다.
미래 방향: 제안된 방법은 교사 - 학생 간 분포 이동 (distribution shift) 을 완화하고, 보상 해킹을 방지하는 데 기여하며, 향후 LLM 의 추론 및 에이전트 학습을 위한 표준적인 증류 기법으로 자리 잡을 잠재력을 가집니다.

결론적으로, 이 연구는 "단일 토큰"에 의존하는 기존의 OPD 방식이 가진 취약점을 극복하고, 분포 기반의 로컬 비교를 통해 장기적인 추론 작업에서 더욱 견고한 학습을 가능하게 하는 중요한 전환점을 제시합니다.