Trust-Region Behavior Blending for On-Policy Distillation Review

0. Introduction

Paper link

Trust-Region Behavior Blending, 줄여서 TRB는 OPD를 “어떤 loss로 teacher를 맞출 것인가”보다 “초기 student rollout을 어디까지 teacher 쪽으로 끌어와도 되는가”의 문제로 다시 보는 논문이다. OPD는 student가 자기 policy로 만든 prefix 위에서 teacher distribution을 맞춘다는 점이 장점이다. 하지만 학습 초반에는 student rollout 자체가 너무 약해서, teacher가 좋은 supervision을 줘도 그 supervision이 걸리는 state가 좋지 않을 수 있다.

이 논문의 흥미로운 점은 objective를 크게 바꾸지 않는다는 데 있다. per-prefix reverse KL OPD loss는 그대로 둔다. 대신 rollout을 수집하는 behavior policy만 초기 warmup 동안 teacher 쪽으로 제한적으로 이동시킨다. 이때 제한은 student-centered KL trust region으로 둔다. 즉 teacher를 그대로 따라가지는 않고, student distribution에서 너무 멀어지지 않는 선에서만 teacher-like sampling을 허용한다.

한 줄 요약: TRB는 OPD 초반에 student rollout이 너무 약한 문제를 해결하기 위해, teacher에 가까우면서도 student-centered KL trust region 안에 있는 behavior policy를 사용하고, warmup이 끝나면 다시 pure student rollout으로 돌아가는 rollout-side warmup 방법이다.

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

• OPD의 핵심 실패 지점이 loss가 아니라 초기 rollout distribution 에 있을 수 있다는 점을 분리해서 보여준다.
• teacher token injection, SFT warmup, temperature warmup 같은 baseline과 비교해 teacher guidance를 얼마나 넣고 언제 빼야 하는가 를 더 명시적으로 다룬다.
• Qwen3 math reasoning distillation setting에서 작은 평균 gain이지만, OPD recipe 설계 관점에서는 꽤 실용적인 signal을 준다.

이 논문의 메시지는 단순하다. OPD에서 on-policy라는 성질은 중요하지만, 학습 초반의 on-policy rollout이 항상 좋은 training state를 만든다는 뜻은 아니다. 그래서 초기에는 teacher가 student를 약하게 끌어주고, 곧바로 그 보조 장치를 제거하는 쪽이 더 안정적인 절충점이 될 수 있다.

1. Problem Setting

1-1. Problem definition

이 논문이 겨냥하는 문제는 OPD early training instability 다.

OPD는 offline KD의 prefix mismatch를 줄이기 위해 student가 직접 생성한 prefix에서 teacher를 맞춘다. 이를 간단히 쓰면 다음과 같은 형태다.

\[L_{OPD}(\theta) = \mathbb{E}_{h \sim d_{\pi_\theta}} \left[ D_{KL}(\pi_\theta(\cdot \mid h) \Vert \pi_T(\cdot \mid h)) \right]\]

여기서 $h$는 student rollout으로 방문한 prefix이고, $\pi_T$는 teacher policy다. 이 구조는 inference time에 student가 실제로 방문할 state에서 학습한다는 점에서 자연스럽다. 문제는 학습 초반이다.

• student가 아직 약하면 rollout prefix가 off-task이거나 low-quality일 수 있다.
• teacher는 강하더라도, 너무 나쁜 prefix 위에서는 useful local signal을 주기 어렵다.
• teacher intervention을 세게 넣으면 prefix quality는 좋아지지만, OPD가 원래 지키려던 student rollout distribution에서 멀어진다.
• 따라서 핵심은 초기 rollout을 개선하되, OPD의 on-policy 성질을 완전히 깨지 않는 것 이다.

논문은 이 문제를 rollout behavior policy의 선택 문제로 본다. 즉 학습 loss를 바꾸는 대신, 어떤 policy로 prefix를 수집할지 를 다시 설계한다.

1-2. Why previous approaches are insufficient

기존 접근은 크게 세 방향으로 볼 수 있다.

Approach	What it changes	Limitation
Vanilla OPD	student rollout 그대로 사용	early rollout이 약하면 teacher signal이 걸리는 state 자체가 나쁨
Veto or target reformulation	visited prefix에서 target distribution 수정	prefix distribution 자체는 이미 student가 방문한 뒤임
SKD style teacher injection	rollout token 일부를 teacher token으로 교체	teacher 개입이 직접적이고 off-policy shift를 제어하기 어려움
SFT warmup	초기 단계를 teacher-generated response로 대체	prompt stream은 맞출 수 있지만 online rollout dynamics는 사라짐
Temperature warmup	sampling temperature를 낮춤	더 보수적인 rollout은 만들지만 teacher closeness를 직접 optimize하지 않음

TRB가 보완하려는 지점은 분명하다. visited prefix에서 target을 고치는 것이 아니라, prefix를 만드는 behavior policy를 바꾼다. 다만 teacher로 완전히 갈아타지 않고, student 주변 trust region 안에서만 움직인다. 그래서 이 논문은 OPD의 target-side method라기보다 rollout-side control method 에 가깝다.

2. Core Idea

2-1. Main contribution

TRB의 기여는 다음 4가지로 정리할 수 있다.

1. Behavior policy를 별도 optimization 대상으로 둔다.
   • student update loss와 rollout policy를 분리한다.
   • update는 기존 reverse KL OPD를 유지한다.
   • rollout collection만 warmup 동안 teacher-guided behavior policy로 바꾼다.
2. student-centered KL trust region을 둔다.
   • behavior policy는 teacher에 가까워지려고 한다.
   • 하지만 student policy에서 허용된 KL budget 이상 벗어나지 못한다.
   • 이 구조가 teacher guidance와 on-policy closeness 사이의 절충을 만든다.
3. closed-form behavior blend를 사용한다.
   • teacher와 student distribution의 exponential interpolation 형태로 behavior policy를 만든다.
   • feasible coefficient는 KL constraint를 만족하도록 binary search로 찾는다.
4. warmup 이후 pure OPD로 되돌린다.
   • KL budget을 선형으로 anneal한다.
   • warmup이 끝나면 behavior policy는 다시 student policy가 된다.
   • teacher guidance는 persistent regularizer가 아니라 early rollout stabilizer다.

2-2. Design intuition

이 논문을 이해하는 가장 쉬운 방법은 OPD를 두 단계로 나눠보는 것이다.

• 어떤 prefix를 방문할 것인가.
• 방문한 prefix에서 teacher distribution을 어떻게 맞출 것인가.

기존 OPD는 두 단계를 모두 student에 맡긴다. TRB는 두 번째 단계는 그대로 두되, 첫 번째 단계만 warmup 동안 살짝 바꾼다. 이때 behavior policy $b$는 다음 문제를 푼다.

\[b^*(\cdot \mid h) = \arg\min_b D_{KL}(b(\cdot \mid h) \Vert \pi_T(\cdot \mid h))\]

subject to

\[D_{KL}(b(\cdot \mid h) \Vert \pi_S(\cdot \mid h)) \le \epsilon_t\]

직관적으로는 teacher에 가장 가까운 sampling distribution을 고르되, student에서 너무 멀리 가지 못하게 하는 것이다. 그 결과 behavior policy는 아래처럼 teacher와 student 사이의 log-space blend로 이해할 수 있다.

\[b_\lambda(y \mid h) = \frac{\pi_S(y \mid h)^{1 - \lambda} \pi_T(y \mid h)^\lambda}{Z_\lambda(h)}\]

여기서 $\lambda$는 trust region constraint를 만족하는 범위에서 가능한 가장 큰 teacher strength로 선택된다. $\epsilon_t = 0$이면 $\lambda = 0$이고, behavior policy는 student와 같다. teacher 자체가 feasible하면 teacher에 더 가까운 policy를 쓸 수 있다.

이 설계의 핵심은 teacher guidance를 token replacement가 아니라 constrained distribution movement로 다룬다 는 점이다. 그래서 intervention이 hard switch가 아니라, prefix마다 student와 teacher 사이에서 budgeted interpolation으로 결정된다.

3. Architecture / Method

3-1. Overview

Item	Description
Goal	early OPD rollout을 teacher-supported state로 이동시키되 on-policy distribution에서 너무 멀어지지 않게 제어
Student update	per-prefix reverse KL OPD 유지
Rollout behavior	warmup 동안 TRB behavior policy 사용
Constraint	student-centered KL trust region
Solver	closed-form log-space blend + binary search
Schedule	KL budget을 warmup 동안 anneal to zero
Baselines	Vanilla OPD, Veto, SKD, Temperature warmup, SFT warmup, Fixed blending

3-2. Module breakdown

1) Per-prefix behavior policy

TRB는 매 prefix $h$에서 next-token behavior policy를 만든다. 이 policy는 teacher에 가까워지려고 하지만, student-centered KL budget을 넘지 않는다.

이 관점에서 teacher는 직접 token을 강제로 넣는 주체가 아니라, behavior policy의 direction을 정하는 reference distribution이다. student는 반대로 trust region의 중심이다. 그래서 TRB는 teacher-following과 student-following 사이의 중간 policy를 매 prefix마다 계산한다.

중요한 점은 update loss가 바뀌지 않는다는 것이다. sampled prefix가 정해지면, student는 여전히 reverse KL OPD objective로 teacher를 맞춘다. 따라서 TRB의 intervention surface는 rollout collection이다.

2) Closed-form blend and binary search

제약 최적화는 teacher와 student distribution의 exponential interpolation 형태로 풀린다.

\[b_\lambda(y \mid h) \propto \pi_S(y \mid h)^{1 - \lambda} \pi_T(y \mid h)^\lambda\]

논문은 feasible한 가장 큰 $\lambda$를 선택한다고 설명한다. 만약 teacher distribution이 trust region 안에 있으면 teacher에 가까운 policy를 그대로 쓸 수 있고, 그렇지 않으면 binary search로 constraint를 만족하는 최대 coefficient를 찾는다.

이게 단순 weighted average보다 중요한 이유는 log probability space에서 teacher와 student를 섞기 때문이다. probability를 선형으로 섞는 방식보다 KL geometry와 잘 맞고, trust region constraint를 직접 해석하기 쉽다.

3) Annealed warmup

TRB는 이 behavior policy를 전체 학습 동안 계속 쓰지 않는다. 논문은 KL budget을 warmup horizon 동안 줄이다가 0으로 만든다.

\[\epsilon_t = \epsilon_0 \max \left(0, 1 - \frac{t}{T_{warmup}} \right)\]

warmup 초기에는 teacher 쪽으로 더 움직일 수 있고, 시간이 갈수록 허용 deviation이 줄어든다. 마지막에는 $\epsilon_t = 0$이므로 pure student rollout으로 돌아간다.

이 부분이 fixed blending과의 핵심 차이다. fixed blending은 같은 solver를 쓰더라도 budget을 계속 유지한다. 논문 결과에서는 같은 solver를 쓰는 fixed blending보다 annealed TRB가 더 좋은 평균 결과를 낸다. 즉 teacher guidance의 존재 자체보다 언제 제거하는가 가 중요하다는 해석이 가능하다.

4) Relationship to SKD and SFT warmup

TRB는 SKD와 비슷하게 teacher를 rollout에 개입시킨다. 하지만 방식이 다르다.

• SKD는 student token이 teacher top-k 밖이면 teacher sample로 교체한다.
• TRB는 token을 사후 교체하지 않고, 처음부터 constrained behavior distribution에서 sample한다.
• SFT warmup은 online rollout을 teacher-generated response로 대체한다.
• TRB는 여전히 rollout-time sampling 구조를 유지한다.

따라서 TRB는 teacher intervention을 가장 직접적인 hard replacement가 아니라, student-local behavior shaping 으로 만든다.

4. Training / Data / Recipe

4-1. Data

논문은 두 가지 Qwen3 math reasoning distillation setting을 사용한다.

Student	Teacher	Evaluation suite
Qwen3-1.7B-Base	Qwen3-8B	MATH500, Olympiad, AMC, AIME24, AIME25
Qwen3-0.6B-Base	Qwen3-4B	GSM8K, MATH500, Olympiad, AMC

Training prompt는 OpenThoughts3-1.2M corpus에서 25,600개를 sampling한다. 모든 training input에는 step-by-step reasoning과 boxed final answer를 요구하는 system prompt를 붙인다.

Reward는 math-verify를 사용한다. correct solution은 1.0, incorrect solution은 0.0으로 처리된다. 다만 TRB의 학습 objective 자체는 RL reward objective가 아니라 reverse KL OPD다. 이 reward는 evaluation and diagnostics에서 correctness를 측정하는 데 쓰인다.

4-2. Training strategy

공통 training recipe는 아래와 같다.

Parameter	Value
Optimizer	AdamW
Weight decay	0.01
Gradient norm clipping	1.0
LR scheduler	Constant
Warmup	15-step cosine warmup
Max prompt length	1024
Max response length	7168
Global batch size	64
Rollouts per prompt	4
PPO epochs	1
Training epochs	1
KL loss type	reverse KL
KL top-k support	16 tokens
Rollout temperature	1.0

구현 stack은 verl pipeline과 SGLang rollout generation을 사용하고, 모든 run은 FSDP2와 8 NVIDIA H100 GPUs에서 수행된다. 논문은 student와 teacher tokenizer의 EOS token id mismatch도 canonicalization해서 behavior construction과 KL evaluation에서 같은 stop event로 비교한다.

4-3. Engineering notes

TRB를 실제로 구현할 때 가장 중요한 비용은 online teacher decoding이다. Vanilla OPD에서는 student rollout을 먼저 만든 뒤 teacher log probability를 batched pass로 계산할 수 있다. TRB는 generation step마다 teacher policy가 필요하므로 student와 teacher가 동시에 resident해야 한다.

즉 TRB는 아래 비용을 추가한다.

• warmup 동안 teacher online forward 필요
• student and teacher co-residency 필요
• teacher KV cache까지 peak memory에 포함
• binary search and log-space blend overhead는 transformer forward에 비해 작음

논문은 teacher FLOPs order 자체는 표준 OPD와 비슷하지만, vanilla OPD의 batched teacher pass보다 wall-clock에서 불리할 수 있다고 설명한다. 따라서 TRB는 serving-friendly trick이라기보다 training warmup trick에 가깝다.

5. Evaluation

5-1. Main results

주요 결과는 Table 1의 pass@1이다. 핵심 평균만 보면 아래와 같다.

Setting	TRB	Vanilla OPD	SKD	Temperature warmup	SFT warmup	Fixed blending
Qwen3-1.7B <- Qwen3-8B	33.2	32.3	32.7	32.8	32.2	32.6
Qwen3-0.6B <- Qwen3-4B	44.4	44.0	44.2	44.0	43.4	43.8

결과를 크게 읽으면 TRB가 두 setting 모두에서 최고 평균을 낸다. 단, margin은 아주 크지 않다. 이 논문은 압도적 성능 향상 논문이라기보다, OPD warmup에서 teacher guidance를 넣는 방법을 더 정교하게 비교한 논문에 가깝다.

개별 benchmark에서는 다른 baseline이 이기는 column도 있다. 예를 들어 Table 1에서 Veto나 fixed blending이 특정 AIME or AMC column에서 강한 경우가 있다. 하지만 평균 기준으로는 TRB가 두 model-pair setting 모두에서 가장 높다.

5-2. What really matters in the experiments

이 논문의 실험에서 중요한 포인트는 최종 평균보다 training dynamics다.

첫째, fixed blending보다 TRB가 낫다. 같은 per-prefix solver를 쓰더라도 teacher guidance를 계속 유지하면 최종 평균이 떨어진다. 이는 teacher intervention이 초반에는 유용하지만, 후반까지 남으면 student trajectory distribution을 방해할 수 있다는 해석을 가능하게 한다.

둘째, Figure 3은 TRB warmup 동안 teacher token-mean entropy가 낮아지고, warmup 이후 vanilla OPD와 비슷해진다고 보고한다. 즉 TRB가 training 전체를 다른 regime으로 바꾸는 것이 아니라, 초반 state distribution을 더 teacher-supported하게 만드는 역할을 한다.

셋째, Figure 4의 step-0 probe는 같은 truncation length에서 TRB prefix가 vanilla OPD prefix보다 teacher and student continuation 모두에서 더 높은 success를 보인다고 설명한다. 이는 TRB가 단순히 teacher probability만 올리는 것이 아니라, 실제로 더 useful한 early prefix를 만드는 쪽에 가깝다는 증거로 읽을 수 있다.

넷째, Appendix diagnostic에서는 trust-region budget이 커질수록 mean teacher log probability와 verifier reward가 올라가지만, teacher-support score의 AUROC는 내려갈 수 있다고 보고한다. 이건 high teacher support와 separability가 항상 같은 의미가 아니라는 점을 보여준다.

6. Limitations

1. math reasoning scope가 좁다.
   • 실험은 두 개 Qwen3-Base student-teacher pair와 math reasoning benchmark에 집중한다.
   • coding, open-ended instruction following, tool use, long-context QA로 그대로 일반화할 수 있는지는 원문 기준 추가 검증이 필요하다.
2. 성능 gain이 크지는 않다.
   • 평균 gain은 분명히 있지만, benchmark별로는 다른 baseline이 이기는 column도 있다.
   • 따라서 TRB를 universal OPD upgrade로 읽기보다는 warmup design의 한 candidate로 봐야 한다.
3. training-time cost가 증가한다.
   • warmup 동안 teacher online decoding과 student-teacher co-residency가 필요하다.
   • teacher model이 훨씬 크거나 memory budget이 빡빡하면 실험보다 비용이 더 커질 수 있다.
4. trust-region schedule tuning이 남는다.
   • initial KL budget과 warmup horizon이 핵심 hyperparameter다.
   • 논문은 sweep으로 좋은 setting을 찾지만, 다른 model gap에서는 다시 tuning이 필요할 수 있다.
5. loss-side OPD 문제를 모두 해결하지 않는다.
   • TRB는 rollout-side method다.
   • teacher target quality, top-k KL approximation, tokenizer mismatch, reward verification 같은 요소는 별도 관리가 필요하다.

7. My Take

7-1. Why this matters for my work

OPD를 실제 post-training recipe로 쓰려면, loss 하나보다 rollout distribution이 더 중요해지는 순간이 있다. 이 논문은 그 지점을 꽤 선명하게 잡는다.

TRB의 진짜 가치는 성능 숫자보다 design principle에 있다. OPD early phase에서는 pure on-policy를 고집하는 것보다, teacher에 가까운 state를 제한적으로 방문하게 하는 것이 더 나을 수 있다. 하지만 그 보조 장치는 반드시 사라져야 한다. 이게 fixed blending 대비 TRB가 주는 가장 실용적인 교훈이다.

7-2. Reuse potential

TRB는 아래 상황에서 특히 재사용 가치가 있어 보인다.

• small student가 large teacher를 따라가야 하는 reasoning OPD
• student rollout이 초반에 off-task로 자주 무너지는 setting
• teacher token injection은 너무 강하고, temperature warmup은 너무 약한 setting
• OPD와 RLVR 사이의 warmup stage를 설계해야 하는 setting
• teacher and student tokenizer mismatch가 있는 distillation setting

실무적으로는 다음 실험을 해볼 만하다.

1. TRB warmup 후 standard OPD를 적용한다.
2. TRB warmup 후 sampled-token OPD or top-k OPD를 비교한다.
3. teacher entropy, student-teacher KL, verifier reward를 warmup 구간에서 같이 추적한다.
4. fixed blending, SKD, temperature warmup과 동일 compute budget으로 비교한다.
5. teacher online decoding cost를 amortize할 수 있는 rollout server 구조를 따로 설계한다.

7-3. Follow-up papers

• Rethinking On-Policy Distillation of Large Language Models: Phenomenology, Mechanism, and Recipe
• Trust Region On-Policy Distillation
• Draft-OPD: On-Policy Distillation for Speculative Draft Models
• MiniLLM: Knowledge Distillation of Large Language Models
• On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes

8. Summary

• TRB는 OPD의 per-prefix reverse KL loss를 바꾸지 않고, early rollout behavior policy만 teacher-guided trust region 안에서 바꾼다.
• behavior policy는 teacher에 가까워지되 student-centered KL budget을 넘지 않는 log-space blend로 계산된다.
• KL budget은 warmup 동안 anneal to zero되고, 이후에는 pure student rollout으로 돌아간다.
• 두 Qwen3 math reasoning distillation setting에서 TRB는 평균 pass@1 기준 비교 method 중 가장 높은 결과를 보인다.
• 다만 gain은 modest하고, online teacher decoding 비용이 있으므로, TRB는 general solution이라기보다 OPD warmup design primitive로 보는 편이 맞다.