Trust-Region Behavior Blending for On-Policy Distillation

기술 요약: 온폴리시 증류를 위한 신뢰 영역 행동 블렌딩 (Trust-Region Behavior Blending)

문제 정의

온폴리시 증류(On-Policy Distillation, OPD)는 오프라인 지식 증류의 고질적인 문제인 프리픽스 불일치(prefix mismatch)를 해결하기 위해, 학생 모델이 현재 자신의 정책에서 샘플링한 프리픽스를 기반으로 학습하면서 더 강력한 교사 모델을 모방하도록 합니다. 그러나 이 방식은 학습 초기 단계에서 치명적인 취약성을 드러냅니다. 즉, 약한 학생 모델은 저품질의 프리픽스를 생성하게 됩니다. 학생의 궤적 분포(trajectory distribution)가 좋지 않을 때, 교사의 감독이 저품질 또는 "사용 불가능한" 프리픽스에 적용되면서 학습 신호의 효과가 제한됩니다.

이 과정을 안정화하려는 기존의 시도들은 종로를 오프폴리시(off-policy)로 이동시키거나(예: 교사 토큰 주입 또는 타겟 분포 변경), 이는 타겟 학습 분포를 교란할 수 있습니다. 반대로, 순수 학생 롤아웃(pure student rollouts)은 분포를 보존하지만 초기 품질이 낮다는 문제가 있습니다. 핵심 과제는 근본적인 프리픽스별 증류 목적을 변경하거나 롤아웃 과정을 영구적으로 오프폴리시로 전환하지 않으면서, 초기 롤아웃의 품질을 개선하는 것입니다.

방법론: 신뢰 영역 행동 블렌딩 (Trust-Region Behavior Blending, TRB)

저자들은 초기 롤아웃 수집 단계 동안 행동 정책(behavior policy)을 제어하기 위해 설계된 웜업 방법인 **신뢰 영역 행동 블렌딩(TRB)**을 제안합니다. TRB는 현재 학생 정책을 중심으로 하는 "신뢰 영역" 내에 엄격히 제한되면서도, 교사에게 최대한 가까운 행동 정책을 찾는 원리에 따라 작동합니다.

1. 프리픽스별 최적화
각 생성 프리픽스 $h$에서, TRB는 다음과 같은 제약 최적화 문제를 풀어 행동 정책 $\mu^*(\cdot | h)$를 정의합니다. 목적 함수는 교사 정책 $\pi_T$에 대한 쿨백-라이블러(KL) 발산을 최소화하는 것이며, 학생 정책 $\pi_S$에 대한 KL 발산 제약을 따릅니다.

$$\mu^*(\cdot | h) = \arg \min_{\mu} D_{KL}(\mu \parallel \pi_T) \quad \text{s.t.} \quad D_{KL}(\mu \parallel \pi_S) \le \epsilon$$

여기서 $\epsilon$은 학생으로부터 허용되는 국소적 편차를 나타냅니다. 이 공식은 샘플링 분포가 교사가 지원하는 토큰 쪽으로 끌어당겨지도록 보장하면서도, 학생 중심의 신뢰 영역 내에 머물도록 합니다.

2. 폐쇄형 해 (Closed-Form Solution)
이 최적화는 행동 정책이 학생과 교사 정책 사이의 기하학적 보간(geometric interpolation)이 되는 폐쇄형 해를 산출합니다:
$$\mu_\beta(a | h) = \frac{\pi_S(a | h)^{1-\beta} \pi_T(a | h)^\beta}{Z_\beta(h)}$$
계수 $\beta^*(h)$는 KL 제약 조건을 만족하는 $[0, 1]$ 범위 내의 가장 큰 가능한 값입니다. 만약 교사가 이미 신뢰 영역 내에 있다면 $\mu^*$는 교사가 되며, $\epsilon=0$이면 학생으로 돌아갑니다.

3. 어닐링된 웜업 (Annealed Warmup)
TRB는 초기 웜업 단계 동안에만 적용됩니다. KL 예산 $\epsilon$은 호라이즌 $K$에 걸쳐 선형적으로 0까지 어닐링됩니다:
$$\epsilon_k = \epsilon_0 \left(1 - \frac{k}{K}\right)$$
이를 통해 학습 초기에는 강한 교사 가이드를 제공하여 프리픽스 품질을 높이고, 학생이 개선됨에 따라 점진적으로 롤아웃 수집을 순수 학생 샘플링으로 복귀시킵니다. 결정적으로, 프리픽스별 역-KL 증류 손실은 변경되지 않으며, 오직 데이터 수집 정책만 수정됩니다.

주요 기여

• 새로운 웜업 전략: TRB는 단순히 토큰을 주입하거나 손실 함수를 변경하는 대신, 명시적인 학생 중심 KL 제약 하에 교사 가이드 정책을 최적화함으로써 초기 롤아웃 행동을 제어하는 방법을 도입했습니다.
• 온폴리시 특성 보존: 지속적으로 교사와 학생 정책을 혼합하거나 타겟 분포를 변경하는 방법과 달리, TRB는 훈련이 결국 순수 온폴리시 롤아웃으로 돌아오도록 보장하여 메인 훈련 단계를 위한 타겟 분포를 보존합니다.
• 이론적 정식화: 본 논문은 보간 계수에 따른 KL 발산의 단조성을 근거로, 학생의 안정성과 교사의 가이드를 균형 있게 맞추는 행동 정책에 대한 폐쇄형 해를 제공합니다.

실험 결과

저자들은 두 가지 수학적 추론 증류 설정에서 TRB를 평가했습니다:

1. Qwen3-1.7B-Base를 Qwen3-8B로부터 증류.
2. Qwen3-0.6B-Base를 Qwen3-4B로부터 증류.

성능:

• TRB는 두 설정 모두에서 비교된 모든 방법 중 가장 강력한 평균 성능을 달 achieved, vanilla OPD, fixed-$\epsilon$ 블렌딩, Veto, SKD(교사 토큰 인터리빙), temperature warmup, SFT warmup을 능가했습니다.
• Qwen3-1.7B 설정에서 TRB는 vanilla OPD의 32.3%와 비교하여 **33.2%**의 평균 pass@1을 달성했습니다.
• Qwen3-0.6B 설정에서 TRB는 **44.4%**를 달성하여 vanilla OPD의 44.0%를 앞질렀습니다.

분석:

• 초기 훈련 역학: 궤적 분석(그림 2)에 따르면, 일부 방법(SKD 등)은 초기에 더 빠르게 상승하지만 반드시 가장 높은 최종 성능을 유지하는 것은 아닙니다. TRB는 우수한 최종 결과를 이끄는 "빠른 시작(fast start)"을 제공합니다.
• 엔트로피 및 프리픽스 품질: 웜업 단계 동안 TRB는 방문된 프리픽스 상의 교사 토큰의 평균 엔트로피를 감소시키는데, 이는 더 높은 품질의 가이드를 나타냅니다. Step-0 프로브(그림 4)는 TRB 하에서 생성된 프리픽스가 교사나 학생에 의해 계속될 때 vanilla OPD 프리픽스보다 더 높은 성공률을 보임을 입증하며, 이는 TRB가 초기 상태 분포를 더 학습 가능한 궤적으로 성공적으로 이동시켰음을 시사합니다.

의의 및 주장

본 논문은 TRB가 온폴리시 목적에서 영구적으로 벗어나지 않으면서 초기 데이터 분포의 품질을 개선함으로써, 초기 온폴리시 증류의 취약성을 효과적으로 해결한다고 주장합니다.

• 제한된 범위: 저자들은 본 연구가 Qwen3-Base 쌍을 사용한 두 가지 수학적 추론 설정과 정답 기반 평가 프로토콜으로 범위가 한정되어 있음을 명시합니다. 특정 웜업 스케줄이 다른 도메인이나 교사-학생 격격에 그대로 전이된다고 주장하지 않습니다.
• 비용 트레이드오프: 논문은 온라인 교사 디코딩과 학생-교사 공존(co-residency) 요구 사항으로 인해 웜업 단계 동안 훈련 시간 비용이 증가함을 인정합니다. 그러나 이러한 비용은 일시적이며, 훈련이 웜업 호라이즌 이후 표준적이고 더 효율적인 OPD 런타임 프로필로 복귀한다는 점을 언급합니다.
• 핵심 통찰: 결과는 교사 가이드 오프폴리시 행동이 학생이 방문한 프리픽스가 교사와 크게 불일치할 때 가장 유익하다는 것을 시사합니다. 일단 이 불일치가 해결되면, 지속적인 오프폴리시 가이드는 수익 체감의 법칙을 따르게 되며, 이는 행동 정책을 학생에게로 다시 어닐링하는 전략을 뒷받침합니다.