Pointwise Mutual Information 기반의 Reasoning RL을 위한 Anti-Self-Distillation

On-policy self-distillation은 외부 교사 없이 토큰 단위의 imitation 신호를 제공하는 방향을 제시하지만, 수학 추론에서의 성능은 불안정하다. PMI 분석은 privileged context가 per-token 신호를 왜곡하는 구조적 바이어스를 밝히고, AntiSD는 gradient 방향을 역전시켜 deliberation 토큰의 기여를 회복한다. 다섯 모델에서 GRPO 대비 학습 속도 향상과 최종 정확도 향상을 보여준다.

출발점은 on-policy self-distillation에서 teacher가 privileged context로 토큰 확률을 조정한다는 점이다. ut = log πT(yt | x, y<t) − log πS(yt | x, y<t)로 정의되는 per-token 신호는 PMI(yt; c | x, y<t)로 해석되며, 이 신호는 privileged context가 특정 토큰의 확률을 높이거나 낮춘다는 정보를 담고 있다. 이로 인해 shortcut 토큰은 과대평가되고 deliberation 토큰은 과소평가되는 구조적 바이어스가 발생한다. AntiSD는 이 신호의 부호를 역전시켜 델리버레이션 토큰의 기여를 증가시키고, JSD를 사용해 두 분포 간 차이의 상한을 갖도록 조절한다. 엔트로피 게이트는 교사 정보가 충분히 남아 있을 때만 AntiSD 신호를 작동시키고, 교사 엔트로피가 낮아지면 자동으로 비활성화한다.

1. 퍼-token 수준의 신호 δt를 ut = tt − st로 정의하고, δt를 −φ(ut)로 바꿔 역전 방향으로 학습을 진행한다. 2) f-divergence로 정의된 DJSD(πS∥πT)에서 φ(u) = 1/2(softplus(u) − log 2)로 정의된 대역형를 사용해 AAntiSDt = −φ(ut) 형태의 어드밴티지를 얻는다. 3) 엔트로피 게이트 g를 도입해 H이 Hwarm 이상일 때만 λmax를 곱한 AntiSD를 활성화하고, τdown 이하로 떨어지면 비활성화한다. 4) 알고리즘 1에 따라 학습 스텝에서 SAD를 적용하고, GRPO와의 결합 형태를 정책기울기로 업데이트한다.

다섯 모델에서 AntiSD가 GRPO의 최고점에 도달하는 속도를 2–10× 가속했고, Avg 기준으로 최종 정확도는 최대 +11.5pp 상승했다. 특히 Qwen3-8B, Qwen3-4B-IT-2507, Olmo3-7B-IT에서 큰 개선이 나타났다. HMMT25/ AIME25 등 벤치마크에서 AntiSD가 GRPO 대비 안정적으로 우위를 유지했으며, pass@k 그래프에서도 AntiSD의 리드가 유지되었다. SD는 교사 엔트로피가 0.1 nat 수준으로 수렴하거나 1 nat를 초과하는 경향이 있어 학습 신호가 불안정했고, No-gate 변형은 모델에 따라 불안정한 구동을 보였다.

LSD의 per-token KL을 역전시키는 AntiSD를 도입하고, DJSD의 gradient를 이용해 각 토큰의 어드밴티지를 φ(ut)로 조절한다. ut은 PMI(yt; c | x, y<t)로 해석되며, 텍스트의 shortcut 토큰의 신호를 상한하고 deliberation 토큰의 신호를 상향한다. 게이트는 엔트로피 H를 기반으로 작동하며, Hwarm 이상에서 λmax를 곱한 AntiSD 신호를 활성화하고 H이 τdown 아래로 떨어지면 비활성화한다. 실험에서 4B–30B 파라미터 규모의 Qwen3/Olmo-3 계열 모델에 적용되었다.

논문은 PMI 기반의 지역적 분석에 기반하며, 전체 최적해를 보장하는 전역적 분석은 제시하지 않는다. Evaluation은 수학 추론 및 코드 추론 벤치마크에 한정되며, 다중 턴 에이전트 설정이나 멀티도메인 확장에 대해서는 추가 연구가 필요하다.