Self-Distilled Policy Gradient

기술 요약: 자기 증류 정책 경사 (Self-Distilled Policy Gradient, SDPG)

문제 정의
수학 및 코드 생성과 같은 추론 작업에서 대규모 언어 모델(LLM)의 사후 학습을 위한 표준으로 검증 가능한 보상 기반 강화 학습(RLVR)이 자리 잡았습니다. GRPO(Group Relative Policy Optimization)와 같은 알고리즘은 규칙 기반 결과 보상(예: 정답/오답)에 최적화하며, 인간의 선호도 주석을 필요로 하지 않습니다. 그러나 RLVR은 두 가지 주요 한계점에 직면해 있습니다:

1. 희소한 보상 (Sparse Rewards): 보상 신호가 시퀀스 단위로 적용되어 특정 추론 단계(토큰)에 대한 공로를 할당하기 어렵게 만들며, 이는 비효율적인 탐색으로 이어집니다.
2. 학습 불안정성 (Training Instability): 학습 초기 단계에서 음의 어드밴티지(negative advantages)로 인해 불안정성이 발생하며, 기초적인 언어 능력이 저하될 위험이 있습니다.

온-폴리시 증류(On-Policy Distillation, OPD)가 조밀한 토큰 수준의 감독을 제공하기 위해 제안되었으나, 전통적인 방식은 외부의 더 큰 교사 모델에 의존하여 높은 메모리 비용을 초천하고 분포 불일치를 유발합니다. 모델이 특권적 문맥(privileged context, 예: 정답 솔루션)을 사용하여 학생이자 교사 역할을 수행하는 최근의 "자기 증류(self-distillation)" 방식은 메모리 문제를 해결하지만, 적절히 제어되지 않을 경우 전역적으로 틀린 궤적에 대해 국소적으로 그럴듯한 토큰을 강화할 위험이 있습니다. 또한, 순수 자기 증류는 모드 붕괴(mode collapse)와 탐색 감소를 초래할 수 있습니다.

방법론: SDPG 프레임워크
본 논문은 정밀한 전체 어휘 온-폴리시 자기 증류와 KL 정규화 정책 최적화를 통합하는 자기 증류 정책 경사(Self-Distilled Policy Gradient, SDPG) 프레임워크를 제안합니다. 핵심 목적 함수는 세 가지 구성 요소를 결합합니다:

1. 결과 기반 정책 경사 ($L_{out}$):
   SDPG는 검증기로부터 오는 희소한 이진 보상 신호를 유지합니다. 이는 GRPO와 유사하게 그룹 내 롤아웃(rollout) 내에서 정규화된 그룹 상대적 어드밴티지($A_{out}$)를 활용하여 학습을 안정화하지만, PPO 스타일의 중요도 비율 클리핑(importance ratio clipping) 없이 직접적인 REINFORCE 스타일의 대리 손실(surrogate loss)로 공식화됩니다.

2. 정밀한 전체 어휘 온-폴리시 증류 ($L_{OPD}$):
   교사 신호를 근사하기 위해 토큰을 샘플링하는 대신, SDPG는 추가적인 문맥 $c$(예: 솔루션 또는 추론 경로)가 주어진 상태에서 학생 분포($p_t$)와 특권적 교사 분포($q_t$) 사이의 정밀한 역(reverse) Kullback-Leibler (KL) 발산(divergence)을 계산합니다.

• 그래디언트 항등식 (Gradient Identity): 저자들은 핵심적인 이론적 결과(Proposition 3.1)를 도출했습니다. 즉, 역-KL 전체 어휘 손실의 학생 정책에 대한 그래디언트는 로컬하게 디태치드 샘플링(detached-sampling) 정책 경사 업데이트와 동일합니다. 이 업데이트에서의 "어드밴티지"는 교사와 학생 확률의 중심화된 로그 비율(centered log-ratio)입니다.
• 게이팅 메커니즘 (Gating Mechanism): 특권적 교사가 잘못된 궤적을 강화하는 것을 방지하기 위해, 증류 손실은 결과 어드밴티지에 의해 게이팅됩니다. 증류는 롤아웃이 올바른 것으로 검증되었을 때($A_{out} > 0$)만 적용됩니다.

3. 참조 정책 정규화 ($L_K$):
   정책이 초기 능력에서 너무 멀어지거나 붕괴되는 것을 방지하기 위해, SDPG는 고정된 참조 정책($\pi_{ref}$)에 업데이트된 정책을 고정합니다. 본 논문은 정규화를 위해 비정규화 KL (Unnormalized KL, UKL) 발산(순방향 또는 역방향)을 사용하며, 이는 비정규화된 척도를 다룰 때 표준 정규화 KL보다 더 우아하고 대칭적인 그래디언트를 제공합니다.

안정화 전략
특권적 신호의 노이즈와 탐색 및 활용 간의 트레이드오프를 관리하기 위해, SDPG는 두 가지 안정화 장치를 도입합니다:

• 양의 어드밴티지 게이팅 (Positive Advantage Gating): 검증기에 의해 승인된 궤적에서만 증류가 발생하도록 보장합니다.
• 웜업-감쇠 스케줄 ($\beta(k)$): 증류 계수 $\beta$를 점진적으로 증가(웜업)시켜 모델이 증류에 의존하기 전에 올바른 궤적을 찾을 수 있도록 하고, 학습 종료 시점에는 이를 감소(감쇠)시킵니다. 이는 모델이 필요한 지식을 내재화한 후 특권적 신호가 모델을 과도하게 제약하는 것을 방지하여 최종 단계의 탐색을 촉진합니다.

주요 기여

1. 이론적 도출: 역-KL 전체 어휘 온-폴리시 증류의 정밀한 로컬 정책 경사 형태를 도출하여, 이것이 중심화된 로그 비율 어드밴티지를 가진 디태치드 샘플링 업데이트와 동일함을 입증했습니다.
2. 새로운 목적 함수: 이진 결과 보상, 게이팅된 전체 어휘 특권 증류, 그리고 고정된 참조에 대한 비정규화 KL 정규화를 결합한 통합된 목적 함수를 제안합니다.
3. 안정화 기술: 밀도 높은 감독과 탐색 사이의 균형을 맞추기 위해 양의 어드밴티지 게이팅과 증류 계수를 위한 웜업-감쇠 스케줄을 도입했습니다.

실험 결과
저자들은 Qwen3-4B를 사용하여 수학적 추론 벤치마크(AIME2024, AIME2025, AMC23)에서 SDPG를 평가했습니다.

• 성능: 두 가지 SDPG 변형(Unnormalized Forward KL 및 Unnormalized Reverse KL 사용) 모두 모든 벤치마크에서 베이스라인인 GRPO와 자기 증류 베이스라인인 RLSD를 능가했습니다.
• 안정성: SDPG 변형들은 GRPO보다 훨씬 빠르게 높은 보상 고점에 도달했습니다. 특히, SDPG는 RLSD가 약 250 스텝 부근에서 모드 붕괴의 징후인 엔트로피 붕괴를 보인 것과 대조적으로, 훈련 전반에 걸쳐 더 높은 액터 엔트로피(actor entropy)를 유지했습니다.
• 효율성: SDPG 방법은 다단계 추론에 충분하면서도 GRPO의 장황한 출력보다 더 간결한 응답 길이를 생성했습니다.

의의
본 논문은 SDPG가 희소 보상 RLVR과 밀도 높은 감독의 자기 증류 사이의 간극을 성공적으로 메웠다고 주장합니다. 특권적 문맥을 통한 자기 증류를 정책 경사 업데이트로 취급하고, 게이팅과 스케줄링을 통해 특권적 신호를 엄격하게 제어함으로써, SDPG는 외부 교사의 메모리 오버헤드나 순수 자기 증류의 불안정성 없이도 탁월한 안정성과 성능을 달로 달성합니다. 이는 모델 자신의 특권적 문맥에서 유도된 밀도 높은 토큰 수준의 감독이 추론 모델을 효과적으로 가이드할 수 있음을 보여줍니다.