기술 요약: 중요한 곳에 롤아웃을 할당하라

문제 제기

강화 학습 (RL) 을 통한 대규모 언어 모델 (LLM) 의 사후 학습, 특히 온라인 온-정책 (on-policy) 설정에서 지배적인 계산 병목 현상은 최적화 단계가 아니라 롤아웃 (샘플링) 생성에 있습니다. 그룹 기반 정책 최적화 방법인 그룹 상대 정책 최적화 (GRPO) 와 같은 방법들은 동일한 프롬프트에 대해 생성된 여러 롤아웃을 비교하여 이점 (advantage) 을 계산합니다. 그러나 기존 접근 방식들은 학습 잠재력과 무관하게 배치 내 모든 프롬프트에 롤아웃 예산을 무분별하게 할당합니다.

이는 상당한 비효율로 이어집니다. 붕괴된 보상 분포를 가진 프롬프트 (모든 롤아웃이 동일한 보상을 얻는 경우, 즉 모두 정답이거나 모두 오답인 경우) 는 이점이 0 이 되어 그라디언트 신호에 기여하지 않지만, 여전히 비싼 샘플링 자원을 소모합니다. 반면, 높은 보상 분산을 가진 프롬프트는 가장 강력한 학습 신호를 제공합니다. 현재 방법들은 이러한 구간을 온라인으로 구분하지 못해 정보 가치가 없는 데이터에 컴퓨팅 자원을 낭비합니다. DAPO 와 같은 과샘플링 전략은 탐색을 개선하지만, 신호가 약한 프롬프트를 필터링하지 않고 샘플링 비용을 선형적으로 증가시킵니다.

방법론: 파일럿 - 커밋 (Pilot-Commit)

저자들은 프롬프트 평가와 활용을 분리하도록 설계된 예산 인식형 롤아웃 할당 프레임워크인 **파일럿 - 커밋 (PC)**을 제안합니다. 이 방법은 그룹 기반 업데이트가 정책 그라디언트 크기가 최대화되는 높은 보상 분산 구간에서 가장 효과적이라는 이론적 발견에 기반합니다.

이 프레임워크는 각 학습 단계에서 두 단계로 구성됩니다:

1. 파일럿 단계: 샘플링 배치 내 모든 프롬프트에 대해 소수의 롤아웃 ($n_{pilot}$) 을 샘플링하기 위해 총 롤아웃 예산의 일부가 사용됩니다. 시스템은 현재 정책 하에서 각 프롬프트의 경험적 성공률 (따라서 보상 분산) 을 추정합니다.

   • 프롬프트는 비대칭 임계값을 기반으로 필터링됩니다. "너무 쉽다"고 판단된 프롬프트 (성공률 $> p_{upper}$) 는 그라디언트 신호가 약해지므로 향후 샘플링에서 제외됩니다. "너무 어렵다"고 판단된 프롬프트 (성공률 $< p_{lower}$) 는 폐기되지 않고 다음 에포크로 연기되는데, 이는 정책이 개선됨에 따라 높은 분산 구간에 도달할 수 있음을 인정하기 위함입니다.
   • 목표 분산 범위 (일반적으로 $p \approx 0.5$ 근처) 에 속하는 프롬프트는 커밋 단계를 위해 선택됩니다.

2. 커밋 단계: 남은 롤아웃 예산 ($n_{commit}$) 은 파일럿 필터를 통과한 프롬프트에만 할당됩니다. 이러한 프롬프트는 학습 데이터를 생성하기 위해 샘플링 자원의 대부분을 받습니다.

3. 학습: 선택된 프롬프트에 대한 파일럿 및 커밋 롤아웃의 합집합을 사용하여 정책을 업데이트합니다.

주요 최적화:

• 바인딩 (Binding): 추론 오버헤드를 줄이기 위해 $t-1$ 단계의 파일럿 단계를 $t$ 단계의 커밋 단계에 바인딩하여 추론 패스를 중첩시키고, 최소한의 오프-정책 지연만 도입합니다.
• 리플레이 버퍼: 파일럿 필터를 통과한 초과 프롬프트는 현재 파일럿 단계에서 생존자가 너무 적을 경우 학습 배치를 채우기 위해 리플레이 버퍼에 저장되며, 노후화된 프롬프트는 부패를 방지하기 위해 제거됩니다.
• 제거 전략: 일관되게 해결된 프롬프트 (모든 파일럿 롤아웃이 정답인 경우) 는 예산 낭비를 방지하기 위해 샘플링 풀에서 영구적으로 제거됩니다.

주요 기여

• 이론적 통찰: 이 논문은 그룹 내 보상 결과가 최대한 불확실할 때 (즉, 높은 분산일 때) GRPO 그룹 그라디언트의 크기가 최대화된다는 것을 증명하는 분석적 증명 (제 3.1 명제) 을 제공합니다. 이는 보상 분산을 프롬프트의 정보성 직접 대변자로 확립합니다.
• 프레임워크 설계: 파일럿 - 커밋은 보조 모델이나 오래된 역사적 통계 없이 프롬프트 정보성에 대한 실시간 온라인 추정에 기반하여 고정된 롤아웃 예산을 동적으로 재할당하는 새로운 2 단계 할당 전략을 도입합니다.
• 시스템 구현: 이 방법은 verl RL 프레임워크의 새로운 레시피로 구현되었으며, 확장성을 보장하기 위해 오프-정책 바인딩 및 리플레이 버퍼와 같은 실용적 최적화를 포함합니다.

실험 결과

저자들은 AIME, AMC, MATH 500, DeepMath 를 포함한 수학 추론 벤치마크에서 15 억에서 140 억 파라미터까지의 모델 규모에 걸쳐 파일럿 - 커밋을 평가했습니다.

• 롤아웃 효율성: 파일럿 - 커밋은 GRPO 및 DAPO 에 비해 누적 롤아웃 수를 현저히 줄이면서 목표 기준 정확도에 도달합니다.
  • GRPO와 비교할 때, PC 는 누적 롤아웃 기준으로 목표 정확도에 1.5 배에서 1.9 배 더 빠르게 도달합니다.
  • DAPO와 비교할 때, PC 는 누적 롤아웃 기준으로 2.3 배에서 4.0 배 더 빠릅니다.
• 월시계 시간: 64 개의 H200 GPU 와 같이 고도로 병렬화된 환경에서도 샘플링 비용 감소는 실질적인 속도 향상으로 이어져, PC 는 단계당 월시계 시간 기준으로 GRPO 보다 1.21 배, DAPO 보다 1.83 배 더 빠릅니다.
• 학습 품질: PC 는 학습 전반에 걸쳐 기준선보다 높은 평균 보상 분산을 유지하며, 이는 업데이트가 더 강력한 그라디언트 신호를 가진 프롬프트에 집중됨을 나타냅니다. 결과적으로 PC 는 동일한 학습 단계 수 내에서 기준선보다 더 높은 피크 정확도를 달성하는 경우가 많습니다.
• 하이퍼파라미터 강건성: 제거 연구는 임계값이 예산을 절망적인 프롬프트에 낭비하거나 경계선상의 정보성 프롬프트를 폐기하는 극단으로 설정되지 않는 한, 파일럿/커밋 분할 비율 및 필터링 임계값의 변화에 대해 이 방법이 강건함을 보여줍니다.

중요성 및 주장

이 논문은 롤아웃 할당을 RL 사후 학습의 "1 급 문제 (first-class problem)"라고 주장합니다. "얼마나 많은 롤아웃을 생성할지"에서 "어떤 프롬프트에 이를 할당할지"로 초점을 이동시킴으로써, 파일럿 - 커밋은 균일 샘플링의 근본적인 비효율성을 해결합니다.

저자들은 그들의 접근 방식이 학습 목적을 변경하거나 고정된 커리큘럼 순서를 부과하지 않는다고 강조합니다. 대신, 이는 진화하는 정책에 맞춰 온라인으로 적응하여 학습 신호가 가장 강한 곳에 컴퓨팅 예산이 사용되도록 보장합니다. 결과는 그룹 기반 RL 사후 학습의 경우 지능적인 샘플링 할당이 최종 모델 성능을 희생하지 않고 종종 향상시키면서 상당한 효율성 향상을 가져올 수 있음을 시사합니다. 이 작업은 생성 속도뿐만 아니라 샘플링되는 데이터의 효과성에 초점을 맞춘 시스템 수준 처리량 최적화에 대한 보완적 개선으로 제시됩니다.