Scaling Reasoning Efficiently via Relaxed On-Policy Distillation

이 논문은 온-정책 증류의 불안정성을 해결하고 샘플 효율성과 추론 속도를 극대화하기 위해, 교사의 보상을 유연하게 활용하는 'REOPOLD'라는 새로운 증류 프레임워크를 제안하고 수학, 시각, 도구 사용 추론 작업에서 기존 방법들을 능가하는 성능을 입증했습니다.

핵심

🎓 핵심 비유: "무조건 따라 하기 vs. 똑똑한 멘토링"

1. 문제: "무조건 따라 하는 학생" (기존 방식)

예전에는 작은 AI(학생) 가 큰 AI(선생님) 의 답변을 그대로 복사하듯 배우는 '온-폴리시 증류 (On-Policy Distillation)' 방식을 썼습니다.

• 상황: 선생님 (32B 모델) 이 "이 문제는 이렇게 풀어야 해!"라고 말하면, 학생 (7B 모델) 은 그 말 한마디 한마디를 절대적인 진리로 여겨 무조건 따라 합니다.
• 문제점:
  • 혼란: 선생님이 실수했거나, 학생이 아직 이해할 수 없는 어려운 부분을 가르치면 학생은 당황해서 망가집니다 (부정적 전이).
  • 공포: "내가 틀린 걸까?"라는 생각이 들면 학생은 스스로 생각할 기회를 잃고, 선생님의 말만 반복하다가 창의성이 사라집니다 (엔트로피 붕괴).
  • 비효율: 선생님이 "물론 1+1=2 지"라고 말한 부분까지 학생이 열심히 공부하면 시간만 낭비합니다.

2. 해결책: REOPOLD (Relaxed On-Policy Distillation)

이 논문은 **"완벽한 복제가 아니라, 유연한 학습"**을 제안합니다. 이를 REOPOLD라고 부릅니다.

🌟 REOPOLD 의 3 가지 비결 (비유 포함):

① "무서운 선생님도 때로는 웃어주는 법" (Reward Clipping)

• 상황: 선생님이 학생의 답을 보고 "너 이거 완전 틀렸어! (부정적 점수 -∞)"라고 너무 극단적으로 비난하면 학생은 겁에 질려 아무것도 못 합니다.
• 해결: REOPOLD 는 "너 이거 틀렸지만, 너무 무서워하지 마. 적당히 점수 깎아줄게."라고 부정적인 점수에도 상한선을 둡니다.
• 효과: 학생이 실수해도 멘탈이 무너지지 않고, 다시 일어설 수 있는 용기를 줍니다.

② "중요한 부분만 집중하기" (Entropy-Guided Sampling)

• 상황: 선생님이 "1+1=2"라고 말했을 때와 "이 복잡한 수학 문제의 미묘한 뉘앙스"를 설명했을 때, 학생은 둘 다 똑같이 열심히 들어야 할까요?
• 해결: REOPOLD 는 **"학생이 가장 헷갈려하고, 선생님과 생각이 다른 부분 (높은 엔트로피)"**만 골라 집중하게 합니다. 이미 아는 쉬운 부분은 건너뛰고, 진짜 고민해야 할 부분만 집중합니다.
• 효과: 공부 시간을 아껴서 진짜 어려운 문제 해결 능력을 키웁니다.

③ "탐험과 정리의 두 단계 학습" (Exploration-to-Refinement)

• 상황: 처음부터 정답만 외우면 창의성이 죽습니다.
• 해결:
  • 1 단계 (탐험): "일단 여러 가지 답을 시도해 봐! 틀려도 괜찮아!"라고鼓励学生 (SFT 방식).
  • 2 단계 (정리): "자, 이제 우리가 찾은 여러 답 중에서 가장 논리적인 것만 골라 다듬자." (RL 방식).
• 효과: 처음에는 자유롭게 사고하다가, 나중에는 그 사고를 정교하게 다듬어 완벽한 지능을 만듭니다.

---

🚀 실제 성과: "작은 AI 가 거인보다 빠르고 똑똑해지다"

이 방법을 적용한 결과, 놀라운 일들이 일어났습니다.

1. 학습 효율성 폭증: 같은 지식을 배우는 데 걸리는 데이터 양이 6.7 배~12 배나 줄었습니다. (예: 100 시간 걸리던 공부를 10 시간 만에 끝냄)
2. 작은 모델의 대박: 320 억 개의 파라미터를 가진 거대 AI(선생님) 와 맞먹는 성능을, **70 억 개 파라미터의 작은 AI(학생)**가 냈습니다.
3. 속도 향상: 작은 AI 는 크기가 작아서 계산이 빠릅니다. 그래서 거대 AI 와 똑같은 성능을 내면서도 3 배 이상 빠르게 답을 내놓습니다.

💡 한 줄 요약

"작은 AI 가 거대 AI 의 지식을 배울 때, 무조건 따라 하는 '복제'가 아니라, 실수는 용납하고 중요한 부분만 골라 배우는 '유연한 멘토링'을 통해, 더 빠르고 더 똑똑하게 성장할 수 있다."

이 논문은 AI 가 더 작고 저렴하면서도, 거대 모델 못지않은 추론 능력을 갖출 수 있는 새로운 길을 열었습니다.

기술 요약

논문 요약: Scaling Reasoning Efficiently via Relaxed On-Policy Distillation (REOPOLD)

이 논문은 제한된 용량을 가진 소형 언어 모델 (SLM) 에도 대형 언어 모델 (LLM) 의 추론 능력을 효과적으로 전이하기 위해 제안된 REOPOLD (Relaxed On-Policy Distillation) 프레임워크를 소개합니다. 기존 온-정책 증류 (On-Policy Distillation) 의 불안정성과 비효율성을 해결하고, 강화학습 (RL) 의 최신 통찰을 활용하여 학습을 안정화하고 샘플 효율성을 극대화하는 방법을 제시합니다.

---

1. 문제 제기 (Problem)

• 소형 모델의 추론 한계: 대형 모델 (DeepSeek-R1, o1 등) 은 강화학습 (RL) 과 테스트 타임 스케일링을 통해 뛰어난 추론 능력을 보이지만, 표현력이 제한된 소형 모델 (SLM) 은 직접적인 보상 최적화 (Direct RL) 가 어렵고 표준 RL 접근법이 비효율적입니다.
• 온-정책 증류의 한계: 최근 온-정책 증류 (학생이 자신의 궤적을 학습하며 교사 모델의 가이드를 받는 방식) 가 주목받고 있으나, 다음과 같은 근본적인 문제가 존재합니다.
  • 불안정성 (Instability): 학생 모델이 교사 모델과 확률 분포가 크게 다른 토큰을 생성할 때, 로그 가능도 비율 (Log-likelihood ratio) 이 $-\infty$에 수렴하여 무거운 꼬리 (Heavy-tail) 를 가진 큰 음수 보상이 발생합니다. 이는 그래디언트 폭발을 유발하고 학습을 불안정하게 만듭니다.
  • 비효율성 (Inefficiency): 대부분의 토큰은 학생과 교사 분포가 유사하여 보상이 0 에 가깝습니다. 이러한 **희미한 신호 (Vanishing rewards)**는 학습을 지연시키고 샘플 효율성을 떨어뜨립니다.
  • 엔트로피 붕괴 (Entropy Collapse): 학습 초기에 모델이 특정 답으로 빠르게 수렴하여 다양한 추론 경로를 탐색하지 못하게 됩니다.

2. 방법론 (Methodology)

저자들은 온-정책 증류를 **정책 최적화 (Policy Optimization)**의 관점에서 재해석했습니다. 교사 - 학생 로그 가능도 비율을 고정된 보상 (Reward) 으로 간주하고, 이를 RL 의 관점에서 안정화하는 REOPOLD를 제안합니다.

핵심 구성 요소:

1. Stop-Gradient 적용 및 RL 동등성:

   • 증류 목적 함수를 RL 의 정책 경사 (Policy Gradient) 형태로 재구성합니다.
   • 보상 항 ($R_{i,t}$) 에 Stop-Gradient를 적용하여, 보상이 모델 파라미터에 의존하는 것을 방지하고 이를 고정된 외적 신호로 취급합니다. 이는 분산 감소 (Variance Reduction) 를 통해 학습 안정성을 확보합니다.

2. 혼합 기반 보상 클리핑 (Mixture-Based Reward Clipping):

   • 문제 해결: 교사 확률이 0 에 가까울 때 발생하는 무거운 꼬리의 음수 보상을 제거합니다.
   • 기법: 교사와 학생 분포의 볼록 혼합 (Convex Mixture) 에서 유도된 이론적 하한값 ($\log \frac{\lambda}{1-\lambda}$) 을 사용하여 보상의 하한을 제한합니다. 이는 그래디언트 폭발을 방지하면서도 RKL 의 모드를 찾는 성격을 유지합니다.

3. 엔트로피 기반 토큰 레벨 동적 샘플링 (Entropy-Guided Token-Level Dynamic Sampling):

   • 문제 해결: 보상이 0 에 가까운 저엔트로피 (확정적) 토큰은 학습 신호가 불필요합니다.
   • 기법: 학생 모델의 엔트로피가 높은 토큰 (불확실성이 높은, 중요한 분기점) 만을 학습에 포함시키는 마스크를 적용합니다. 이는 그래디언트 희석 (Gradient Dilution) 을 막고 학습 효율을 높입니다.

4. 탐색 - 정제 다단계 학습 전략 (Exploration-to-Refinement Multi-Stage Training):

   • Phase 1 (Exploration): 학습 초기에는 음수 보상을 필터링하여 모델이 다양한 해답을 탐색하도록 유도합니다 (SFT 와 유사).
   • Phase 2 (Refinement): 학습 후기에는 엔트로피 기반 마스크를 적용하여 중요한 불확실한 토큰에 대해 정밀한 보상을 주고 정책을 정제합니다 (RL 과 유사).

3. 주요 기여 (Key Contributions)

• 증류의 이론적 분석: 온-정책 증류가 RL 과 이론적으로 동등하며, Stop-Gradient 가 제어 변수 (Control Variate) 역할을 하여 학습을 안정화함을 증명했습니다.
• REOPOLD 프레임워크: 보상 클리핑, 토큰 레벨 동적 샘플링, 다단계 학습을 통합하여 증류 과정의 불안정성을 해결하고 효율성을 극대화했습니다.
• 성능 및 확장성: 수학, 시각적 추론, 에이전트 도구 사용 등 다양한 작업에서 기존 RL 및 증류 기법보다 우수한 성능을 입증했습니다.

4. 실험 결과 (Results)

• 샘플 효율성 (Sample Efficiency):

  • 수학 추론 (AIME-25) 에서 REOPOLD 는 기존 RL 방법 (ProRL) 보다 6.7 배, 다른 증류 방법 (DeepScaleR) 보다 12 배 더 높은 샘플 효율성을 보였습니다.
  • 1.5B 모델이 300 스텝 만에 1.5B 모델 기준 12 배 효율로 수렴하는 등 빠른 학습 속도를 입증했습니다.

• 테스트 타임 스케일링 (Test-Time Scaling):

  • 시각적 추론 (Geometry3K, MathVerse) 에서 7B 학생 모델이 32B 교사 모델의 성능을 따라잡거나 능가했습니다.
  • 추론 속도 향상: 7B 모델이 32B 교사 모델 대비 약 3.3 배 빠른 추론 속도를 유지하면서 동등한 정확도를 달성했습니다.

• 다양한 작업에서의 우위:

  • 수학: AIME, MATH-500 등 다양한 벤치마크에서 SFT 및 RKL 베이스라인을 상회했습니다.
  • 시각적 추론: Qwen2.5-VL 모델 (3B, 7B) 을 대상으로 한 실험에서 NoisyRollout 및 GRPO 를 능가했습니다.
  • 에이전트 도구 사용: 이미지 조작 도구를 사용하는 Pixel-Reasoner 태스크에서도 REOPOLD 가 GRPO 보다 우수한 일반화 성능을 보였습니다.

5. 의의 및 결론 (Significance)

이 논문은 엄격한 모방 (Strict Imitation) 을 완화하는 것이 소형 모델의 추론 능력을 확장하는 핵심임을 강조합니다.

1. 안정성 확보: 기존 증류 방식의 불안정성 (음수 보상, 엔트로피 붕괴) 을 RL 기법을 차용하여 체계적으로 해결했습니다.
2. 효율성 극대화: 불필요한 학습 신호를 제거하고 중요한 토큰에 집중함으로써, 적은 데이터와 계산 자원으로 대형 모델의 능력을 소형 모델에 전이할 수 있음을 증명했습니다.
3. 실용적 가치: 7B 모델이 32B 모델 수준의 추론 능력을 가지면서도 추론 속도가 훨씬 빠르다는 점은, 실제 배포 환경에서 비용 효율적인 고성능 추론 모델 구축을 가능하게 합니다.

결론적으로, REOPOLD 는 소형 언어 모델의 추론 능력을 확장하기 위한 새로운 표준 (State-of-the-Art) 프레임워크로 자리 잡았으며, 증류와 강화학습의 경계를 허물어 효율적인 모델 학습을 위한 중요한 통찰을 제공합니다.