프론티어 중심 힙 샘플링과 온정책 쿼리 확장으로 효율적인 LLM 강화학습

초록

HeaPA는 제한된 크기의 프롬프트 풀을 유지하면서 힙 기반 경계 샘플링으로 학습 가능한 난이도 구역에 집중하고, 정책이 직접 생성한 새로운 쿼리를 비동기 검증을 통해 풀에 추가한다. 이를 통해 RL‑VR 훈련 시 롤아웃 비용을 크게 절감하면서도 정확도를 향상시킨다.

상세 분석

HeaPA는 LLM 강화학습(RL‑VR)에서 가장 비용이 많이 드는 롤아웃 생성 단계의 효율성을 극대화하기 위해 세 가지 핵심 메커니즘을 결합한다. 첫 번째는 이중 힙 구조를 이용한 경계 샘플링이다. 풀에 저장된 각 프롬프트는 평균 보상 ˜r (−1~1)으로 요약되며, 이 값을 기준으로 ‘hard’와 ‘easy’ 두 힙으로 나뉜다. 힙의 경계에 위치한 중간 난이도(˜r ≈ 0) 아이템을 우선적으로 선택함으로써, 너무 쉬워서 학습 신호가 부족하거나 너무 어려워서 거의 모두 실패하는 프롬프트에 대한 불필요한 롤아웃을 방지한다. 또한 Cold Queue를 도입해 아직 검증되지 않은 시드 프롬프트를 빠르게 초기화하고, 초기 단계에서 풀 전체에 대한 난이도 분포를 빠르게 구축한다.

두 번째는 온정책 쿼리 증강이다. 현재 정책 πθ가 직접 새로운 프롬프트를 생성하고, 이를 경량 검증기(또는 교사 모델)에게 비동기적으로 보내 정답을 얻는다. 이 과정은 별도의 교사 로스 계산 없이 정답 레이블만 확보하므로 교사 비용과 지연을 최소화한다. 생성된 쿼리는 즉시 풀에 삽입되지만, 초기에는 ˜r이 정의되지 않은 상태로 Cold Queue에 배치되어 빠른 검증 후 평균 보상이 업데이트된다.

세 번째는 위계 인식(pool‑statistic) 재추정 및 제어된 재삽입이다. 증강된 쿼리는 종종 동일 템플릿이나 작은 변형을 공유하므로, 단순히 풀에 삽입하면 샘플링 분포가 급격히 변동할 위험이 있다. HeaPA는 증강 그래프 Gₜ 를 유지해 각 쿼리의 혈통(lineage)을 추적하고, 혈통 전체에 걸쳐 ˜r을 재계산한다. 이렇게 하면 상관관계가 높은 아이템이 동시에 풀에 과다하게 반영되는 현상을 완화한다. 또한, 학습이 끝난 레코드를 Archive에 보관했다가 일정 시점에 Controlled Reinsertion 전략으로 다시 풀에 넣어, 풀 크기가 고정된 상황에서도 다양성을 유지한다.

HeaPA는 기존의 GRPO, DAPO와 같은 그룹 기반 RL‑VR 레시피와 optimizer‑agnostic하게 결합될 수 있다. 실험에서는 Qwen2.5‑7B 모델을 두 개의 수학 데이터셋(DAPO‑Math, OpenR1‑Math)에 적용했으며, 7개의 벤치마크에서 기존 파이프라인 대비 동일 목표 정확도에 도달하는 데 필요한 롤아웃 토큰 수를 평균 20‑30% 절감했다. 특히 모델 규모가 커질수록(7B→13B) 경계 샘플링의 효과가 두드러져, 대형 모델에서는 효율성이 더욱 확대된다. 프로파일링 결과, 힙 연산·아카이브 관리·비동기 검증이 전체 실행 시간에 차지하는 비중은 5% 이하로, 실제 학습 속도에 거의 영향을 주지 않는다.

핵심 인사이트는 프론티어‑중심 샘플링이 학습 효율을 크게 좌우한다는 점과, 온정책 데이터 증강이 정답 레이블을 자동으로 확보하면서도 현재 정책의 능력에 맞는 난이도 분포를 유지한다는 점이다. 이러한 설계는 정적 풀에 의존하던 기존 방법들의 한계를 넘어, 동적으로 진화하는 프롬프트 풀을 안정적으로 운영할 수 있게 만든다.