
LLM을 특정 목적에 맞게 훈련시키는 **Post-training** 과정은 크게 두 가지 갈래로 나눌 수 있다. 하나는 학생 모델 스스로의 경험을 바탕으로 학습하는 **On-policy** 방식이고, 다른 하나는 교사 모델의 정답을 모방하는 **Off-policy** 방식이다. 기존의 방식들은 각기 뚜렷한 장단점을 가진다. **On-Policy Distillation(혹은 GKD)** 은 이 두 가지 방식의 장점만을 결합하여 학습 효율을 극대화한 방법론이다.

---

## 1. 기존 Post-training 방식의 딜레마

모델의 Post-training 방식은 크게 두 가지로 나뉜다. 두 방식 모두 명확한 한계를 가진다.

### Off-policy Distillation (SFT)

고성능 교사 모델(Teacher model)이 미리 생성해 둔 완벽한 경로(Trajectory)와 중간 사고 과정(Thinking steps)을 학생 모델이 그대로 따라가며 출력 분포(Output distribution)를 맞추도록 학습하는 방식이다. 매 토큰마다 전체 확률 분포(Logit distillation)를 참고하기도 한다.

- **장점:** 매 토큰마다 정답을 알려주므로 밀도 높은 지도(Dense supervision)가 가능하다.
- **단점:** 노출 편향(Exposure bias) 문제가 발생한다.
- **비유:** 학생이 일타 강사의 완벽한 풀이 과정을 눈으로 보며 베껴 쓰는 것과 같다. 교사가 자주 마주하는 상황(Context)에만 노출되다 보니, 실전에서 학생 본인만의 사소한 실수로 교사의 궤도에서 벗어났을 때 이를 스스로 수습하고 정답으로 되돌아오는 방법을 전혀 모르게 된다.

### On-policy RL (강화학습)

학생 모델(Student model)이 스스로 텍스트를 생성하며(Rollout/Trajectory), 그 결과에 대해 보상을 받는 방식이다.

- **장점:** 학생 모델이 스스로 만들어낸 분포 안에서 학습하므로, 자신의 실수로부터 궤도를 수정하고 복구하는 방법을 직접적으로 학습한다.
- **단점:** 피드백이 너무 희소(Sparse feedback)하다.
- **비유:** 모델이 복잡한 수학 문제를 풀 때, 10단계의 풀이 과정을 거쳐 '21'이라는 오답을 냈다고 가정해 보자. RL 기반 방식에서는 그저 "최종 답이 틀렸다"라는 결과만 알려줄 뿐, 10단계 중 어느 단계에서 어떤 치명적인 논리적 오류가 있었는지는 토큰 단위로 알려주지 못한다.

---

## 2. On-Policy Distillation의 핵심 아이디어

On-policy distillation은 위 두 가지 방법의 장점을 융합한다.

> **핵심 원리:** 학생 모델이 스스로 텍스트 경로를 생성하도록 둔다(On-policy). 단, 고성능 교사 모델이 그 생성된 경로의 매 토큰(Token)마다 실시간으로 점수를 매겨 지도한다(Distillation).

앞선 수학 문제의 예시로 돌아가 보자. On-policy distillation 환경에서는 학생 모델이 스스로 자기 방식대로 문제를 풀어나간다. 이때 훌륭한 과외 선생님(Teacher model)이 옆에 밀착해서 채점을 진행한다. 학생이 오답으로 향하는 치명적인 실수를 한 토큰에는 즉각적으로 강한 감점을 주고, 올바른 논리를 전개한 토큰에는 보상을 준다. 결과적으로 학생은 "자신의 행동 반경(자신의 분포)" 안에서 실수를 복구하는 법을 배우면서도, "밀도 높은 피드백(Dense reward signal)"을 받게 된다.

---

## 3. 체스로 이해하는 비유

이 세 가지 개념의 차이는 '체스 학습법'에 비유하면 아주 쉽게 이해된다.

> **Off-policy Distillation** = "그랜드마스터의 대국 관전하기"
>
> 엄청나게 훌륭한 수를 보며 배우지만, 실전에서 초보자인 내가 이상한 위기 상황에 처했을 때 어떻게 탈출해야 하는지는 배울 수 없다.

> **On-policy RL** = "코치 없이 혼자 체스 두기"
>
> 내가 직접 두면서 감을 익히지만(On-policy), 피드백은 경기가 끝난 후 '승/패'로만 주어진다. 어느 수가 패배의 원인이었는지 아무도 짚어주지 않는다.

> **On-Policy Distillation** = "내 경기를 보며 매 수마다 평가해 주는 개인 코치"
>
> 내가 직접 둔 엉망진창인 수(**Trajectory**)를 보면서, 코치(Teacher)가 매 수마다 "이건 치명적 실수야!", "이건 훌륭한 수야!"라고 실시간으로 채점(**Dense Reward**)해 준다. 나의 실제 수준에 맞는 상황에서 매우 촘촘한 피드백을 받게 된다.

---

## 4. 기술적 디테일: Reverse KL Divergence

On-Policy Distillation의 핵심은 학생 모델($\pi_\theta$)이 스스로 생성한 Trajectory 상의 **모든 토큰마다 Teacher 모델($\pi_{teacher}$)이 즉각적인 피드백을 준다는 것**이다.

이때 피드백을 계산하는 핵심 손실 함수(Loss function)로 **Reverse KL Divergence**를 사용한다.

$$KL(\pi_\theta \| \pi_{teacher}) = \mathbb{E}_{x \sim \pi_\theta} \left[ \log \pi_\theta(x_{t+1} \mid x_{1..t}) - \log \pi_{teacher}(x_{t+1} \mid x_{1..t}) \right]$$

### 동작 원리

학생 모델이 직접 샘플링한 Trajectory $x \sim \pi_\theta$ 위에서, 각 시점 $t$마다 "Teacher 모델이라면 이 상황에서 다음에 어떤 토큰을 뱉을 것인가?"를 계산하여 학생의 확률 분포를 Teacher에 맞추도록 강제한다.

### 왜 Reverse KL인가?

일반적인 SFT(Forward KL)는 Teacher의 모든 가능성을 평균적으로 따라가려 하지만, Reverse KL은 **Mode-seeking** 특성을 갖는다. 즉, 학생 모델이 여러 어설픈 옵션 사이에서 확률을 분산시키며 방황하지 않고, Teacher의 가장 확실하고 바람직한 행동(Mode) 하나를 집중적으로 좇도록 유도하여 노이즈를 줄여준다.

---

## 5. On-Policy Distillation의 주요 이점

**1. 강력한 학습 신호 (Dense Reward)**

결과물 끝에서만 보상을 받는 일반적인 On-policy RL과 달리, 모든 토큰(추론 단계)마다 보상이 주어진다. 복잡한 추론 과정 중 정확히 어느 지점에서 경로를 이탈했는지 짚어줄 수 있다.

**2. 실수 극복 능력 향상**

자신이 직접 만든 낯선 상황(On-policy)에서 출발하여 다시 정답 경로를 찾아가는 과정을 학습하므로, 긴 문맥의 작업에서도 성능 붕괴(**Compounding Error**) 현상을 방지한다.

**3. 압도적인 계산 효율성 (Compute Efficiency)**

Teacher 모델이 직접 전체 텍스트를 끝까지 샘플링할 필요 없이, 학생이 만든 Trajectory에 대한 로그 확률(Logprobs)만 한 번 계산해 주면 된다. 별도의 보상 모델(Reward Model)도 필요하지 않다. 이로 인해 기존 SFT 방식 대비 **FLOPs 기준 9배에서 최대 30배 높은 효율**을 보여준다.


On-Policy Distillation 개념 정리

참고 자료