How to Allocate, How to Learn? Dynamic Rollout Allocation and Advantage Modulation for Policy Optimization

1. 연구 배경 및 문제 정의 (Problem)

확산형 보상 (Verifiable Rewards) 을 활용한 강화학습 (RLVR) 은 대형 언어 모델 (LLM) 의 추론 능력을 향상시키는 데 효과적이지만, 현재 두 가지 근본적인 한계에 직면해 있습니다.

1. 비효율적인 자원 할당 (Resource Allocation):
   • 기존 방법 (예: GRPO) 은 학습 인스턴스 (문제) 간에 균일한 롤아웃 (rollout) 예산을 할당합니다.
   • 그러나 문제마다 기울기 (gradient) 정보의 이질성이 존재합니다. 고난이도 문제는 학습 신호가 풍부할 수 있지만 추정이 불안정할 수 있고, 너무 쉬운 문제는 학습 신호가 약합니다. 균일한 할당은 이러한 정보와 잡음 (noise) 간의 트레이드오프를 무시하여 학습 효율을 저하시킵니다.
2. 정책 최적화의 동적 불안정성 (Policy Optimization Dynamics):
   • 소프트맥스 (Softmax) 정책 구조의 한계: 높은 확률 (고신뢰도) 로 정답을 예측하는 행동은 기울기 크기가 자연스럽게 감소하여 학습 신호가 약해집니다 (Gradient Attenuation).
   • 학습 불안정: 반대로, 과도한 기울기 업데이트는 학습을 불안정하게 만듭니다.
   • 기존 방법들은 비율 클리핑 (ratio clipping) 이나 샘플 재가중치 등 거시적인 개입을 시도했으나, 토큰 수준의 기울기 감쇠와 엔트로피 변화 간의 이론적 관계를 체계적으로 해결하지 못했습니다.

2. 제안 방법론: DynaMO (Methodology)

저자들은 DynaMO(Dynamic Rollout Allocation and Advantage Modulation) 라는 이중 전략 (Dual-pronged) 최적화 프레임워크를 제안합니다. 이는 시퀀스 레벨과 토큰 레벨에서 각각 작동합니다.

A. 시퀀스 레벨: 동적 롤아웃 할당 (Dynamic Rollout Allocation)

• 이론적 기반: 전체 추정 분산을 최소화하는 관점에서 최적의 롤아웃 할당 문제를 수식화했습니다. 균일한 할당이 비최적 (suboptimal) 임을 증명하고, 분산을 최소화하는 할당 방정식을 유도했습니다.
• 베르누이 분산 (Bernoulli Variance) 프록시: 실제 기울기 분산을 계산하는 것은 비용이 많이 들기 때문에, 이진 보상 (성공/실패) 에 기반한 베르누이 분산을 계산 가능한 프록시로 사용합니다.
  • $P_i = \frac{k_i(G_i - k_i)}{G_i(G_i - 1)}$ (여기서 $k_i$는 정답 수, $G_i$는 총 롤아웃 수)
  • 성공 확률이 0.5 에 가까울 때 (즉, 정답과 오답이 균형을 이룰 때) 분산이 최대가 되며, 이는 모델이 가장 많이 학습할 수 있는 "학습 가능 구간 (capability gap)"을 의미합니다.
• 워터 레벨 할당 (Water-level Allocation): 알고리즘 1 에서와 같이, 각 문제의 베르누이 분산에 비례하여 예산을 동적으로 재분배하되, 최소 ($G_{min}$) 와 최대 ($G_{max}$) 롤아웃 수의 제약을 적용하여 과소/과대 샘플링을 방지합니다.

B. 토큰 레벨: 기울기 인식 이점 변조 (Gradient-Aware Advantage Modulation)

소프트맥스 정책의 구조적 한계를 해결하기 위해 보상 (Compensation) 과 안정화 (Stabilization) 메커니즘을 통합했습니다.

1. 기울기 보상 (Gradient Compensation):
   • 이론: 높은 신뢰도 (높은 엔트로피가 낮은 상태) 의 정답 행동은 기울기 크기가 작아 학습이 느려집니다.
   • 해결: 엔트로피에 반비례하는 보상 인자 ($\beta_{comp}$) 를 도입하여, 높은 확률의 정답 토큰에 대해 학습 신호를 증폭시킵니다.
   • 수식: $g(H) = 1 + \alpha \cdot \frac{H_{max} - H}{H_{max} - H_{min}}$
2. 업데이트 크기 안정화 (Update Magnitude Stabilization):
   • 이론: 기울기 크기와 정책 엔트로피 변화 사이에는 상한 관계가 존재합니다. 엔트로피의 급격한 변화는 과도한 기울기 업데이트를 의미하며 이는 학습 불안정을 초래합니다.
   • 해결: 토큰별 엔트로피 변화량 ($\Delta H$) 을 모니터링하여, 과도한 변화를 보이는 토큰에 대해 변조 인자 ($\beta_{stab}$) 를 통해 업데이트 크기를 억제합니다.
3. 통합 이점 변조: 최종 이점 (Advantage) 은 원래 이점에 보상과 안정화 인자를 곱하여 계산됩니다.
   • $A_{final} = A \cdot \beta_{comp} \cdot \beta_{stab}$

3. 주요 기여 (Key Contributions)

1. 균일 할당의 비최적성 증명 및 분산 최소화 할당 도출: 정책 기울기 방법론을 위해 베르누이 분산을 경량 프록시로 활용한 동적 할당 전략을 이론적으로 증명했습니다.
2. 기울기 - 엔트로피 관계 정립: 소프트맥스 정책에서 기울기 크기가 엔트로피에 의해 상한이 결정됨을 이론적으로 분석하고, 이를 기반으로 한 보상 및 안정화 메커니즘을 설계했습니다.
3. 실험적 검증: 다양한 수학 추론 벤치마크와 LLM 규모 (1.5B, 7B, 14B) 에서 기존 강력한 베이스라인 (GRPO, Entropy Loss 등) 을 일관되게 상회하는 성능을 입증했습니다.

4. 실험 결과 (Results)

• 벤치마크: AIME24/25, AMC23, MATH500, Minerva, Olympiad 등 6 가지 수학 추론 벤치마크에서 평가.
• 모델: Qwen2.5-Math (1.5B, 7B), Qwen3-14B.
• 성과:
  • Qwen2.5-Math-7B: 평균 Pass@1 에서 기존 GRPO(37.2%) 대비 **41.6%**로 향상되었으며, 다른 모든 베이스라인 (Clip-Higher, Entropy Advantages 등) 보다 우수했습니다.
  • 확장성 (Scaling): 모델 크기가 커질수록 (1.5B → 7B → 14B) DynaMO 와 GRPO 간의 성능 격차가 더욱 벌어져, 제안된 메커니즘이 대규모 모델에서 더 효과적으로 작동함을 보였습니다.
  • 학습 동역학: DynaMO 는 GRPO 에 비해 기울기 노름 (Gradient Norm) 의 급격한 스파이크와 엔트로피의 불안정한 변동을 억제하여 더 안정적이고 매끄러운 학습 곡선을 보였습니다.
• Ablation Study: 동적 할당 (DRA), 기울기 보상 (GC), 안정화 (UMS) 구성 요소 중 하나라도 제거하면 성능이 저하되었으며, 특히 여러 요소를 제거할 때 성능 하락이 가산적 이상으로 발생하여 구성 요소 간의 시너지 효과를 입증했습니다.

5. 의의 및 결론 (Significance)

• 이론적 엄밀성: RLVR 의 자원 할당과 정책 최적화 동역학에 대한 기존 경험적 접근을 넘어, 분산 최소화 이론과 기울기 - 엔트로피 관계에 기반한 엄밀한 이론적 토대를 제공했습니다.
• 효율성과 안정성: 계산 자원을 가장 학습에 효과적인 문제 (높은 분산, 균형된 성공/실패 비율) 에 집중시키고, 토큰 단위에서 학습 신호를 보정하여 모델의 수렴 속도와 안정성을 동시에 개선했습니다.
• 범용성: 특정 아키텍처에 의존하지 않고 알고리즘 레벨 (롤아웃 할당 및 이점 계산) 에서 작동하므로, 다양한 LLM 과 작업에 적용 가능한 잠재력을 가집니다.

요약하자면, DynaMO는 "어떻게 자원을 할당할 것인가 (Dynamic Allocation)"와 "어떻게 학습할 것인가 (Advantage Modulation)"라는 두 가지 핵심 질문에 대한 이론적으로 입증된 해답을 제시함으로써, LLM 의 추론 능력 향상을 위한 RLVR 의 새로운 표준을 제시합니다.