맥락 기반 롤아웃 밴딧으로 강화학습 검증 보상 최적화

초록

본 논문은 검증 가능한 보상(RLVR) 환경에서 롤아웃 데이터를 효율적으로 활용하기 위해 롤아웃 스케줄링을 컨텍스트 밴딧 문제로 정의하고, 이를 해결하는 신경망 기반 스케줄러(CBS)를 제안한다. CBS는 동일 그룹 내의 품질 이질성을 고려한 노이즈‑aware 선택과, 과거 롤아웃을 재사용하는 글로벌 적응 메커니즘을 하나의 프레임워크에 통합한다. 이론적으로 서브리니어 레그레트 경계를 증명하고, 버퍼 크기 확대가 성능 상한을 높임을 보인다. 실험에서는 6개의 수학 추론 벤치마크와 세 가지 RLVR 최적화 방법에 대해 일관된 성능 향상과 학습 효율 개선을 확인하였다.

상세 분석

본 연구는 기존 RLVR 방식이 롤아웃을 무차별적으로 사용하고, 최신 배치만을 한 번씩 소비한다는 두 가지 근본적인 한계를 지적한다. 첫 번째는 동일 프롬프트에 대해 생성된 다중 응답이 품질 차이를 보이지만, 모두 동일한 가중치로 정책 업데이트에 사용된다는 점이다. 이는 노이즈가 섞인 라벨을 제공해 정책을 왜곡하고, 불필요한 연산 비용을 초래한다. 두 번째는 과거에 생성된 고품질 롤아웃이 재활용되지 않아 샘플 효율이 저하된다는 문제다. 이를 해결하기 위해 저자는 롤아웃 선택을 ‘컨텍스트 밴딧’으로 모델링한다. 각 롤아웃을 하나의 팔(arm)로 보고, 해당 팔의 보상은 현재 정책 파라미터 θ_{t-1}에서 θ_t 로 업데이트될 때 얻는 성능 향상으로 정의한다. 이렇게 하면 롤아웃 선택이 즉각적인 정책 개선에 직접 연결된다.

CBS 스케줄러는 두 단계로 구성된다. (1) 롤아웃 인코더 f(·)를 통해 10차원 컨텍스트 벡터를 생성하고, (2) 다층 퍼셉트론(MLP) 기반 보상 예측기 s_ϕ가 각 벡터에 대해 미래 유틸리티를 추정한다. 예측된 보상이 높은 롤아웃을 상위 K개 선택함으로써 그룹 내 노이즈를 억제한다. 선택된 롤아웃 집합에 대해 전체 그룹 보상 R(·)을 계산하고, 이를 각 롤아웃의 어드밴티지 |A_i|에 비례해 샘플 수준 보상 R(·)으로 분배한다. 이렇게 얻은 샘플 보상은 스케줄러 파라미터 ϕ를 온라인 SGD 방식으로 업데이트하는데, 손실 함수는 예측 보상과 실제 보상의 제곱 차이이다. 이 과정은 매 학습 라운드마다 수행되며, 스케줄러는 데이터 분포와 정책 변화에 적응한다.

이론적 분석에서는 단일 롤아웃 선택 상황을 가정하고, 컨텍스트 밴딧의 표준 가정(선형 보상 모델, 서브가우시안 노이즈 등) 하에 레그레트가 O(√T)로 수렴함을 증명한다. 또한, 롤아웃 버퍼 크기 L을 늘릴 경우, 선택 가능한 최적 팔의 집합이 확대되어 정책이 도달할 수 있는 성능 상한이 증가한다는 정리를 제시한다. 이는 기존 방법이 최신 배치만 사용함으로써 얻을 수 있는 최적점보다 더 높은 성능을 이론적으로 보장한다는 의미다.

실험에서는 MATH, GSM8K, MMLU 등 6개의 수학·논리 추론 벤치마크와 GRPO, DAPO, GSPO 세 가지 RLVR 최적화 알고리즘을 조합했다. CBS를 적용한 경우 평균 2~4%의 정확도 상승과, 동일 에포크 수 대비 30% 이하의 연산량 감소를 기록했다. 특히, 초기 학습 단계에서 고품질 롤아웃을 빠르게 식별해 정책이 급격히 향상되는 현상이 관찰되었으며, 버퍼 재사용을 통해 후반 학습에서도 꾸준한 성능 개선이 유지되었다. Ablation 연구에서는 (i) 인트라그룹 선택만 적용했을 때와 (ii) 글로벌 재사용만 적용했을 때 각각의 효과를 분리 분석했으며, 두 기능을 결합했을 때 시너지 효과가 가장 크게 나타났다. 한계점으로는 현재 10차원 컨텍스트가 제한적일 수 있어, 더 풍부한 정책‑특화 메트릭을 포함한 고차원 표현이 향후 연구 과제로 남는다.