기술 요약: RL 사후 학습을 위한 교차 에포크 적응형 롤아웃 최적화 (CERO)

1. 문제 정의

GRPO와 같은 알고리즘을 사용하는 대규모 언어 모델(LLM)의 RL 사후 학습(post-training)은 일반적으로 각 프롬프트당 고정된 수의 롤아웃을 샘플링하는 방식에 의존합니다. 이러한 균등 할당 전략은 서로 다른 프롬프트가 제공하는 학습 신호의 이질성을 무시합니다. 어떤 프롬프트는 사소하며(일관된 성공 또는 실패를 보이며 분산이 낮음), 어떤 프롬프트는 너무 어렵거나 이미 마스터되어 한계 효용이 감소합니다. 결과적으로, 전체 롤아웃 예산의 상당 부분이 학습 가치가 제한적인 프롬프트에 낭비됩니다.

본 논문은 이를 고정된 글로벌 예산 $B$ 하에서의 교차 에포크 온라인 자원 할당 문제로 정의합니다. 프롬프트 내에서 수행되는 이전의 배치 내 재할당이나 생성 후 필터링 방식과 달리, 본 연구는 프롬프트와 훈련 에포크 모두에 걸친 롤아웃 할당 문제를 다룹니다. 목표는 프롬프트 수준의 **한계 효용 체감(diminishing returns)**을 고려하여 수집된 샘플의 누적 유틸리티를 최대화하는 것입니다. 이 문제는 $\sum_{i,k} N_{i,k} \le B$라는 글로벌 예산 제약 조건 하에서 프로프트별 유틸리티 함수 $U_i(\sum_k N_{i,k})$의 합을 최대화하는 것으로 수학적으로 정의됩니다. 여기서 $N_{i,k}$는 에포크 $k$에서 프롬프트 $i$에 할당된 롤아웃 수입니다.

2. 방법론: CERO

저자들은 추정된 프롬프트 정보성(informativeness)에 따라 롤아웃을 적응적으로 할당하도록 설계된 베이지안 프라이멀-듀얼(Bayesian primal-dual) 온라인 알고리즘인 CERO(Cross-Epoch Adaptive Rollout Optimization)를 제안합니다.

2.1 정보성 및 유틸리티 설계

• 베타 사후 분포 모델링: 각 프롬프트 $x_i$에 대해, CERO는 롤아웃 이력 $H_i^k$를 바탕으로 잠재적 성공 확률 $p_i$에 대한 베타 사후 분포를 유지하고 업데이트합니다.
• 정보성 점수 ($q_{k,i}$): 단순히 성공 확률의 점 추정치를 사용하는 대신, CERO는 기대 베르누이 분산(expected Bernoulli variance)을 추정합니다: $q_{k,i} = E[p_i(1-p_i) | H_i^k]$. 이 점수는 $p_i \approx 0.5$(중간 난이도)일 때 최대화되며, 쉬운 프롬프트($p_i \approx 1$)나 어려운 프롬프트($p_i \approx 0$)에서는 0에 수렴합니다. 이는 가장 불확실하고 따라서 학습 정보가 많은 프롬프트를 효과적으로 식별합니다.
• 유틸리티 함수: 오목하고 포화되는 유틸리티 함수 $U_i(n) = 1 - \exp(-\eta q_i n)$이 정의됩니다 (여기서 $n$은 누적 롤아웃 수). 이는 한계 효용 체감을 포착합니다: 정보가 풍부한 프롬프트에 대한 초기 롤아웃은 매우 가치 있지만, 더 많은 샘플이 수집됨에 따라 한계 가치는 감소합니다.

2.2 Fenchel-Dual 재정식화를 통한 최적화

유틸리티 함수이 시간적으로 결합되어 있어(누적 할당량에 의존함) 직접적인 온라인 최적화가 어렵습니다. CERO는 Fenchel-dual 재정식화를 통해 이를 해결합니다:

• 문제는 프롬프트 수준의 듀얼 변수 $\theta_i$와 글로벌 예산 승수 $\mu$를 포함하는 듀얼 목적 함수로 변환됩니다.
• 이 재정식화는 시간적으로 결합된 문제를 듀얼 변수에 의해 조정되는 에포크별 할당 하위 문제로 분해합니다.

2.3 온라인 프라이멀-듀얼 알고리즘

CERO는 기존 RL 파이프라인(예: GRPO) 위에 드롭인(drop-in) 레이어로 작동합니다:

1. 할당 결정: 각 에포크에서 각 프롬프트에 대해 알고리즘은 선형 하위 문제를 풉니다. 추정된 한계 가치(듀얼 변수 $\theta_i$)가 현재 글로벌 예산 가격($\mu$)보다 크면 해당 프롬프트는 최대 허용 롤아웃($N_{max}$)을 받으며, 그렇지 않으면 0을 받습니다.
2. 피드백 업데이트: 응답을 생성하고 보상을 관찰한 후, 프롬프트에 대한 베타 사후 분포가 업데이트되고 정보성 점수 $q_{k,i}$가 갱신됩니다.
3. 듀얼 변수 업데이트:
   • 프롬프트 수준 ($\theta$): 할당된 롤아웃을 현재 정보성 추정치에 따른 유틸리티와 균형을 맞추기 위해 투영 온라인 경사 하강법(projected online gradient descent)을 통해 업데이트됩니다.
   • 예산 수준 ($\mu$): 실제 롤아웃 소비량과 목표 예산 비율 사이의 불일치를 기반으로 업데이트됩니다. 소비량이 비율을 초과하면 $\mu$가 증가하여 향후 할당을 더 선택적으로 만들고, 미달하면 $\mu$가 감소합니다.

3. 주요 기여

1. 문제 정의: 본 논문은 LLM RL 사후 학습의 롤아웃 할당을 프롬프트별 한계 효용 체감을 가진 교차 에포크 고정 예산 온라인 최적화 문제로 공식화하고, 공유된 글로벌 예산을 통해 결정 사항을 시간축 전체에 명시적으로 결합한 최초의 연구입니다.
2. 알고리즘 설계: 프롬프트 수준의 롤아웃 가치에 대한 사후 추정치를 사용하여 예산을 가이드하는 베이지안 프라이멀-듀얼 알고리즘인 CERO를 도입했습니다. 이는 학습 가능성의 대리 지표로 사후 기대 베르누이 분산을 활용하며, 시간적 비분리성을 처리하기 위해 Fenchel-dual 재정식화를 사용합니다.
3. 이론적 및 경험적 검증:
   • 이론: 고정된 프롬프트 유틸리티 하에서, 저자들은 오프라인 할당 벤치마크에 대해 $O(\sqrt{K})$ 리그렛 바운드(regret bound)를 증명했습니다 (여기서 $K$는 에포크 수).
   • 실험: 네 가지 오픈 웨이트 LLM(1.5B에서 7B 파라미터 범위)에 대해 수학적 추론 작업(AIME, AMC)에서 실시한 광범위한 실험을 통해 CERO가 바닐라 GRPO를 일관되게 능가함을 입증했습니다.

4. 실험 결과

저자들은 DeepSeek-R1-Distill-1.5B, Qwen3-4B (Base 및 Instruct), Qwen2.5-Math-7B를 사용하여 DAPO-Math-17K 데이터셋에 대해 CERO를 평가했습니다.

• 성능: CERO는 모든 모델과 벤치마크에서 바닐라 GRPO 대비 일관된 개선을 달이트했습니다. 모델에 따라 평균 이득은 +1.53에서 +6.04 포인트 사이였습니다.
• 효율성 분석: "유효 프롬프트 비율"(성공과 실패 응답을 모두 포함하여 0이 아닌 그래디언트를 생성하는 배치의 프롬프트 비율)이 GRPO보다 CERO에서 현저히 높았습니다.
  • 바닐라 GRPO에서는 훈련이 진행됨에 따라 유효 프롬프트 비율이 급격히 떨어졌는데, 이는 모델이 많은 프롬프트를 마스터하여 롤아웃 그룹이 퇴화되었음을 나타냅니다.
  • CERO는 예산을 마스터된 프롬프트에서 불확실성이 더 높은 프롬프트로 적응적으로 이동시킴으로써 더 높은 유효 프롬프트 비율을 유지했습니다.
• 오버헤드: 할당 업데이트 메커니즘은 무시할 만한 계산 오버헤드(업데이트당 1초 미만)를 도입합니다.

5. 의의 및 주장

본 논문은 적응형 롤아웃 예산 책정이 RL 사후 학습의 샘플 효율성을 개선하기 위한 원칙적이고 효과적인 방향임을 주장합니다. 균등한 할당에서 벗어남으로써, CERO는 고정된 생성 예산이 가장 정보가 풍부한 훈련 신호를 제공하는 프롬프트에 사용되도록 보장합니다.

저자들은 CERO를 기본 정책 그래디언트 목적 함수를 변경하지 않으면서 어떤 프롬프트가 롤아웃을 받고 얼마나 많이 받을지를 수정하는 "드롭인(drop-in)" 적응형 데이터 수집 레이어로 포지셔닝합니다. 이 연구는 현재 RL 사후 학습의 비효율성이 단순히 정책 업데이트 규칙 때문이 아니라, 최적화되지 않은 데이터 수집 전략에서 비롯된다는 점을 강조합니다.

저자들이 언급한 한계점:

• 이론적 리그렛 보장은 고정된 프롬프트 유틸리티를 가정하지만, 구현에서는 동적인 에포크별 플러그인 사후 추정치를 사용합니다. 따라서 분석이 정책 업데이트 및 사후 분포 갱신으로 인해 발생하는 비정상성(non-stationarity)을 완전히 포착하지는 못합니다.
• 경험적 검증은 현재 수학적 추론 작업에 국한되어 있으며, 다른 도메인 및 보상 모델에 대한 광범한 검증이 향후 과제로 제시되었습니다.