작고 일반화 가능한 프롬프트 예측 모델은 대형 추론 모델의 효율적인 강화학습 사후 훈련을 유도할 수 있다

Small Generalizable Prompt Predictive Models Can Steer Efficient RL Post-Training of Large Reasoning Models

강화학습은 대형 언어 모델의 추론 능력을 향상시키지만, 롤아웃 집약적인 최적화 과정으로 인해 높은 계산 비용을 수반하는 경우가 많습니다. 온라인 프롬프트 선택은 정보가 풍부한 프롬프트를 우선시하여 훈련 효율성을 높이는 타당한 해결책을 제시합니다. 그러나 기존의 방법들은 비용이 많이 드는 정확한 평가에 의존하거나, 프롬프트 간 일반화가 부족한 프롬프트별 예측 모델을 구축하는 데 그치고 있습니다. 본 연구에서는 공유된 최적화 이력을 바탕으로 훈련된 경량 생성 모델을 사용하여 프롬프트 난이도에 대한 베이지안 추론을 수행하는 '일반화 가능한 예측 프롬프트 선택(GPS)' 기법을 소개합니다. 정보가 풍부한 프롬프트 배치를 선택하기 위해 중간 난이도 우선순위 지정 및 이력 기반 다양성을 배치 획득 원칙에 통합하였습니다. 또한 이 소형 예측 모델은 테스트 시점에도 일반화되어 효율적인 계산 할당을 가능하게 합니다. 다양한 추론 벤치마크에 대한 실험 결과, GPS는 우수한 베이스라인 방법들에 비해 훈련 효율성, 최종 성능 및 테스트 시점 효율성 면에서 상당한 개선을 보이는 것으로 나타났습니다.

Reinforcement learning enhances the reasoning capabilities of large language models but often involves high computational costs due to rollout-intensive optimization. Online prompt selection presents a plausible solution by prioritizing informative prompts to improve training efficiency. However, current methods either depend on costly, exact evaluations or construct prompt-specific predictive models lacking generalization across prompts. This study introduces Generalizable Predictive Prompt Selection (GPS), which performs Bayesian inference towards prompt difficulty using a lightweight generative model trained on the shared optimization history. Intermediate-difficulty prioritization and history-anchored diversity are incorporated into the batch acquisition principle to select informative prompt batches. The small predictive model also generalizes at test-time for efficient computational allocation. Experiments across varied reasoning benchmarks indicate GPS's substantial improvements in training efficiency, final performance, and test-time efficiency over superior baseline methods.