Where Should RL Post-Training Compute Go? Model Size, Search, Learning, and Feedback

기술 요약: RL 포스트 트레이닝 연산량은 어디로 가야 하는가?

문제 정의
강화 학습(RL) 포스트 트레이닝은 추론, 계획 및 피드백 기반 로봇 학습 파이프라인을 위해 파운데이션 모델을 적응시키는 데 점점 더 많이 활용되고 있습니다. 그러나 현재의 스케일링 분석은 일반적으로 이러한 프로세스의 비용을 단일 집계 지표인 총 FLOPs로 요약합니다. 본 논문은 이러한 보고 방식이 정보가 불충분하다고 주장합니다. 고정된 포스트 트레이닝 예산 하에서, 연산 할당은 다음 네 가지 결합된 차원을 포함하는 제약된 자원 할당 문제입니다:

1. 모델 크기: 훈련 중인 정책(policy)의 용량.
2. 탐색/롤아웃(Search/Rollout): 후보 궤적(trajectory) 생성.
3. 학습/업데이트(Learning/Updates): 정책의 경사 기반 최적화.
4. 피드백/보상 평가(Feedback/Reward Evaluation): 보상 모델 또는 검증기(verifier)의 추론 비용.

다루고자 하는 핵심 질문은 다음과 같습니다: 고정된 FLOP 예산 하에서, 성능을 극대화하기 위해 연산을 이 구성 요소들 사이에 어떻게 배분해야 하는가? 저자들은 더 큰 정책이 토큰당 더 많은 연산을 소비하므로, 동일한 예산 하에서 더 적은 수의 롤아웃과 업데이트 단계를 구매하게 된다고 상정합니다. 또한, 학습된 피드백 메커니즘(예: 프로세스 보상 모델, PRM)의 도입은 롤아웃 및 업데이트 예산과 경쟁하는 새로운 연산 비용을 발생시키며, 이는 총 FLOP 보고 방식에서는 가려지는 트레이드오프입니다.

방법론
본 연구는 GRPO(Group Relative Policy Optimization) 포스트 트레이닝을 위해 설계된 FLOP 계측 프레임워크를 소개합니다. 총 연산량($C_{total}$)은 세 가지 별개의 항으로 분해됩니다:
$$C_{total} = C_{search} + C_{learning} + C_{reward}$$
여기서:

• $C_{search}$: 자기회귀적 롤아웃 생성.
• $C_{learning}$: 정책 최적화 (LoRA 어댑터 파라미터 및 동결된 백본 비용 고려).
• $C_{reward}$: 보상 또는 검증기 추론.

주요 할당 변수는 총 연산량에 대한 학습 연산의 비율로 정의되는 업데이트 분율($\rho$)입니다. 저자들은 Polaris-53K 수학적 추론 데이터셋을 사용하여 LoRA가 적용된 Qwen2.5 정책을 통해 실험을 수행했습니다. 실험 설계는 다음과 같습니다:

• 단계 A: 탐색-학습 할당을 연구하기 위해 모델 크기(1.5B, 3B, 7B)를 변화시킴.
• 단계 B & D: 피드백 비용이 할당 경계(allocation frontiers)를 어떻게 변화시키는지 평가하기 위해 다양한 보상 시스템(sparse, structured, dense, proxy-PRM, real PRM)을 비교함.
• 평로(Evaluation): 지표로는 네이티브 보상(native training reward), 다운스트림 pass@1 정확도(GSM8K 및 MATH-500), 그리고 "공통 심판(common judge)" 프로세스 품질 점수를 포함합니다.

방대한 탐색 비용 문제를 해결하기 위해, 저자들은 진단용 파일럿 그리드 프로토콜인 **RACE (Reward-Aware Compute Allocation)**를 도입합니다. RACE는 작은 규모의 IsoFLOP 파일럿 그리드를 실행하여 로컬 보상-할당 반응을 추정하고, 해당 설정이 "탐색 중심(search-heavy)", "업데이트 중심(update-heavy)", 또는 "보상 중심(reward-heavy)" 영역에 속하는지 식별합니다. 저자들은 RACE가 영역 식별을 위한 진단 도구이며, 홀드아웃(held-out) 성능 향상을 보장하는 도구는 아니라고 명시적으로 밝힙니다.

핵심 기여

1. FLOP 계측 프레임워크: RL 포스트 트레이닝 연산을 롤아웃, 업데이트, 보상 평가 구성 요소로 공식적으로 분해하여, 메커니즘적으로는 다르지만 동일한 총 FLOPs를 가진 실행 건들을 비교할 수 있게 합니다.
2. 결합된 모델 크기와 할당: 모델 선택과 훈련 할당이 불가분함을 보여주는 실증적 증거를 제시합니다. 고정된 예산 하에서 더 큰 정책은 본질적으로 더 적은 업데이트와 롤아웃을 구매하게 되며, 이는 최적의 할당 경계를 변화시킵니다.
3. 보상 및 타겟 의존성: "최적"의 할당은 보상 시스템(규칙 기반 vs PRM)과 평가 대상(네이티브 보상 vs 다운스트림 정확도 vs 프로세스 품질)에 따라 달라짐을 입증합니다. 예를 들어, PRM 기반 피드백은 보상 모델 추론에 예산의 일부가 소비되므로 최적의 업데이트 분율($\rho$)을 낮추는 방향으로 이동합니다.
4. RACE 프로토콜: 작은 파일럿 그리드를 통해 할당 영역을 식별하는 진단 방법으로, 검증 실행을 대체하기보다는 가이드하기 위한 용도입니다.

결과

• 조건부 할당 경계: 최적의 할당은 고정된 점이 아니라 조건부 경계입니다. 탐색 중심 설정은 학습 용량을 과소 활용하며, 업데이트 중심 설정은 탐색 제한(exploration-limited) 상태가 됩니다.
• 보상 시스템의 영향: 규칙 기반 보상(sparse, structured, dense)은 높은 업데이트 분율($\rho \approx 0.72$)을 선호하는 경향이 있습니다. 반면, PRM 스타일의 피드백(proxy 및 real)은 고정된 예산이 이제 보상 모델 추론을 커버해야 하므로, 최적의 네이티브 보상 경계를 더 낮은 업데이트 분율($\rho \approx 0.44–0.45$) 쪽으로 이동시킵니다.
• 타겟 간 불일치: 서로 다른 평가 타겟은 서로 다른 최적 할당을 선택합니다.
  • 네이티브 보상: 훈련 중에 사용된 특정 피드백 신호를 최대화하는 할당을 선호합니다.
    설정은 보상 시스템 전반에 걸쳐 높은 업데이트 영역(예: $a180$)을 선호하는 경우가 많습니다.
  • 공통 심판: 밀집형(dense) 및 PRM 보상에 대해 중간 정도의 업데이트 영역(예: $a100$)에서 정점을 찍는 경우가 많습니다.
    이러한 불일치는 한 지표를 최대화하는 레시피가 다른 지표를 최대화하지 않을 수 있음을 나타냅니다.
• RACE 성능: RACE는 검증 설정에서 올바른 영역(높음 vs 낮음 업데이트)을 높은 정확도(동일 영역 일치율 1.0)로 식별해냈으나, 베이스라인 대비 홀드아웃 보상의 전역적 향상을 보장하지는 못했습니다.

의의 및 주장
본 논문은 RL 포스트 트레이닝이 단순한 스케일링 법칙이라기보다 제약된 자원 할당 문제라고 겸손하게 주장합니다. 저자들은 총 FLOPs만을 보고하는 것은 모델 용량, 경험 생성, 피드백 평가 사이의 메커니즘적 트레이드오프를 숨긴다고 주장합니다.

이 연구의 의의는 커뮤니티가 **할당을 인지하는 연산량(allocation-aware compute)**을 보고할 것을 권장하는 데 있습니다. 즉, 총 FLOPs와 함께 롤아웃, 업데이트, 보상 평가의 구체적인 내역과 그에 따른 업데이트 분율 $\rho$를 함께 보고해야 한다는 것입니다. 이러한 투명성은 로봇 학습 시스템에서 매우 중요합니다. 왜냐나 RL 포스트 트레이닝을 통해 추론 또는 피드백 모듈을 개선할 때, 제한된 자원을 더 큰 모델, 더 광범위한 탐색, 또는 더 강력한 피드백 메커니즘 중 어디에 투자할지 결정해야 하기 때문입니다. 본 논문은 보편적인 "최선의" 레시피는 존재하지 않으며, 최적의 할당은 모델 크기, 보상 설계, 그리고 최적화하고자 하는 특정 평가 타겟에 공동으로 의존한다는 결론을 내립니다.