Trajectory-Refined Distillation Review

0. Introduction

Paper link

Code link

On-policy distillation, 줄여서 OPD는 student가 직접 생성한 rollout 위에서 teacher의 token distribution을 받아 학습한다. Off-policy reference만 따라가는 SFT보다 student가 실제로 방문하는 state를 다룰 수 있고, trajectory 끝의 scalar reward만 받는 RLVR보다 훨씬 촘촘한 supervision을 제공한다. 최근 reasoning model의 post-training recipe에서 OPD와 on-policy self-distillation, 즉 OPSD가 자주 등장하는 이유다.

그런데 이 설명에는 중요한 전제가 숨어 있다. Teacher가 student rollout의 각 prefix를 볼 때마다, 그 prefix를 더 나은 방향으로 이끄는 유용한 distribution을 제공할 수 있어야 한다. Student가 초반에 잘못된 reasoning branch로 들어간 뒤에도 이 전제가 유지될까.

Trajectory-Refined Distillation은 이 질문에 부정적으로 답한다. 잘못된 prefix가 이미 논리적 모순이나 계산 오류를 포함해 단순 continuation만으로 정답에 도달할 수 없으면, teacher는 두 행동 사이에서 갈라진다. 현재 문맥을 자연스럽게 이어 가거나, 갑자기 “Wait” 또는 “Actually” 같은 correction onset을 내놓고 되돌아가야 한다. 전자는 오류를 유지하고, 후자는 student 분포에서 매우 낮은 확률의 mode일 수 있다.

더 큰 문제는 token-level KL이 correction onset을 한 번 추천한다고 해서 그 이후의 correction path까지 가르쳐 주지 못한다는 점이다. 다음 token의 teacher query는 방금 제안한 correction 위가 아니라, 원래 student rollout의 다음 잘못된 prefix 위에서 다시 계산된다. 따라서 teacher는 매 position에서 비슷한 correction onset만 반복해서 추천하고, 올바른 multi-step recovery는 하나의 일관된 trajectory로 펼쳐지지 않는다.

논문은 이 현상을 prefix failure라고 부른다. 그리고 loss clipping, top-k truncation, token importance reweighting 같은 기존 처방은 잘못된 prefix를 그대로 둔 채 각 token loss의 크기만 조절하므로 문제의 scale이 맞지 않는다고 주장한다.

TRD의 해법은 단순하다. KL loss를 계산하기 전에 teacher가 student의 raw rollout을 먼저 다시 쓴다. Student가 만든 trajectory를 $y_o$, teacher가 수정한 trajectory를 $y_r$라고 하면 pipeline은 다음과 같다.

\[\text{student rollout } y_o \to \text{teacher refinement } y_r \to \text{KL distillation on } y_r\]

이렇게 하면 distillation context가 원래의 실패 trajectory가 아니라, teacher가 실제로 펼친 correction trajectory를 따라간다. 핵심 변화는 새로운 divergence나 복잡한 reward가 아니다. Supervision을 계산하는 trajectory 자체를 먼저 고치는 것이다.

한 줄 요약: TRD는 잘못된 student prefix 위에서 token-level teacher signal이 분절되는 prefix failure를 지적하고, raw rollout을 teacher가 trajectory 단위로 수정한 뒤 그 refined trajectory 위에서 KL distillation을 수행한다.

이 논문을 지금 볼 가치가 있는 이유는 다음과 같음.

• OPD의 failure를 loss instability가 아니라 supervision context의 구조적 mismatch로 설명한다.
• Forward KL, reverse KL, clipping, top-k가 서로 다른 현상을 보이는 이유를 하나의 prefix failure 관점으로 연결한다.
• OPD와 parameter-sharing OPSD에 같은 trajectory-level intervention을 적용한다.
• Math와 code에서 single-attempt accuracy뿐 아니라 Pass@K coverage, trajectory correctness, trajectory length, wall-clock cost를 함께 본다.
• Agent, search, tool-use처럼 잘못된 early action이 긴 suffix 전체를 오염시키는 문제에도 재사용할 수 있는 설계 원리를 제공한다.

1. Problem Setting

1-1. On-policy distillation의 기본 구조

Standard knowledge distillation은 fixed data prefix에서 teacher와 student distribution을 맞춘다. OPD는 prefix를 student가 직접 생성한다는 점이 다르다. Prompt $x$가 주어지면 student policy가 rollout을 샘플링한다.

\[y_o \sim \pi_\theta(\cdot \mid x)\]

그다음 teacher는 student가 방문한 각 prefix $y_{o,\lt t}$에서 token distribution을 계산한다. Representative objective는 다음처럼 쓸 수 있다.

\[\mathcal{L}_{\mathrm{OPD}}(\theta) = \mathbb{E}_{x \sim \mathcal{D},\, y_o \sim \pi_\theta(\cdot \mid x)} \left[ \frac{1}{|y_o|} \sum_{t=1}^{|y_o|} D\left( \pi_\theta(\cdot \mid x,y_{o,\lt t}) \,\|\, \pi_T(\cdot \mid x,y_{o,\lt t}) \right) \right]\]

여기서 $D$는 forward KL이나 reverse KL이 될 수 있다. 실제 recipe에서는 full vocabulary distribution을 사용해 sample-level gradient variance를 줄이는 경우가 많다.

OPSD는 별도 teacher model을 두지 않고 같은 parameter를 공유한다. Student branch는 문제만 보고, teacher branch는 reference solution 같은 privileged information을 추가로 본다. Teacher branch에는 stop-gradient를 적용한다.

\[\pi_S = \pi_\theta(\cdot \mid x,y_{\lt t})\] \[\pi_T = \operatorname{sg}\left[ \pi_\theta(\cdot \mid x,y^*,y_{\lt t}) \right]\]

이 구조는 external teacher 없이도 같은 model이 더 풍부한 context에서 만든 distribution을 plain-context policy에 전달할 수 있다는 장점이 있다.

OPD와 OPSD가 매력적인 이유는 다음 세 가지다.

Property	SFT	RLVR	OPD / OPSD
Prefix source	Fixed expert data	Current policy	Current policy
Main signal	Hard target token	End-of-trajectory reward	Dense token distribution
On-policy state coverage	낮음	높음	높음
Feedback density	높음	낮음	높음
Main risk	Distribution mismatch	Sparse reward	Teacher signal quality on failed prefixes

하지만 dense signal이 항상 useful signal이라는 보장은 없다. Prefix failure는 바로 이 지점에서 발생한다.

1-2. Prefix failure

논문은 student prefix $y_{o,\lt t}$가 다음 조건을 만족할 때 prefix failure가 발생한다고 본다.

• Prefix 안에 reference와 모순되는 reasoning error가 이미 들어 있다.
• 단순 continuation만으로는 정답에 도달하기 어렵다.
• Correct solution을 만들려면 retract, backtrack, reflection, branch replacement가 필요하다.

수학 문제에서 초반에 식의 부호를 잘못 잡았거나 잘못된 lemma를 채택한 경우가 대표적이다. Code에서는 data structure choice나 algorithmic complexity assumption이 잘못되어 suffix만 고쳐서는 통과할 수 없는 경우가 해당된다.

이 상황에서 teacher가 할 수 있는 행동은 대략 두 mode로 나뉜다.

1. Wrong-continuation mode
   • 현재 prefix와 문법적, 논리적 continuity를 유지한다.
   • Student가 이미 높은 probability를 둔 영역이다.
   • 하지만 최종 정답으로 이어지지 않는다.
2. Correction mode
   • “Wait”, “Actually”, “We need to reconsider” 같은 epistemic token으로 방향을 튼다.
   • 이후 오류를 retract하고 새로운 derivation을 시작한다.
   • Student distribution에서는 낮은 probability의 OOD-like mode일 수 있다.

Teacher distribution을 두 mode의 mixture로 단순화하면 다음처럼 생각할 수 있다.

\[\pi_T = \alpha \pi_{\mathrm{continue}} + (1-\alpha)\pi_{\mathrm{correct}}\]

Forward KL은 teacher probability로 weighting되므로 correction mode를 cover하려는 압력이 커진다. Correction onset이 student에게 매우 낮은 probability라면 해당 token의 KL이 지나치게 커져 training을 불안정하게 만들 수 있다. 그래서 token-level KL clipping이 등장한다.

Reverse KL은 student probability로 weighting되므로 student가 이미 선택한 wrong-continuation mode에 집중하기 쉽다. Correction token은 student probability가 작아 gradient에서 약하게 반영된다. 그래서 teacher top-k나 token importance reweighting으로 teacher-preferred mode를 다시 강조하려는 시도가 등장한다.

이 관점에서 서로 반대처럼 보이는 처방들이 이해된다.

KL regime	Prefix failure에서 나타나는 문제	Typical token-level fix
Forward KL	Low-probability correction token이 loss를 지배	High-KL token clipping
Reverse KL	Wrong continuation이 loss를 지배하고 correction은 약함	Teacher top-k, importance reweighting

그러나 두 처방 모두 같은 한계를 공유한다. Teacher를 query하는 prefix는 여전히 원래의 실패 trajectory다.

1-3. 왜 perfect teacher도 correction path를 충분히 전달하지 못하는가

가장 중요한 부분은 teacher quality와 별개로 생기는 context mismatch다.

가령 position $t$에서 원래 prefix가 $y_{o,\lt t}$이고, teacher가 correction onset token $c_t$를 추천했다고 하자. Ideal supervision은 다음 context에서 이어져야 한다.

\[(y_{o,\lt t}, c_t)\]

그리고 다음 correction token $c_{t+1}$은 아래 context에서 학습되어야 한다.

\[(y_{o,\lt t}, c_t, c_{t+1})\]

즉 correction path 자체가 autoregressive하게 펼쳐져야 한다.

하지만 standard token-level OPD는 frozen rollout $y_o$ 위에서 teacher distribution을 계산한다. Position $t+1$에서 teacher가 보는 context는 correction을 반영한 prefix가 아니라 아래와 같다.

\[(y_{o,\lt t}, y_{o,t})\]

Position $t+2$에서는 더 깊어진 원래 failure prefix를 본다.

\[(y_{o,\lt t}, y_{o,t}, y_{o,t+1})\]

그래서 ideal supervision pair와 실제 supervision pair는 첫 correction onset 뒤 즉시 갈라진다.

\[\text{ideal:} \quad (y_{o,\lt t},c_t) \to (y_{o,\lt t},c_t,c_{t+1}) \to \cdots\] \[\text{fragmented:} \quad (y_{o,\lt t},c_t) \to (y_{o,\lt t},y_{o,t},c_t) \to (y_{o,\lt t},y_{o,t},y_{o,t+1},c_t) \to \cdots\]

Token clipping은 이 pair들의 weight를 낮출 수 있다. Top-k는 target support를 바꿀 수 있다. Importance weighting은 어떤 position을 강조할지 바꿀 수 있다. 하지만 correction path 위의 새로운 pair를 생성하지는 않는다.

논문의 핵심 diagnosis는 여기다.

Prefix failure는 bad token weight 문제가 아니라 bad supervision context 문제다.

2. Core Idea

2-1. Main contribution

TRD는 OPD pipeline 앞에 trajectory refinement stage를 추가한다.

1. Student가 raw on-policy rollout $y_o$를 생성한다.
2. Teacher가 $y_o$를 읽고 더 나은 trajectory $y_r$로 다시 쓴다.
3. Standard full-vocabulary KL distillation을 $y_r$의 prefix 위에서 수행한다.

이를 식으로 단순화하면 다음과 같다.

\[y_o \sim \pi_\theta(\cdot \mid x)\] \[y_r \sim q_T(\cdot \mid x,y_o,\mathcal{I})\] \[\mathcal{L}_{\mathrm{TRD}}(\theta) = \mathbb{E} \left[ \frac{1}{|y_r|} \sum_{t=1}^{|y_r|} D_{\mathrm{KL}} \left( \pi_T(\cdot \mid \mathcal{C}_T,y_{r,\lt t}) \,\|\, \pi_\theta(\cdot \mid x,y_{r,\lt t}) \right) \right]\]

여기서 $\mathcal{I}$는 refinement에 제공되는 auxiliary information이고, $\mathcal{C}_T$는 distillation teacher의 context다. OPD와 OPSD에서 이 값이 다르다.

Setting	Refinement teacher	Refinement input	Distillation teacher
OPD	별도 stronger teacher	Problem plus raw rollout $y_o$	별도 teacher
OPSD	Student와 같은 backbone	Problem plus reference $y^*$ plus raw rollout $y_o$	Same model with privileged context

공개 code의 rewrite prompt를 보면 OPD와 OPSD 차이가 명확하다. OPD teacher는 problem과 initial response를 보고 overall structure를 가능한 한 유지하면서 computation 또는 logic error를 고친다. OPSD teacher는 여기에 expert solution을 함께 보고 reference와 일관된 solution으로 다시 쓴다.

이 구분은 중요하다. OPD의 gain을 hidden reference leakage로만 설명할 수 없고, OPSD의 gain은 privileged reference를 trajectory-level output으로 변환하는 과정과 연결된다.

2-2. Design intuition

TRD의 설계 직관은 세 가지다.

1) Correction을 token이 아니라 path로 materialize한다

Standard OPD는 correction onset의 probability를 높일 수 있지만, correction 뒤에 이어질 reasoning state를 직접 방문하지 않는다. TRD는 teacher rewrite를 통해 correction path 전체를 실제 token sequence $y_r$로 만든다.

이후 KL은 다음 pair에서 계산된다.

\[(x,y_{r,\lt t}) \to \pi_T(\cdot \mid x,y_{r,\lt t})\]

즉 teacher와 student 모두 correction이 누적된 prefix를 본다. Fragmented gradient가 하나의 coherent trajectory supervision으로 바뀐다.

2) Student rollout을 anchor로 유지한다

Teacher가 reference solution을 처음부터 새로 쓰는 대신 student의 $y_o$를 rewrite하도록 한 이유는 on-policy character를 보존하기 위해서다.

• Student가 사용하는 terminology와 reasoning style을 일부 유지할 수 있다.
• Correct intermediate step은 재사용할 수 있다.
• Student가 전혀 방문하지 않은 expert-only format으로 급격히 이동하는 것을 줄일 수 있다.
• Error가 발생한 부분을 중심으로 local 또는 structural correction을 수행할 수 있다.

논문은 이를 on-policy support 안의 refinement로 설명한다. 다만 실제 implementation에서 hard likelihood threshold, edit-distance constraint, rejection sampling gate가 명시적으로 support를 enforce하는 것은 아니다. Student rollout을 prompt anchor로 사용하고 useful structure를 유지하라는 instruction이 practical mechanism이다.

3) Correct rollout도 다시 써 exploration을 넓힌다

TRD는 incorrect rollout만 수정하는 error repair method가 아니다. Raw rollout이 verifier를 통과한 경우에도 teacher는 더 짧거나 다른 valid derivation을 제시할 수 있다.

이 특성은 distillation dataset의 support를 넓힌다.

• 같은 problem에 대한 alternative proof path가 생긴다.
• Redundant reasoning을 제거한 shorter trajectory가 생긴다.
• Student가 우연히 맞힌 brittle solution을 더 stable한 path로 바꿀 수 있다.
• Pass@K에서 서로 다른 solution mode를 찾을 가능성이 커진다.

따라서 TRD의 목표는 단순한 trajectory denoising이 아니라 corrective exploration plus coherent distillation이다.

3. Architecture / Method

3-1. Overview

Item	Description
Goal	Prefix failure가 만든 fragmented token supervision을 coherent trajectory supervision으로 변환
Raw trajectory	Current student가 생성한 $y_o$
Refined trajectory	Teacher가 $y_o$를 rewrite한 $y_r$
OPD teacher	Separate Qwen3-8B teacher
OPSD teacher	Student와 같은 model, privileged reference context 사용
Main distillation loss	Full-vocabulary forward KL on $y_r$
Main tasks	Competition math and code generation
Main evaluation	Avg@16, Pass@16, 일부 Pass@128, verifier pass rate, trajectory length
Main distinction	Loss weight를 바꾸지 않고 supervision prefix 자체를 교체

3-2. Stage 1: Raw on-policy rollout

첫 stage는 vanilla OPD와 같다. Current student policy에서 rollout을 샘플링한다.

\[y_o \sim \pi_\theta(\cdot \mid x)\]

이 stage가 중요한 이유는 refinement source가 fixed expert corpus가 아니라 current policy의 behavior이기 때문이다. Training이 진행되어 student가 바뀌면 $y_o$의 error pattern과 reasoning style도 바뀐다. 따라서 refinement target도 current student competence boundary를 따라 이동한다.

Raw rollout은 verifier로 correct 또는 incorrect를 판정할 수 있지만, TRD는 correctness flag만으로 token loss를 masking하는 데 그치지 않는다. Full response를 rewrite input으로 사용한다.

3-3. Stage 2: Teacher-guided trajectory refinement

Teacher는 raw rollout을 보고 refined rollout을 생성한다.

OPD refinement

OPD에서는 별도 teacher가 problem과 initial solution을 본다. Public prompt의 핵심 instruction은 다음과 같다.

• Overall structure와 reasoning path를 가능한 한 보존한다.
• Computation 또는 logic error를 찾고 수정한다.
• Correct intermediate step과 meaningful work를 유지한다.
• Rewritten solution만 출력한다.

Code task에서도 같은 원리를 쓴다. Teacher는 problem과 initial Python solution을 보고 correctness issue와 edge case를 수정하되 useful approach를 보존한다.

OPSD refinement

OPSD에서는 같은 backbone이 teacher role을 맡지만, reference solution $y^*$를 privileged input으로 추가로 본다.

• Reference가 target reasoning과 method를 제공한다.
• Raw rollout이 student의 current style과 visited path를 제공한다.
• Teacher는 둘을 결합해 reference-consistent rewrite를 만든다.

이 구조는 SFT와 다르다. SFT는 $y^$ 자체를 hard target으로 사용한다. TRD-OPSD는 $y^$를 teacher의 private guide로 쓰고, 실제 student target은 raw rollout을 기반으로 다시 생성된 $y_r$다.

3-4. Stage 3: KL distillation on the refined trajectory

Refined trajectory가 만들어지면 standard token-level distillation을 수행한다. 중요한 차이는 teacher query와 student query가 $y_{r,\lt t}$를 context로 사용한다는 점이다.

\[\mathcal{L}_{\mathrm{TRD}} = \frac{1}{|y_r|} \sum_{t=1}^{|y_r|} D_{\mathrm{KL}} \left( \pi_T(\cdot \mid \mathcal{C}_T,y_{r,\lt t}) \,\|\, \pi_\theta(\cdot \mid x,y_{r,\lt t}) \right)\]

Public recipe는 TRD variant에 full-vocabulary forward KL을 사용한다. 비교군은 다음을 포함한다.

Variant	Trajectory	Loss intervention
Forward KL	$y_o$	None
Forward KL plus clipping	$y_o$	High token KL cap
Reverse KL	$y_o$	None
Reverse KL plus top-k	$y_o$	Teacher support truncation
TRD	$y_r$	Forward KL on refined context

이 비교는 trajectory intervention과 token-level intervention을 분리해 본다는 점에서 의미가 있다.

3-5. Prefix failure를 어떻게 바꾸는가

Standard OPD에서는 teacher가 아래와 같은 frozen context sequence를 본다.

\[y_{o,\lt 1}, y_{o,\lt 2}, \ldots, y_{o,\lt T}\]

TRD에서는 refinement stage에서 teacher가 새 trajectory를 실제로 생성한다. Distillation stage의 context는 다음과 같다.

\[y_{r,\lt 1}, y_{r,\lt 2}, \ldots, y_{r,\lt T_r}\]

따라서 correction onset 뒤의 token도 correction을 반영한 prefix에 conditioned된다. Teacher가 “Wait”를 추천한 다음 position에서 원래 error를 다시 보지 않고, “Wait”가 포함된 state에서 다음 correction token을 지도한다.

이 차이는 작아 보이지만 gradient가 평가되는 state distribution을 바꾼다.

• Token-level fix: 같은 state에서 loss weight만 변경
• TRD: Teacher와 student가 방문하는 supervision state를 변경

Reinforcement learning 관점으로 보면 reward shaping보다 trajectory relabeling에 가깝다. Offline RL 관점에서는 failed behavior를 teacher-guided improved behavior로 바꾸는 relabeling과 닮았고, imitation learning 관점에서는 DAgger-style on-policy state coverage와 expert correction을 결합한 형태로 볼 수 있다.

4. Training / Data / Recipe

4-1. Models and data

논문은 OPD와 OPSD를 모두 평가한다.

Setting	Student	Teacher	Training data
OPD math	Qwen3-1.7B, Qwen3-4B-Instruct-2507	Qwen3-8B	DeepScaleR, 약 40K problems
OPD code	Qwen3-1.7B, Qwen3-4B-Instruct-2507	Qwen3-8B	TACO, 약 25K problems
OPSD math	Qwen3-4B-Instruct-2507, Qwen3-8B	Same backbone with privileged reference	DeepScaleR

Math evaluation은 AIME24, AIME25, HMMT25, BeyondAIME, AMOBench를 사용한다. Code evaluation은 HumanEval+, MBPP+, LiveCodeBench v6를 사용한다.

Public repository에는 다음 component가 공개되어 있다.

• OPD and OPSD trainer
• Forward KL, reverse KL, clipping, top-k variant
• Raw rollout $y_o$ preparation
• Refined rollout $y_r$ prompt preparation
• Math benchmark runner
• EvalPlus and LiveCodeBench helper
• Qwen3 model별 run wrapper
• verl 기반 distributed training runtime

4-2. Training pipeline

Practical pipeline은 세 번의 주요 model operation으로 나뉜다.

Step 1. Generate $y_o$

Current student checkpoint로 training prompt에 대한 rollout을 생성한다. Reward 또는 verifier outcome과 initial response를 parquet 형태로 저장한다.

Step 2. Generate $y_r$

y_r_prepare.py가 raw rollout을 rewrite prompt로 바꾼다. OPD와 OPSD에 따라 prompt가 다르다.

• OPD: problem plus initial response
• OPSD: problem plus expert solution plus initial response

Teacher rollout을 실행해 refined response를 저장한다.

Step 3. KL training

Public wrapper의 TRD path는 다음 설정을 사용한다.

• KL_TYPE=forward
• KL_METHOD=full_vocab
• Y_MODE=y_r
• TEACHER_TRAINING_PROMPT=refine

예를 들어 Qwen3-4B-Instruct math OPD wrapper는 Qwen3-8B teacher, DeepScaleR training source, AIME24/25, HMMT25, BeyondAIME, AMOBench evaluation을 지정한다.

4-3. Why full-vocabulary forward KL

Trajectory를 고쳤다고 해서 hard-label SFT로 끝내지 않고 teacher distribution을 다시 사용하는 이유는 refined token 뒤의 alternative probability structure를 전달하기 위해서다.

Hard target은 $y_{r,t}$ 하나만 본다.

\[\mathcal{L}_{\mathrm{SFT}} = - \sum_t \log \pi_\theta(y_{r,t} \mid x,y_{r,\lt t})\]

Full-vocabulary KL은 같은 prefix에서 teacher가 허용하는 alternative token과 uncertainty를 함께 전달한다.

\[\mathcal{L}_{\mathrm{FKL}} = \sum_t \sum_{v \in \mathcal{V}} \pi_T(v \mid \mathcal{C}_T,y_{r,\lt t}) \log \frac{ \pi_T(v \mid \mathcal{C}_T,y_{r,\lt t}) }{ \pi_\theta(v \mid x,y_{r,\lt t}) }\]

TRD의 message는 token distribution matching이 쓸모없다는 것이 아니다. Distribution matching을 올바른 trajectory context에서 수행해야 한다는 것이다.

4-4. Trajectory length and compute

Refinement stage는 teacher rollout을 한 번 더 요구한다. 따라서 generation cost만 보면 vanilla OPD보다 비싸다. 하지만 refined trajectory가 raw trajectory보다 훨씬 짧아지면 이후 KL training의 sequence cost가 줄 수 있다.

Qwen3-8B OPSD training corpus 분석에서는 다음 변화가 보고된다.

Metric	Raw $y_o$	Refined $y_r$	Reference $y^*$
Verifier pass rate	65.8%	81.4%	-
Median length	약 7.7K tokens	약 0.88K tokens	약 0.49K tokens

Median 기준으로 $y_r$는 $y_o$보다 약 9배 짧다. 이는 raw reasoning의 반복, dead end, unnecessary verification을 teacher rewrite가 제거했기 때문이다.

다만 이 숫자를 inference acceleration로 읽으면 안 된다. Shorter $y_r$는 training target compression이다. Distilled model이 serving에서 항상 9배 짧게 답한다는 결과는 아니다.

4-5. Wall-clock cost is regime-dependent

8x H100 80GB single-node 환경에서 보고된 end-to-end timing은 model과 setting에 따라 다르다.

Setting	Vanilla	TRD	Interpretation
OPSD Qwen3-8B	약 9h40	약 9h20	Shorter $y_r$가 extra refinement cost를 상쇄
OPSD Qwen3-4B	약 4h20	약 5h30	Refinement overhead가 더 큼
OPD Qwen3-4B	약 5h20	약 9h00	Separate teacher generation cost가 크게 반영

따라서 TRD를 “추가 비용 없이 더 좋은 distillation”이라고 일반화하면 안 된다. Correctness gain, teacher serving cost, trajectory compression, KL sequence length을 함께 계산해야 한다.

5. Evaluation

5-1. Metrics

논문은 각 prompt에 대해 16개 sample을 생성하고 두 metric을 본다.

• Avg@16: 16개 sample의 평균 correctness. Single-attempt accuracy의 Monte Carlo estimate에 가깝다.
• Pass@16: 16개 중 하나라도 정답이면 성공. Reasoning coverage와 exploration breadth를 본다.

Pass@16만 높고 Avg@16이 낮다면 rare lucky sample이 존재할 수 있다. Avg@16과 Pass@16이 함께 오르면 typical sample quality와 support coverage가 동시에 좋아졌다고 볼 수 있다.

5-2. Main math results

대표 결과를 base model 대비 TRD로 요약하면 다음과 같다.

Setting	Benchmark	Metric	Base	TRD	Gain
OPD Qwen3-4B-Instruct	AMOBench	Pass@16	23.1	35.9	+12.8 points
OPSD Qwen3-8B	HMMT25	Pass@16	66.7	76.3	+9.6 points
OPSD Qwen3-8B	AMOBench	Pass@16	41.0	61.5	+20.5 points
OPSD Qwen3-8B	AMOBench	Avg@16	15.9	17.3	+1.4 points

AMOBench에서 Qwen3-8B OPSD Pass@16은 41.0에서 61.5로 오른다. Absolute gain은 20.5 points이고 relative gain은 50%다.

이 결과에서 중요한 것은 Pass@16 gain이 Avg@16 gain보다 훨씬 크다는 점이다. TRD가 typical sample의 correctness를 조금 높이는 동시에, 여러 sample을 뽑았을 때 적어도 하나의 valid reasoning path를 찾는 probability를 크게 넓혔다고 해석할 수 있다.

논문은 AMOBench에서 OPD Qwen3-4B의 Pass@128도 46.7에서 53.8로 증가한다고 보고한다. Sample budget을 더 크게 늘려도 coverage gain이 남는다는 증거다.

5-3. What the baselines tell us

TRD의 비교군은 단순 vanilla KL뿐 아니라 token-level correction method를 포함한다.

• Forward KL
• Forward KL plus clipping
• Reverse KL
• Reverse KL plus teacher top-k
• Base model
• TRD on $y_r$

Reported OPSD Avg@16 block에서 TRD는 base model 아래로 떨어지지 않는 반면, 일부 token-level baseline은 task에 따라 base보다 낮아진다. 이는 prefix failure가 있는 상황에서 loss engineering만으로 dense supervision을 안정화하기 어렵다는 논문의 diagnosis와 일치한다.

그러나 “trajectory refinement면 항상 이긴다”고 단정하면 안 된다. Gain은 difficult benchmark에서 특히 크고, easier benchmark에서는 margin이 작을 수 있다. 또한 Pass@K는 sampling temperature와 decoding budget의 영향을 받는다.

5-4. Trajectory quality analysis

Raw rollout과 refined rollout을 직접 비교한 분석이 이 논문의 주장을 가장 잘 뒷받침한다.

1. Verifier pass rate improves
   • Raw $y_o$: 65.8%
   • Refined $y_r$: 81.4%
2. Trajectory becomes much shorter
   • Raw median: 약 7.7K tokens
   • Refined median: 약 0.88K tokens
   • Reference median: 약 0.49K tokens
3. Epistemic correction-token concentration drops
   • Failed-prefix $y_o$에서는 teacher가 small residual probability budget의 큰 부분을 correction onset token에 반복해서 할당한다.
   • Refined $y_r$에서는 이 concentration이 크게 줄어든다.
4. Teacher-student gap becomes useful again
   • Raw failed trajectories에서는 privileged teacher와 student가 같은 wrong continuation에 붙어 perplexity gap과 KL signal이 거의 사라질 수 있다.
   • Refined trajectory에서는 teacher context가 correction path를 따라가므로 informative difference가 복원된다.

즉 performance gain과 mechanism analysis가 연결된다. TRD는 단지 더 좋은 target response를 만드는 것이 아니라, teacher distribution이 informative하게 평가될 state sequence를 다시 만든다.

5-5. Code generation result and the LiveCodeBench warning

Code generation에서는 결과가 더 혼합되어 있다. HumanEval+와 MBPP+에서는 TRD가 경쟁력 있는 결과를 보이지만, LiveCodeBench에서는 모든 distillation variant가 base model을 완전히 회복하지 못한다.

Qwen3-4B-Instruct Pass@16 기준 대표 수치는 다음과 같다.

Model	LiveCodeBench Pass@16
Base	55.2
TRD	54.1

이 결과는 중요한 caution이다.

• Strong teacher가 raw code의 subtle bug를 항상 고치는 것은 아니다.
• Reference-free OPD rewrite는 hidden test distribution을 알 수 없다.
• Teacher가 code를 더 짧게 정리하면서 edge case를 잃을 수 있다.
• Math verifier보다 code execution feedback가 더 복잡하고 task coverage가 넓다.
• Trajectory-level correction이 benchmark distribution shift를 자동으로 해결하지 않는다.

따라서 code domain에서 TRD를 적용하려면 rewrite 이후 unit test, static check, compiler feedback, hidden-like adversarial test를 gating에 포함하는 편이 안전하다.

5-6. Which result matters most

가장 눈에 띄는 숫자는 AMOBench Pass@16 +20.5 points지만, method를 판단할 때는 세 결과를 함께 봐야 한다.

1. Base to TRD performance gain
   • Final model이 실제로 좋아졌는가.
2. Raw to refined verifier gain
   • Upstream trajectory construction이 의도대로 동작하는가.
3. Raw to refined length reduction
   • Extra teacher rollout이 downstream KL cost와 data quality를 어떻게 바꾸는가.

이 세 축을 함께 보면 TRD는 단순히 더 많은 teacher compute를 쓴 method가 아니다. Teacher compute를 token-level marginal query가 아니라 coherent rewrite에 쓰고, 그 결과를 shorter and cleaner distillation path로 바꾼 method다.

6. Limitations

1. Extra teacher generation cost
   • Raw rollout 뒤에 refined rollout을 한 번 더 생성해야 한다.
   • OPD에서는 separate teacher serving cost가 커질 수 있다.
   • Wall-clock timing은 setting에 따라 vanilla보다 빠르기도, 상당히 느리기도 하다.
2. Reference or verifier dependency
   • OPSD refinement는 expert solution 같은 privileged information의 품질에 의존한다.
   • Math에서는 boxed answer verifier가 비교적 명확하지만, open-ended agent task에서는 reference와 verifier를 정의하기 어렵다.
3. On-policy support가 hard constraint로 enforce되지 않는다
   • Raw rollout을 anchor로 사용하고 structure preservation을 instruction으로 요구한다.
   • 하지만 explicit likelihood floor, edit-distance budget, support rejection rule은 public recipe의 핵심 mechanism으로 보이지 않는다.
   • Strong teacher가 student capability에서 너무 먼 rewrite를 만들 가능성을 별도로 측정해야 한다.
4. Teacher style homogenization risk
   • 모든 rollout을 같은 teacher가 다시 쓰면 다양한 student derivation이 teacher-preferred style로 수렴할 수 있다.
   • Pass@K gain이 보고되지만, broader domain에서 semantic diversity와 surface diversity를 별도로 봐야 한다.
5. Qwen3, math, code 중심 검증
   • 다른 architecture, multimodal reasoning, tool-use agent, long-horizon interaction에서 같은 prefix failure와 gain이 재현되는지는 추가 실험이 필요하다.
6. Code generalization is not uniformly positive
   • LiveCodeBench에서 base model을 회복하지 못한 결과가 있다.
   • Easy unit-test benchmark의 gain을 recent real-world coding benchmark로 일반화하면 안 된다.
7. Pass@16 is not free inference performance
   • Pass@16은 16개 sample 중 하나가 맞을 probability다.
   • Production에서 16개 rollout과 verifier를 사용할 수 없다면 Avg@16이나 single-sample behavior가 더 중요할 수 있다.
8. Training target compression과 serving latency는 다르다
   • $y_r$가 $y_o$보다 약 9배 짧다는 것은 training corpus 분석이다.
   • Distilled model의 inference response가 같은 비율로 짧아진다는 의미가 아니다.
9. Current paper status
   • 현재 arXiv v1이고 under review 상태다.
   • Revised version에서 table, setup, prompt, release scope가 바뀔 수 있다.

7. My Take

7-1. Why this matters for my work

TRD의 가장 재사용 가능한 통찰은 loss를 더 정교하게 만들기 전에 supervision context가 올바른 state sequence를 지나고 있는지 먼저 확인해야 한다는 점이다.

많은 post-training failure는 loss coefficient 문제처럼 보인다.

• KL이 너무 크다.
• Entropy가 너무 낮다.
• 특정 token이 gradient를 지배한다.
• Teacher와 student disagreement가 사라진다.

하지만 state sequence 자체가 잘못되어 있으면 coefficient tuning은 같은 실패 context 안에서 신호의 볼륨만 바꾼다. TRD는 context distribution을 바꾼 뒤 익숙한 full-vocabulary KL을 사용한다.

이 원리는 agent training에서 특히 중요하다. Browser agent가 첫 search query를 잘못 선택하고, coding agent가 초반에 wrong file을 root cause로 정하고, robot policy가 early grasp pose를 잘못 잡으면 뒤의 action을 token 단위로 고치는 것만으로는 recovery trajectory를 학습하기 어렵다.

7-2. Reuse potential

TRD를 다른 domain에 옮긴다면 다음 component가 필요하다.

Component	Math TRD	General agent adaptation
Raw trajectory	Reasoning response	Observation, thought, tool action sequence
Teacher refinement	Solution rewrite	Trajectory repair or replanning
Reference	Expert solution	Goal state, successful trace, environment feedback
Verifier	Final answer checker	Unit test, task reward, simulator, judge
Distillation	Token KL on $y_r$	Action/token KL on repaired trajectory

실무에서는 raw teacher rewrite를 그대로 쓰기보다 quality gate를 추가하는 것이 좋다.

1. Outcome gate
   • $y_r$가 verifier를 통과할 때만 사용한다.
2. Support gate
   • Student likelihood, edit distance, representation distance를 측정해 지나치게 먼 rewrite를 제거한다.
3. Minimal-edit preference
   • Correct prefix와 useful subroutine을 최대한 유지하도록 teacher prompt와 score를 설계한다.
4. Diversity gate
   • 모든 trajectory가 같은 template로 collapse하지 않는지 semantic cluster와 lexical diversity를 본다.
5. Cost-aware routing
   • Correct and concise $y_o$는 그대로 사용하고, failed or inflated rollout만 teacher rewrite에 보낸다.

이렇게 하면 teacher compute를 prefix failure probability가 높은 sample에 집중할 수 있다.

7-3. What I would measure next

후속 실험에서는 다음 질문이 중요하다.

1) Rewrite distance와 gain의 관계

$y_o$와 $y_r$ 사이의 edit distance 또는 model likelihood gap이 너무 작으면 correction이 부족하고, 너무 크면 off-policy expert imitation에 가까워질 수 있다. Performance가 가장 좋은 distance range가 존재하는지 확인할 필요가 있다.

2) Error-localized refinement

Full trajectory rewrite 대신 verifier나 critic이 failure span을 찾고 해당 suffix만 다시 생성하면 teacher cost와 style homogenization을 줄일 수 있다.

\[y_r = y_{o,\lt t^*} \oplus \operatorname{Rewrite}(y_{o,\ge t^*})\]

여기서 $t^*$는 estimated failure point다. 논문은 exact prefix-failure localization이 combinatorial하게 어렵다고 지적하지만, learned critic이나 process reward model로 approximation할 수 있다.

3) Online routing between $y_o$ and $y_r$

모든 sample을 rewrite하기보다 다음 score로 routing할 수 있다.

\[s(x,y_o) = \lambda_1 \cdot \text{verifier failure} + \lambda_2 \cdot \text{trajectory inflation} + \lambda_3 \cdot \text{teacher disagreement}\]

Score가 높은 sample만 TRD를 적용하면 비용 효율이 좋아질 수 있다.

4) Multi-turn environment recovery

Single-response reasoning에서는 rewrite가 전체 trajectory를 대체할 수 있다. Real environment에서는 과거 action을 되돌릴 수 없을 수 있다. 이 경우 TRD는 counterfactual corrected trace, resettable simulator, recovery policy training과 결합해야 한다.

7-4. Follow-up papers

• On-Policy Distillation of Language Models: student-generated state에서 teacher distribution을 맞추는 OPD의 기본 framing을 이해하는 데 필요하다.
• Self-Distilled Reasoner: privileged information을 사용하는 OPSD와 token loss clipping의 배경을 이해하는 데 유용하다.
• Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes: high-loss token, signal collapse, top-k 같은 token-level remedy와 비교하기 좋다.
• TIP: Token Importance in On-Policy Distillation: disagreement와 entropy를 이용한 token reweighting이 trajectory intervention과 어떻게 다른지 볼 수 있다.
• Privileged Information Distillation for Language Models: OPSD와 privileged teacher context를 broader agent setting에서 연결해 읽기 좋다.

8. Summary

• OPD는 student rollout 위에서 dense teacher signal을 제공하지만, early reasoning error가 생기면 teacher가 wrong continuation과 correction 사이의 mixture가 될 수 있다.
• Standard token-level KL은 correction onset 뒤에도 원래 failed prefix를 계속 사용하므로 multi-step recovery gradient가 fragmented된다.
• TRD는 raw rollout $y_o$를 teacher가 refined trajectory $y_r$로 다시 쓴 뒤, $y_r$의 prefix 위에서 full-vocabulary forward KL을 수행한다.
• Math에서는 특히 difficult benchmark의 Pass@K와 reasoning coverage가 크게 개선되고, refined trajectory는 더 높은 verifier pass rate와 훨씬 짧은 median length를 보인다.
• Extra teacher rollout cost, LiveCodeBench regression, support enforcement, verifier dependency를 함께 고려해야 하며, 핵심 takeaway는 loss보다 먼저 supervision trajectory를 고치는 것이다.