FLOPs 재활용: 매우 낮은 확률로 생성된 초기 데이터에 기반하여 강화 학습을 확장하여 어려운 문제를 해결하는 방법

Reuse your FLOPs: Scaling RL on Hard Problems by Conditioning on Very Off-Policy Prefixes

대규모 언어 모델(LLM) 추론을 위한 일반적인 강화 학습(RL) 방법은 올바른 on-policy 데이터가 드물고, 정책 경사가 사라지며, 학습이 정체되는 어려운 문제에서 컴퓨팅 자원을 낭비합니다. 보다 효율적인 RL을 위해, 우리는 이전 추론 또는 RL 훈련 과정에서 생성된 데이터를 활용하여 off-policy 데이터 형태로 재사용하는 방법을 고려합니다. 일반적인 off-policy 방법은 off-policy 데이터에 대한 감독 학습을 수행하지만, 이는 RL 최적화 과정에서 불안정성을 야기합니다. 우리는 PrefixRL을 제안합니다. PrefixRL은 성공적인 off-policy 데이터의 초기 부분(prefix)에 조건을 부여하고, 이를 기반으로 on-policy RL을 수행하여 나머지 부분을 완성함으로써 off-policy로 인한 불안정성을 해결합니다. PrefixRL은 off-policy prefix의 길이를 조절하여 문제의 난이도를 조절함으로써 어려운 문제에 대한 학습 신호를 향상시킵니다. 우리는 PrefixRL의 목표 함수가 표준 RL 목표 함수와 일관성을 유지하며, 더 높은 샘플 효율성을 가진다는 것을 증명합니다. 실험적으로, 우리는 '백-일반화(back-generalization)' 현상을 발견했습니다. 즉, prefix가 추가된 문제에 대해서만 학습한 모델이 분포 외부의 prefix가 없는 문제에 대해서도 뛰어난 성능을 보이며, 학습된 전략이 종종 prefix에 포함된 전략과 다르다는 것을 확인했습니다. 우리의 실험에서는, 기본 모델을 사용하여 rejection sampling을 통해 off-policy 데이터를 생성하여 자체 개선 루프를 구축했습니다. 어려운 추론 문제에서, PrefixRL은 가장 강력한 기준 모델(off-policy 데이터에 대한 SFT 후 RL)보다 2배 더 빠르게 동일한 훈련 보상을 달성했으며, 초기 rejection sampling에 사용된 컴퓨팅 비용까지 고려하더라도 3배 더 높은 최종 보상을 달성했습니다. 이러한 성능 향상은 다른 모델 패밀리에서 파생된 off-policy 데이터를 사용할 때에도 유지되며, PrefixRL의 실용적인 적용 가능성을 검증합니다.

Typical reinforcement learning (RL) methods for LLM reasoning waste compute on hard problems, where correct on-policy traces are rare, policy gradients vanish, and learning stalls. To bootstrap more efficient RL, we consider reusing old sampling FLOPs (from prior inference or RL training) in the form of off-policy traces. Standard off-policy methods supervise against off-policy data, causing instabilities during RL optimization. We introduce PrefixRL, where we condition on the prefix of successful off-policy traces and run on-policy RL to complete them, side-stepping off-policy instabilities. PrefixRL boosts the learning signal on hard problems by modulating the difficulty of the problem through the off-policy prefix length. We prove that the PrefixRL objective is not only consistent with the standard RL objective but also more sample efficient. Empirically, we discover back-generalization: training only on prefixed problems generalizes to out-of-distribution unprefixed performance, with learned strategies often differing from those in the prefix. In our experiments, we source the off-policy traces by rejection sampling with the base model, creating a self-improvement loop. On hard reasoning problems, PrefixRL reaches the same training reward 2x faster than the strongest baseline (SFT on off-policy data then RL), even after accounting for the compute spent on the initial rejection sampling, and increases the final reward by 3x. The gains transfer to held-out benchmarks, and PrefixRL is still effective when off-policy traces are derived from a different model family, validating its flexibility in practical settings.