Train at Moving Edge: Online-Verified Prompt Selection for Efficient RL Training of Large Reasoning Model

Summary

이 논문은 GRPO 같은 RL 기반 LLM 추론 학습에서 많은 프롬프트가 실제로는 거의 학습 신호를 주지 못한다는 문제를 다룬다. 저자들은 프롬프트의 유틸리티가 균일하지 않으며, 가장 강한 학습 신호는 너무 쉽지도 어렵지도 않은 중간 난이도와 높은 불확실성이 만나는 “learning edge”에 집중된다고 분석한다. 또한 이 경계는 학습이 진행되며 계속 이동하므로, 과거 통계만으로 프롬프트를 고르면 빠르게 stale해진다는 점을 보인다.

이를 위해 제안된 HIVE는 2단계 프롬프트 선택 프레임워크다. 1단계에서는 과거 reward trajectory와 response entropy를 사용해 저비용으로 후보를 거칠게 추리고, 2단계에서는 현재 모델로 계산한 prompt entropy를 이용해 실제로 지금도 유용한 샘플만 온라인으로 다시 검증한다. 핵심은 비싼 롤아웃 전에 단일 forward pass 수준의 가벼운 검증을 넣어, stale한 샘플과 zero-gradient 프롬프트를 제거하는 데 있다.

실험은 여러 수학 추론 벤치마크와 다양한 Qwen/DeepSeek/Llama 계열 모델에서 수행되었으며, HIVE는 Dynamic Sampling(DS)과 GRESO 대비 비슷하거나 더 나은 정확도를 유지하면서 롤아웃 수와 총 학습 시간을 크게 줄였다. 특히 Qwen2.5-Math-7B에서는 DS 대비 최대 3.4× 적은 롤아웃으로 더 높은 평균 정확도를 달성했고, 전체적으로 최대 3.8× 롤아웃 속도 향상과 2.2× 수준의 총 학습 시간 단축을 보였다.

Problem

RL 기반 LLM 후처리 학습에서는 프롬프트마다 여러 개의 롤아웃을 생성해야 하므로 계산 비용이 매우 크다. 하지만 실제로는 너무 쉬운 문제나 현재 모델로는 풀 수 없는 너무 어려운 문제는 GRPO에서 reward 분산이 0에 가까워져 gradient가 사라지고, 학습에 거의 기여하지 않는다. 기존 방법은 정적 휴리스틱에 의존하거나, 현재 유용성을 판단하기 위해 추가 롤아웃을 수행해 오히려 비용을 키우거나, 과거 학습 기록만 사용해 시간이 지나며 stale해지는 문제가 있었다. 이 논문은 “현재 시점의 학습 경계에 있는, 가장 학습 효율이 높은 프롬프트를 롤아웃 전에 어떻게 싸게 고를 것인가?”를 해결하려 한다.

Method

HIVE는 두 단계로 프롬프트를 선택한다. 1단계 History-Informed Selection에서는 각 프롬프트의 과거 reward trajectory에서 연속 zero-variance 횟수를 세어, 오랫동안 정보가 없었던 프롬프트의 선택 확률을 낮춘다. 동시에 직전 epoch의 response entropy를 이용해 불확실성이 높은 프롬프트를 우선시하고, 두 점수를 선형 결합해 후보 집합을 만든다. 이때 exploration probability를 easy/hard 샘플에 대해 적응적으로 조절해 학습 경계가 이동해도 완전히 버려지지 않게 한다.

2단계 Online-Verified Selection에서는 후보 프롬프트에 대해 현재 모델의 prompt entropy를 계산하고, 현재 후보 집합의 median 이상인 샘플만 최종적으로 롤아웃에 넘긴다. 저자들은 prompt entropy가 response entropy의 순위를 높은 확률로 보존한다는 rank consistency를 이론적으로 제시해, 이 값이 현재 유틸리티의 효율적인 proxy가 될 수 있음을 주장한다. 결과적으로 HIVE는 “과거 정보로 넓게 추리고, 현재 엔트로피로 정확히 검증하는” soft-to-hard 선택 구조를 통해 rollout 비용을 줄인다.

Results

DAPO+MATH에서 Qwen2.5-Math-7B 기준, HIVE는 평균 정확도 59.7%로 DS의 57.8%, GRESO의 58.6%를 넘으면서 롤아웃 수를 13.1M에서 3.9M으로 줄여 DS 대비 3.4× 효율을 보였다. 같은 설정에서 총 학습 시간은 DS 198.4시간, GRESO 112.3시간 대비 HIVE 85.8시간으로 감소했고, rollout 시간은 153.7시간에서 40.2시간으로 줄어 최대 3.8× rollout speedup을 달성했다. Qwen2.5-Math-1.5B, DeepSeek-R1-Distill-Qwen-1.5B, Llama3.2-3B-Instruct에서도 대체로 비슷하거나 더 나은 평균 성능을 유지하면서 1.6×~2.5× 수준의 rollout 절감 효과를 보였다.

OpenR1 학습에서도 비슷한 경향이 나타났다. 예를 들어 Qwen2.5-Math-7B에서 HIVE는 평균 55.3%로 DS 55.0%, GRESO 55.1%와 유사하거나 소폭 우세하면서 롤아웃 수를 11.4M에서 2.5M으로 줄여 4.7× 효율을 기록했다. 추가 분석에서는 Stage 2의 온라인 검증 비용이 step당 0.82초로 전체 시간의 0.4% 미만에 불과했고, zero-variance 프롬프트 비율도 학습 전반에 걸쳐 가장 낮게 유지되어 실제로 더 많은 계산이 유효한 샘플에 집중됨을 보여준다.

Key Insights

RL에서 중요한 것은 롤아웃 수 자체가 아니라, 현재 모델에 대해 non-zero gradient를 주는 프롬프트 비율이라는 점을 잘 보여준다.
“중간 난이도 + 높은 엔트로피”라는 learning edge 개념은 curriculum learning과 uncertainty sampling을 RL post-training 맥락에서 자연스럽게 연결한다.
과거 메타데이터는 공짜에 가깝지만 금방 stale해지고, 온라인 지표는 정확하지만 비싸다는 trade-off를 2단계 구조로 균형 있게 풀었다.
response rollout 없이 prompt entropy만으로 현재 유틸리티를 근사한 점이 실용적이며, 실제 시스템 관점에서 추가 오버헤드가 거의 없다.
추론 모델 학습에서 데이터 선택이 단순한 데이터 정제가 아니라, 학습 중 계속 움직이는 target을 추적하는 online control 문제라는 관점을 제시한다.

Limitations

논문에서도 인정하듯 현재 평가는 텍스트 기반 대형 추론 모델과 수학 문제 중심 벤치마크에 제한되어 있어, 멀티모달 모델이나 더 일반적인 에이전트/도구 사용 환경으로의 일반화는 검증되지 않았다. 또한 HIVE는 λ, Δp 같은 추가 하이퍼파라미터를 도입하며, 적응 메커니즘이 있더라도 모든 모델 크기와 학습 스케일에서 최적인지는 확인되지 않았다. 더불어 prompt entropy가 response-side utility를 잘 대변한다는 가정은 수학 추론처럼 구조가 비교적 명확한 도메인에서 특히 잘 맞을 가능성이 있어, 자유형 생성 과제에서는 효과가 달라질 수 있다.