교사의 확률에서 학습 신호 얻기: OPD

앞 글에서는 과제 보상에서 어드밴티지를 구하는 두 방법을 살펴봤습니다. 같은 질문의 다른 응답들과 비교하거나, critic의 미래 보상 예측과 비교하는 방법이었습니다. 이번에는 학습 신호의 출처를 확장해보겠습니다. 더 잘하는 교사 모델이 있다면, 학생이 선택한 각 토큰을 교사는 얼마나 선호하는지 읽을 수 있습니다.

On-policy distillation(OPD)은 학생이 직접 만든 경로에서 교사의 확률을 이용해 학습하는 방법입니다. 먼저 왜 학생의 경로를 사용하는지 살펴보고, 교사와 학생의 확률 차이를 작은 숫자로 계산하겠습니다. 이어서 어떤 확률로 차이를 가중하는지에 따라 forward KL과 reverse KL을 구별하고, 마지막에는 Miles v0.1.0에서 이 신호를 과제 어드밴티지에 더하는 방식과, 과제 신호 없이 교사만으로 학습하는 경우를 연결하겠습니다.

학생이 실제로 만드는 문맥에서 배우기

주어진 답을 따라 학습하는 SFT를 먼저 떠올려보겠습니다. 프롬프트 P에 대한 학습용 답이 u → a → c라는 세 토큰이라고 하겠습니다. 학생은 P에서 u를, P + u에서 a를, P + u + a에서 c를 예측하도록 학습합니다. 답의 출처는 사람일 수도 있고 다른 모델일 수도 있습니다.

그런데 실제로 생성하게 했더니 학생이 두 번째 토큰에서 a 대신 b를 골랐다고 해보겠습니다. 세 번째 선택의 문맥은 P + u + b가 됩니다. 주어진 답만 따라갈 때 보았던 P + u + a와는 다른 문맥입니다. 생성이 길어질수록 학생이 스스로 방문하는 문맥과 학습용 답의 문맥이 더 달라질 수 있습니다.

OPD에서는 학생이 u → b → d를 생성했다면 그 경로를 그대로 사용합니다. 교사에게도 학생이 실제로 거친 문맥을 주고, 그때 학생이 선택한 토큰의 확률을 계산하게 합니다.

그림의 OPD 경로에 있는 pT(d)는 “교사가 자기 답을 작성할 때 d를 얼마나 자주 쓰는가”가 아닙니다. 학생의 문맥 P + u + b가 주어졌을 때 교사가 d에 부여하는 확률입니다. 학생이 흔히 만드는 실수나 우회 경로도 교사의 신호를 얻는 입력이 됩니다. 이것이 주어진 답의 경로만 따라가는 학습과의 차이입니다. Thinking Machines — On-Policy Distillation

그림의 연결선은 문맥별 평가 관계를 나타냅니다. 매 토큰을 생성할 때마다 교사 서버에 한 번씩 통신해야 한다는 뜻은 아닙니다. 학생이 생성한 토큰 이력을 모아 교사에 넣고, 각 위치의 로그확률을 순전파로 함께 구할 수 있습니다.

선택한 토큰의 확률을 비교하기

응답 u → b → d의 두 번째 토큰 b를 따라가보겠습니다. 두 모델에 똑같이 P + u를 넣습니다. 학생은 b에 0.2, 교사는 0.4의 확률을 부여합니다. 각 모델은 자기 가중치로 히든 스테이트, LM head의 점수, softmax 분포를 계산합니다. 그 두 출력에서 같은 토큰 ID인 b의 확률을 꺼내 비교합니다.

이번 글에서는 선택한 토큰 하나의 로그확률을 비교하는 방식을 사용하겠습니다. 학생 확률은 학습 신호를 만들기 위해 이번 업데이트 전에 고정한 값입니다. 토큰별 신호 dₜ를 다음처럼 계산합니다. ln은 자연로그입니다.

dₜ = ln(교사 확률 / 고정한 학생 확률)

P+u에서 b:   교사 0.4, 학생 0.2 → d₂ ≈ +0.693
P+u+b에서 d: 교사 0.2, 학생 0.4 → d₃ ≈ −0.693

다음 위치에서는 문맥이 P + u + b이고 선택 토큰은 d입니다. 학생 확률은 0.4, 교사 확률은 0.2라서 ln(0.2 / 0.4) ≈ −0.693이 됩니다. dₜ는 증류 신호를 나타내는 기호이고, 토큰 이름 d는 설명용 어휘 항목입니다. 교사가 읽는 앞부분은 교사가 따로 작성한 답이 아니라 학생이 선택했던 토큰입니다.

이 값을 2편의 어드밴티지처럼 사용하면, +0.693은 선택 확률을 높이는 방향으로, −0.693은 낮추는 방향으로 작용합니다. 교사가 학생의 선택에 더 높은 확률을 주었는지 낮은 확률을 주었는지가 부호를 정합니다.

여기에서 교사가 선호한다는 것과 실제 과제의 정답이라는 것은 구별해야 합니다. 교사도 틀릴 수 있고, 잘못된 앞 문맥이 주어지면 그 문맥을 자연스럽게 이어 쓰는 토큰에 높은 확률을 줄 수도 있습니다. dₜ는 토큰의 참 기여도나 정답성 판정이 아니라 두 모델의 확률 차이에서 얻은 학습 신호입니다.

어느 분포를 기준으로 비교하는가

앞에서는 선택한 토큰 하나의 확률을 비교했습니다. 이제 같은 문맥에서 가능한 다음 토큰 전체로 시야를 넓혀보겠습니다. KL divergence는 두 확률 분포의 차이를 하나의 수로 나타냅니다. 교사를 기준으로 차이를 평균내는가, 학생을 기준으로 평균내는가에 따라 값이 달라집니다.

이번에는 별도의 문맥 P + v에서 후보가 a, b, c 세 개뿐인 작은 예시를 보겠습니다. 교사의 확률은 49%, 49%, 2%이고 학생은 78%, 2%, 20%입니다. 두 모델에 넣는 문맥은 같으며, 그림의 왼쪽과 오른쪽도 같은 두 분포를 사용합니다.

Forward KL은 교사 확률로 가중합니다. 교사가 자주 선택하는 토큰의 차이가 큰 비중으로 들어갑니다. 그림의 b는 교사 확률이 49%인데 학생 확률은 2%뿐입니다. 교사가 가진 선택지 하나를 학생이 거의 놓치고 있는 셈입니다. 교사의 확률을 $p_{\mathrm{T}}$, 학생의 확률을 $p_{\mathrm{S}}$라고 쓰면 계산은 다음과 같습니다. 식에서는 같은 문맥이라는 조건을 생략했으며, x는 다음 토큰 후보입니다.

$$\mathrm{KL}(T\Vert S)=\sum _x{p}_{\mathrm{T}}(x)\ln \frac{p_{\mathrm{T}}(x)}{p_{\mathrm{S}}(x)}$$

b의 항은 0.49 × ln(0.49 / 0.02) ≈ 1.57입니다. 이 예시에서는 학생이 b에 너무 작은 확률을 준 차이가 크게 들어갑니다.

Reverse KL은 학생 확률로 가중합니다. 학생이 자주 선택하는 토큰을 교사도 비슷하게 선호하는지 살펴봅니다. c는 학생 확률이 20%인데 교사 확률은 2%입니다. 학생이 자주 내놓을 수 있는 선택을 교사는 거의 지지하지 않는 상황입니다.

$$\mathrm{KL}(S\Vert T)=\sum _x{p}_{\mathrm{S}}(x)\ln \frac{p_{\mathrm{S}}(x)}{p_{\mathrm{T}}(x)}$$

c의 항은 0.20 × ln(0.20 / 0.02) ≈ 0.46입니다. 비율을 뒤집는 것과 함께, 앞에 곱하는 확률도 교사에서 학생으로 바뀝니다. 두 KL 모두 강조한 토큰만이 아니라 a, b, c의 항을 전부 더합니다. 각 항은 음수일 수 있지만, 모든 항을 합친 정확한 KL은 0 이상입니다. 이 예시의 전체 forward KL은 약 1.29, reverse KL은 약 0.76입니다. 이 숫자의 크기만으로 어느 학습법이 더 좋다고 판단할 수는 없습니다. GKD의 KL 정의와 학습 구성

앞 절의 OPD 신호는 reverse KL의 로그비와 부호가 반대입니다. 학생이 뽑은 토큰마다 “log 교사 확률 − log 학생 확률”을 신호로 주고, 그 선택을 늘리거나 줄이도록 학습합니다. 이 신호를 같은 문맥의 학생 분포로 평균내면 reverse KL의 음수가 됩니다. 선택 토큰 하나에서 +0.693이라는 신호가 나오는 것과, 정확한 KL이 음수가 될 수 없다는 사실은 모순이 아닙니다. Thinking Machines의 reverse KL 설명

또한 학생이 문맥을 생성한다는 것과 reverse KL을 사용한다는 것은 별개의 선택입니다. 학생이 만든 문맥에서도 전체 후보의 교사·학생 확률을 구해 forward KL을 계산할 수 있습니다. 이 글은 그중 학생이 선택한 토큰의 로그비를 학습 신호로 쓰는 reverse-KL 방식에 집중합니다.

Reverse KL은 반드시 한 가지 답에 집중하게 할까요?

Mode는 확률 분포의 봉우리를 뜻합니다. 교사가 선호하는 방식이 두 봉우리에 나뉘어 있고 학생이 봉우리 하나로만 표현할 수 있다고 해보겠습니다. 양쪽을 넓게 덮으면 그 사이의 낮은 확률 구간에도 학생 확률을 배정하게 됩니다. Reverse KL에서는 교사가 거의 지지하지 않는 그 구간이 부담이 되어, 한 봉우리 안에 집중하는 쪽이 유리해질 수 있습니다. 이것이 mode-seeking 성향입니다.

반대로 forward KL에서는 교사의 다른 봉우리를 놓치는 것이 부담이므로, 여러 봉우리를 함께 덮는 mode-covering 성향이 나타날 수 있습니다. 다만 학생이 표현할 수 있는 분포를 제한한 예시입니다. 학생이 교사 분포를 그대로 표현할 수 있다면 두 KL 모두 그 분포에서 0이 됩니다. 따라서 reverse KL을 쓴다는 이유만으로 반드시 하나의 토큰이나 답으로 몰린다고 일반화하지는 않겠습니다. GKD Figure A.16, LLM 증류에서 KL을 분석한 연구

과제 어드밴티지와 교사 신호를 더하기

OPD는 critic이나 GRPO와 반드시 하나만 골라야 하는 선택지는 아닙니다. critic과 그룹 비교는 과제 보상으로부터 비교 신호를 추정하는 방식입니다. OPD는 교사 분포라는 학습 목표를 추가하는 경로이므로 함께 사용할 수 있습니다.

Miles v0.1.0에서 학생이 실제로 선택한 토큰을 비교하는 기본 OPD 경로를 보겠습니다. 이를 sampled-token OPD라고 부릅니다. 먼저 선택한 방식으로 과제 어드밴티지 $A_{\text{task}}$를 구하고, 그 뒤에 교사 신호를 더해 $A_{\text{total}}$을 만듭니다.

$$A_{\text{total}}=A_{\text{task}}+\lambda d_t$$

λ는 교사 신호의 비중입니다. 앞 글의 GAE에서 사용한 λ와는 별개인 OPD 계수입니다. Miles 설정에서는 --opd-kl-coef에 해당합니다. 이 글에서는 0.2라는 설명용 값을 사용하겠습니다. 아래 숫자는 합산 직후의 값이며, 예시에서는 뒤에 추가 정규화를 하지 않는다고 가정합니다. 그림의 가로축은 확률이나 학습 시간이 아니라 어드밴티지 값입니다. 두 토큰이 같은 과제 신호에서 출발해 교사 신호만큼 이동하는 모습을 보겠습니다.

그림 2의 b와 d로 돌아와, 과제에서 두 토큰에 모두 +0.5라는 신호를 주었다고 해보겠습니다. 주황색 점이 이 출발점입니다. b의 교사 신호가 +0.693이면 0.5 + 0.2 × 0.693 ≈ 0.639입니다. 토큰 d의 교사 신호가 −0.693이면 약 0.361입니다. 보라색 화살표가 교사 항의 덧셈이고, 초록색 점이 합산 결과입니다. 교사 신호가 과제의 양수 신호를 강화하거나 약화합니다.

토큰 d에서도 합산 결과는 아직 양수입니다. 교사 신호가 음수라는 사실만으로 최종 정책 손실도 그 토큰을 억제한다고 판단할 수는 없습니다. 이 구성에서는 결합된 $A_{\text{total}}$의 부호와 크기를 먼저 봐야 합니다. 계수나 두 신호의 크기에 따라서는 최종 부호도 바뀔 수 있습니다.

토큰 b를 따라가면, 이제 고정한 계수 약 0.639가 기본 손실의 log pθ(b | P, u)에 곱해집니다. 역전파는 현재 학생의 확률 계산을 통해 LM head와 이전의 학습 가능한 층으로 흐르고, optimizer가 학생 가중치를 바꿉니다. $A_{\text{total}}$이나 교사, 기록된 토큰 ID를 미분하는 경로는 아닙니다.

GRPO와 OPD를 함께 쓰면 GRPO의 응답별 공통 어드밴티지에 토큰마다 다른 교사 항이 더해집니다. Miles는 이 합산 뒤에 설정에 따라 어드밴티지를 정규화할 수 있습니다. 이후의 확률비, clipping, 손실 집계는 해당 학습 구성의 정책 손실 경로를 따릅니다. 합산 위치와 신호의 부호는 Miles v0.1.0 OPD 구현에서 확인할 수 있습니다.

과제 점수 없이도 정책 손실이 생기는 이유

과제 어드밴티지를 0으로 두면 어떻게 될까요? 그림 4의 주황색 출발점을 +0.5에서 0으로 옮겨 생각하면 됩니다. 교사 항은 같으므로 위의 계산에서 $A_{\text{task}}$만 사라지고, λ × dₜ는 남습니다. 두 토큰의 최종 신호는 약 +0.139와 −0.139가 됩니다.

학생의 정책 손실은 이 신호와 학생의 현재 확률 계산을 이용합니다. 기본 설명식이라면 −A_total × log p_current이고, 확률비와 clipping을 쓰는 구성이라면 그 구성의 목적함수에 $A_{\text{total}}$이 들어갑니다. 따라서 과제 신호가 0이어도 교사와 학생의 확률 차이로 학습할 수 있습니다.

교사가 직접 최종 손실을 출력하는 것은 아닙니다. 교사는 확률 정보를 제공하고, 학습 코드는 그 정보로 신호를 만든 뒤 정해진 정책 손실을 계산합니다. 2편에서 보상이 손실 그 자체는 아니었던 것과 같은 구분입니다.

Miles의 pure OPD 예시는 과제 보상을 0으로 돌려주는 경로를 제공합니다. GRPO처럼 각 과제 보상에서 그룹 평균을 빼는 구성에서는 과제 어드밴티지가 0이 되고 교사 항이 남습니다. 다만 보상을 0으로 설정하면 어떤 어드밴티지 추정 방식에서든 즉시 $A_{\text{task}}$가 0이 된다고 일반화해서는 안 됩니다. critic의 예측이나 다른 보상·규제 항이 남아 있는 구성에서는 실제 계산 경로를 확인해야 합니다. 이 OPD 합산은 어드밴티지를 바꾸고 return 목표는 그대로 두므로, 교사 신호가 critic의 학습 목표에도 자동으로 들어가는 것은 아닙니다. Miles v0.1.0 rollout 및 보상 처리

언제 계산하고 무엇을 고정하는가

한 번의 학습 경로를 시간 순서로 정리하면 다음과 같습니다.

1. 학생이 질문에서 응답을 생성하고 문맥·토큰 이력을 남깁니다.
2. 같은 이력에서 교사의 토큰별 확률을 구합니다.
3. 학습 신호용으로 고정한 학생 확률과 비교해 dₜ를 만들고 과제 어드밴티지에 더합니다.
4. 학습 중인 학생의 현재 확률로 정책 손실을 계산하고, 학생 가중치에 역전파합니다.

신호를 계산할 때의 학생 확률과, 정책 손실에서 미분하는 현재 학생 확률을 구별해야 합니다. Miles는 설정에 따라 생성 때 기록한 로그확률 또는 업데이트 전 학생 모델을 평가한 로그확률을 신호 계산에 사용하고, 그 신호를 고정합니다. 정책이 신호 자체를 바꿔 손실을 줄이게 하는 것이 아니라, 주어진 신호에 따라 자신의 선택 확률을 바꾸게 합니다. Miles v0.1.0 advantage 계산

교사의 계산 위치에도 선택지가 있습니다. Miles의 SGLang 교사 모드는 외부 교사 서버에서 구한 정보를 rollout 데이터에 담고, Megatron 교사 모드는 학습 쪽에서 교사의 순전파로 구합니다. 두 경우 모두 이 정책 손실로 교사 가중치를 갱신하지 않습니다. 또한 이 토큰별 비교는 학생과 교사가 같은 토큰 ID를 같은 의미로 해석할 수 있다는 조건을 사용합니다. Miles v0.1.0 OPD 문서

교사 모델은 어떤 역할인가

이제 이름이 비슷하게 등장했던 모델들을 역할로 구별할 수 있습니다.

역할	이번 연재에서 제공하는 정보
보상 모델·검증기	과제의 기준에 따른 결과 점수
Critic	현재 문맥에서 미래에 받을 return의 예측
고정 reference	학습 중 유지하려는 비교 기준의 확률
OPD 교사	학생이 방문한 문맥에서 학습 목표가 되는 토큰 확률

고정 reference와 교사는 둘 다 확률을 제공할 수 있지만, 학습에서 맡는 목적이 다릅니다. reference는 정한 기준에서의 이탈을 관리하고, 교사는 학생에게 익히게 할 행동 분포를 제공합니다. 같은 체크포인트를 어느 역할에 사용하는지는 학습 설계의 선택입니다.

OPD 전체가 반드시 이 글의 어드밴티지 합산식으로만 구현되는 것은 아닙니다. 분포 차이에 대한 증류 손실을 직접 구성하거나, 선택한 토큰 외에 여러 후보의 확률을 사용하는 방식도 있습니다. Miles에도 top-k 후보를 이용하는 별도의 신호 계산이 있습니다. 이번 글은 학생 문맥, 고정된 교사 신호, 학생의 정책 손실이라는 연결을 보기 위해 기본 sampled-token 경로로 한정했습니다.

교사는 한 명일 필요도 없습니다. 수학 문제는 수학 교사에게, 코드 문제는 코드 교사에게 보내는 식으로 확장할 수 있습니다. Miles의 다중 교사 라우팅에서는 샘플 하나에 교사 하나를 선택하며, 모든 교사의 확률을 반드시 평균내는 것은 아닙니다. GLM-5는 이전 SFT·RL 단계의 체크포인트들을 교사로 삼아, 후속 학습 중 약해진 능력을 회복하는 cross-stage OPD를 설명합니다. GLM-5 §3.5

다음 글에서는 RL과 OPD를 전체 학습 전략에 배치하겠습니다. DeepSeek V4의 전문 모델 학습과 다중 교사 통합에서 시작해, OPD를 RL 앞에 두는 경우와 MiMo의 여러 prefix를 이용하는 확장을 비교하겠습니다.