Adaptive Rollout Allocation for Online Reinforcement Learning with Verifiable Rewards

이 논문은 고정된 롤아웃 할당의 비효율성을 해결하기 위해 경량 가우시안 프로세스 모델을 통해 각 프롬프트의 성공 확률을 예측하고 이를 기반으로 분산을 최소화하는 볼록 최적화를 수행하여 롤아웃을 동적으로 할당하는 'VIP' 전략을 제안하여 온라인 강화학습의 샘플링 효율성을 크게 향상시킵니다.

핵심

1. 문제 상황: 모든 학생에게 똑같은 양의 문제지를 주는 비효율

지금까지 AI 를 가르치는 방법 (GRPO 같은 기존 기술) 은 마치 수학 선생님이 100 명의 학생에게 모두 똑같이 16 문제씩 풀게 하는 것과 같습니다.

• 쉬운 문제: 학생이 이미 100% 정답을 아는 문제입니다. 16 번이나 풀어도 새로운 배움이 없습니다. (시간 낭비)
• 너무 어려운 문제: 학생이 100% 틀리는 문제입니다. 16 번이나 풀어도 답을 못 냅니다. (시간 낭비)
• 적당한 문제: 학생이 50% 는 맞고 50% 는 틀리는 문제입니다. 이 문제들을 풀 때 가장 많은 배움이 일어납니다.

기존 방식은 이 '쉬운 문제'와 '어려운 문제'에도 똑같은 시간과 노력 (컴퓨터 자원) 을 써버려서, 정작 중요한 '적당한 문제'를 가르칠 시간이 부족해지는 문제가 있었습니다.

2. VIP 의 해결책: 똑똑한 자원 배분 시스템

이 논문에서 제안한 VIP(Variance-Informed Predictive allocation) 는 **"각 학생의 실력을 미리 예측해서, 필요한 만큼만 문제를 내주는 똑똑한 조교"**입니다.

1 단계: 학생의 실력을 '예측'하다 (가우시안 프로세스)

조교는 학생들의 과거 풀이 기록을 보고, "아, 이 학생은 A 문제는 90% 확률로 풀고, B 문제는 50% 확률로 풀겠구나"라고 예측합니다.

• 여기서는 복잡한 수학 공식 대신, **구름 모양의 지도 (가우시안 프로세스)**를 그려서 학생들의 실력 변화를 부드럽게 예측합니다.

2 단계: '현금'을 가장 필요한 곳에 '투자'하다 (최적화)

조교는 총 100 개의 문제지 (컴퓨터 자원) 만 가지고 있습니다. 이걸 어떻게 나누어야 할까요? VIP 는 다음과 같이 분배합니다.

• 이미 다 아는 학생 (쉬운 문제): 문제지를 2~3 장만 줍니다. (이미 답을 알기 때문에 더 풀 필요 없음)
• 아예 모르는 학생 (너무 어려운 문제): 문제지를 3~4 장만 줍니다. (아무리 풀어도 안 되니 너무 많은 시간을 쓸 필요 없음)
• 아슬아슬한 학생 (적당한 문제): 문제지를 20~30 장이나 줍니다. 이곳에서 가장 큰 성장 (학습 효과) 이 일어나기 때문입니다.

이 과정을 수학적으로 계산하여, 전체적인 학습의 '흔들림 (분산)'을 가장 작게 만드는 최적의 배분을 찾아냅니다.

3. 왜 이것이 중요한가요? (결과)

이 방식을 적용한 실험 결과, 같은 양의 컴퓨터 자원 (시간과 돈) 으로 훨씬 더 똑똑한 AI를 만들 수 있었습니다.

• 기존 방식: 100 달러를 100 명에게 1 달러씩 나눠줌. (누구도 제대로 배움)
• VIP 방식: 100 달러를 배울 의욕이 있고, 배울 수 있는 10 명에게 10 달러씩 집중 투자함. (그 10 명이 크게 성장하여 전체 평균이 급상승)

4. 요약: 한 줄로 정리하면?

"AI 를 가르칠 때, 모든 문제에 똑같은 시간을 쓰지 말고, AI 가 '가장 헷갈려하는 문제'에 집중해서 시간을 투자하면, 더 빠르고 똑똑하게 성장할 수 있다."

이 기술은 수학 문제 풀이뿐만 아니라, AI 가 검색 도구를 사용하거나 복잡한 작업을 수행할 때도 적용되어, 더 적은 비용으로 더 높은 성능을 내는 AI 시대를 여는 중요한 열쇠가 될 것입니다.

기술 요약

VIP: 가변적 확률 분포를 기반으로 한 검증 가능한 보상 강화학습을 위한 적응형 롤아웃 할당 기술 요약

이 논문은 **검증 가능한 보상 (Verifiable Rewards, RLVR)**을 사용하는 강화학습 (RL) 에서 발생하는 **샘플링 효율성 (Sampling Efficiency)**의 병목 현상을 해결하기 위해 제안된 **VIP (Variance-Informed Predictive allocation)**라는 새로운 프레임워크를 다룹니다.

1. 문제 정의 (Problem)

대규모 언어 모델 (LLM) 의 후학습 (Post-training) 과정에서 RLVR (예: 수학 문제 풀이, 도구 활용) 은 인간 피드백 (RLHF) 에 비해 외부 검증 로직을 통해 보상을 얻기 때문에 선호되지만, 계산 비용이 매우 높습니다.

• 기존 방법의 한계: GRPO(Group Relative Policy Optimization) 와 같은 그룹 기반 정책 최적화 알고리즘은 모든 학습 프롬프트에 대해 **고정된 수의 롤아웃 (rollout, 생성 횟수)**을 할당합니다.
• 비효율성: 모든 프롬프트가 동일한 정보를 제공하는 것은 아닙니다. 모델이 이미 정답을 잘 아는 문제 (성공 확률 $p \approx 1$) 나 전혀 풀지 못하는 문제 ($p \approx 0$) 에 많은 롤아웃을 할당하는 것은 계산 자원의 낭비이며, 학습 진행을 방해할 수 있습니다.
• 핵심 과제: 제한된 계산 예산 (Budget) 내에서 각 프롬프트에 얼마나 많은 롤아웃을 할당해야 **정책 업데이트의 기댓값 기울기 분산 (Expected Gradient Variance)**을 최소화하여 학습 효율을 극대화할 수 있을까요?

2. 방법론 (Methodology)

VIP 는 두 가지 핵심 단계로 구성됩니다: 가우시안 프로세스 (GP) 기반 예측과 볼록 최적화 기반 할당.

2.1. 기울기 분산 분석 (Gradient Variance Analysis)

논문은 GRPO 와 RLOO(Reinforcement Learning with Online Optimization) 와 같은 알고리즘에서 프롬프트별 기울기 분산이 **해당 프롬프트의 성공 확률 ($p$)**과 **롤아웃 수 ($n$)**에 어떻게 의존하는지 이론적으로 분석했습니다.

• 핵심 발견: 기울기 분산은 $p(1-p)$에 비례합니다. 즉, 성공 확률이 0.5 에 가까울 때 분산이 최대가 되며, 0 또는 1 에 가까울 때 분산이 줄어듭니다.
• 의미: 모델이 "어느 정도 풀 수 있지만 확실하지 않은" (중간 난이도) 프롬프트에 더 많은 샘플링 자원을 할당해야 전체 학습의 분산을 줄일 수 있습니다.

2.2. 예측 단계: 가우시안 프로세스 (Gaussian Process, GP)

각 학습 단계에서 모델의 가중치가 변함에 따라 프롬프트별 성공 확률도 동적으로 변합니다. 이를 예측하기 위해 VIP 는 다음과 같이 작동합니다:

• GP 모델: 프롬프트 임베딩 (Embedding) 을 입력으로 받아 현재 모델의 성공 확률을 예측하는 비모수적 GP 모델을 사용합니다.
• 재귀적 업데이트: 이전 단계의 롤아웃 결과 (성공/실패) 를 관측하여 GP 의 사후 분포 (Posterior) 를 업데이트하고, 다음 단계의 예측에 활용합니다. 이는 모델의 능력 변화에 적응적으로 대응할 수 있게 합니다.

2.3. 할당 단계: 볼록 최적화 (Convex Optimization)

예측된 성공 확률 ($\hat{p}_q$) 을 기반으로 각 프롬프트 $q$에 할당할 롤아웃 수 $n_q$를 결정합니다.

• 목적 함수: 미니배치 내 모든 프롬프트의 기울기 분산 합을 최소화.
• 제약 조건: 총 롤아웃 수 ($C$) 는 고정되어 있으며, 각 프롬프트당 최소 ($L$) 와 최대 ($U$) 롤아웃 수가 존재합니다.
• 해법:
  1. 정수 제약 조건을 완화한 연속 최적화 문제를 풀어서 최적의 $n_q$를 구합니다 (라그랑주 승수법 및 이분법 사용).
  2. 구해진 실수 해를 정수로 변환하기 위해 **그리디 기반 반올림 휴리스틱 (Greedy Rounding Heuristic)**을 적용하여 실제 실행 가능한 할당량을 도출합니다.

3. 주요 기여 (Key Contributions)

1. 이론적 분석: 그룹 기반 RL 알고리즘 (GRPO, RLOO) 에서 기울기 분산과 성공 확률 간의 정량적 관계를 수학적으로 증명했습니다.
2. 적응형 예측 프레임워크: GP 를 활용하여 동적인 학습 환경에서 프롬프트별 성공 확률과 분산을 실시간으로 추정하는 방법을 제시했습니다.
3. 최적 할당 알고리즘: 예측된 분산을 기반으로 계산 예산을 최소화하는 볼록 최적화 문제를 설계하고, 이를 효율적으로 해결하는 알고리즘을 개발했습니다.
4. 실증적 검증: 수학 추론 및 도구 활용 작업에서 기존 균일 할당 및 휴리스틱 방법 대비 우수한 성능을 입증했습니다.

4. 실험 결과 (Results)

논문은 **수학적 추론 (Mathematical Reasoning)**과 도구 활용 추론 (Tool-Augmented Reasoning) 두 가지 벤치마크에서 VIP 를 평가했습니다.

• 데이터셋 및 모델: DAPO-MATH-17k, AIME2024/2025, Bamboogle, MuSiQue 등. Qwen2.5-Math (1.5B, 7B), Llama-3.2-3B 등 다양한 모델 사용.
• 성능 향상:
  • 수학 추론: VIP 를 적용한 RLOO 및 Dr. GRPO 는 균일 할당 (Uniform) 대비 Pass@32 및 Mean@32 점수에서 일관된 개선을 보였습니다. 특히 1.5B 와 3B 와 같은 상대적으로 작은 모델에서 성능 향상이 두드러졌습니다 (예: Qwen2.5-Math-1.5B 에서 Pass@32 가 12.3%p 향상).
  • 도구 활용: Bamboogle 벤치마크에서 정답 정확도 (EM) 와 검색 품질 (F1@5, Precision@5) 이 동시에 향상되었습니다.
• 효율성: VIP 는 전체 RL 학습 시간에 **1% 미만 (약 0.8~1.1%)**의 추가 오버헤드만 발생시키며, 계산 자원을 효율적으로 재분배함으로써 더 높은 성능을 달성했습니다.
• Ablation Study: GP 예측기나 적응형 할당기를 제거한 변형 모델들은 VIP 전체 모델보다 성능이 현저히 낮아, 두 구성 요소 모두 필수적임을 확인했습니다.

5. 의의 및 결론 (Significance)

• 자원 효율성: 제한된 컴퓨팅 자원을 가진 환경에서 LLM 의 RL 학습 효율을 극대화하는 새로운 패러다임을 제시합니다.
• 적응형 학습: 모델의 학습 진행도에 따라 "어떤 문제가 중요한지"를 동적으로 파악하여 자원을 집중하는 지능형 학습 파이프라인을 구현했습니다.
• 확장성: 검증 가능한 보상 (RLVR) 환경에서 효과적이지만, 향후 인간 피드백 (RLHF) 이나 노이즈가 있는 보상 환경으로도 확장 가능한 잠재력을 가집니다.

결론적으로, VIP 는 고정된 샘플링 전략의 비효율성을 해결하고, 분산 (Variance) 을 고려한 예측적 할당을 통해 RL 기반 LLM 학습의 속도와 성능을 동시에 개선하는 획기적인 방법론입니다.