Prune as You Generate: Online Rollout Pruning for Faster and Better RLVR

Summary

이 논문은 RLVR(Reinforcement Learning with Verifiable Rewards)에서 널리 쓰이는 GRPO, DAPO 같은 방법이 프롬프트마다 많은 rollout을 생성해야 하므로 계산량이 매우 크고, 동시에 그룹 내 보상이 거의 전부 정답이거나 전부 오답이 되는 경우가 많아 학습 신호가 약해진다는 문제를 다룬다. 저자들은 이를 해결하기 위해 ARRoL(Accelerating RLVR via online Rollout Pruning)이라는 온라인 rollout pruning 기법을 제안한다. 핵심은 생성이 진행되는 도중 부분 rollout의 최종 성공 확률을 가볍게 예측하고, 그 예측값을 바탕으로 일부 rollout을 조기에 중단시켜 계산을 줄이면서도 남는 샘플들의 정오답 비율을 더 균형 있게 만드는 것이다.

방법적으로는 백본 LLM 위에 2-layer MLP 형태의 quality head를 붙여 부분 추론 경로의 품질 점수를 예측하고, 이를 온라인 calibration으로 성공 확률로 변환한 뒤 목표 keep ratio와 목표 정답 비율에 맞게 생존 확률을 정해 pruning한다. 또한 pruning을 inference engine 내부에서 수행하고, 살아남은 rollout만 다시 묶어 log-prob 계산과 policy update를 수행하는 시스템 설계를 함께 제시한다. 이로써 단순히 업데이트 단계만 줄이는 것이 아니라 실제 생성 단계의 wall-clock 시간까지 줄인다.

실험에서는 Qwen-3 및 LLaMA-3.2 계열 1B~~8B 모델에서 GRPO와 DAPO 모두에 대해 일관된 개선을 보였다. 평균 정확도는 GRPO 기준 +2.30~~+2.87, DAPO 기준 +2.99 향상되었고, 학습 속도는 최대 1.7배 빨라졌다. 더 나아가 학습된 quality head를 테스트 시 다수 샘플의 가중 투표에 활용해 test-time scaling에서도 최대 +8.33의 추가 정확도 향상을 얻었다.

Problem

RLVR에서는 각 프롬프트에 대해 여러 개의 rollout을 생성하고 그룹 상대적 보상으로 policy를 업데이트하는데, 이 과정은 생성 비용과 후속 log-prob/역전파 비용이 모두 크다. 게다가 보상이 0/1의 검증 가능한 형태인 경우 한 그룹 안의 샘플들이 거의 모두 맞거나 모두 틀리는 현상이 자주 발생해 reward variance가 낮아지고, GRPO류 알고리즘의 advantage가 거의 0이 되어 학습 신호가 약해진다. 기존의 효율화 방법 중 상당수는 생성이 끝난 뒤 rollout을 걸러내기 때문에 실제 generation cost를 줄이지 못하거나, 과거 히스토리에 의존해 cold-start 환경에서 제한적이라는 문제가 있다.

Method

ARRoL은 생성 중간 길이 (L_{detect}=512) 에 도달한 부분 rollout을 quality head로 평가해 조기 pruning하는 방식이다. quality head는 백본의 hidden representation을 입력받아 해당 rollout의 최종 정답 확률을 예측하며, 최종 reward로부터 자연스럽게 supervision을 얻어 online으로 학습된다. 저자들은 단순한 trace confidence보다 학습 가능한 quality head가 정답 여부와 더 잘 상관된다고 보인다.

예측된 점수는 online binned probability estimator로 보정되어 성공 확률 (q_i)로 변환된다. 이후 목표 keep ratio (\kappa)와 목표 positive ratio (\rho) (실험에서는 둘 다 0.5)에 맞게 각 rollout의 survival probability를 정한다. 직관적으로는 현재 그룹이 너무 “정답 쪽” 또는 “오답 쪽”으로 치우치지 않도록 일부 rollout을 제거해 남는 샘플들의 정오답 비율을 더 균형 있게 만든다. 시스템 측면에서는 vLLM 백엔드 내부에서 pruning을 즉시 반영해 request pool에서 제거하고, 프런트엔드에서는 생존 rollout만 재배치해 log-prob 계산과 policy update를 수행한다. 학습된 quality head는 추론 시에도 각 후보 답변의 가중 투표 점수로 재사용된다.

Results

GRPO 설정에서 ARRoL은 Qwen-3 1.7B/4B/8B와 LLaMA-3.2-1B 전반에 걸쳐 평균 정확도를 각각 34.79→37.09, 51.01→53.54, 56.66→59.53, 14.63→17.49로 향상시켰고, 속도는 1.61x~1.67x 빨라졌다. 특히 어려운 벤치마크에서 개선폭이 컸는데, 예를 들어 Qwen-3-8B에서 AIME’24는 56.67→66.67, AIME’25는 40.00→46.67이었다. DAPO에서도 Qwen-3-1.7B 기준 평균 정확도가 36.43→39.42로 +2.99 상승했고, 1.70x speedup을 달성했다.

효율 분해 결과 Qwen-3-1.7B 기준 generation 시간은 106.82초→72.96초(1.46x), log-prob 계산은 18.40초→10.02초(1.84x), update는 63.05초→30.26초(2.08x)로 줄었다. 무작위 pruning과 비교했을 때도 ARRoL은 더 높은 정확도와 더 균형 잡힌 그룹 보상 분포를 보였다. 또한 테스트 시 quality head 기반 가중 투표는 majority vote와 DeepConf를 모두 능가했으며, 예를 들어 Qwen-3-8B에서 AMC’23/AIME’24/AIME’25는 75.0/23.3/26.7(majority), 80.0/23.3/23.3(DeepConf), 85.0/33.3/33.3(ARRoL)로 향상되었다.

Key Insights

RLVR에서 “더 많은 rollout”이 항상 좋은 것이 아니라, 보상 분산이 살아 있는 rollout 집합이 더 중요한 학습 신호를 줄 수 있다는 점이 핵심이다.
생성 후 pruning이 아니라 생성 중 pruning을 해야 실제 wall-clock generation cost까지 줄일 수 있다.
token log-prob 기반 confidence는 reasoning 품질과 어긋날 수 있으며, 부분 추론의 성공 가능성을 직접 학습하는 head가 더 효과적이다.
quality head를 학습용 pruning에만 쓰지 않고 test-time scaling의 voting weight로 재사용한 점이 실용적이다.
효율성과 성능 향상을 동시에 얻은 이유는 단순 계산 절감이 아니라, 그룹 내 정오답 균형을 인위적으로 맞춰 advantage 신호를 강화했기 때문이다.

Limitations

논문은 주로 수학 문제처럼 정답 검증이 명확한 RLVR 환경에서만 평가되었으며, UI interaction이나 tool-use agent 같은 더 일반적인 reward 기반 환경으로의 확장 가능성은 주장하지만 실험적으로 검증하지는 않았다. 또한 pruning이 완전히 즉시 일어나는 것은 아니고, 모든 rollout이 최소한 detection length인 512 토큰까지는 생성되어야 하므로 generation 단계의 가속은 후속 단계만큼 크지 않다. quality head의 예측 품질과 calibration이 학습 초기에 불안정할 수 있어 cold-start 구간이 필요하다는 점도 운영상 제약이다.