기술 요약: PS-PPO: Critic-Free RLHF를 위한 Prefix-Sampling PPO

1. 문제 정의

인간의 선호도에 맞춰 대규모 언어 모델(LLM)을 정렬하는 방법으로 강화 학습으로부터의 인간 피드백(RLHF)이 표준적인 방법으로 사용되고 있습니다. 그러나 RLHF로 LLM을 학습시키는 과정은 별도의 가치 네트워크(value network)를 학습시키지 않는 critic-free 방식(예: GRPO, RLOO)에서도 상당한 계산 병목 현상에 직면합니다.

이러한 critic-free 방식에서는 생성된 전체 궤적(completion)에 대해 단일 스칼라 보상이 할당됩니다. 이 보상은 어드밴티지(advantage)를 계산하기 위해 궤적 내의 모든 토큰에 균등하게 브로드캐스트됩니다. 결과적으로, 정책 업데이트를 위해서는 모든 롤아웃(rollout)에 대해 **전체 궤적 역전파(full-trajectory backpropagation)**가 필요합니다. 저자들은 이것이 비효럽적이라고 주장하는데, 많은 추론 작업(예: 단계별 수학 풀이)에서 중간 단계의 접두사(prefix)가 최종 결과를 결정하기에 충분한 정보를 이미 포함하고 있는 경우가 많기 때문입니다. 저자들의 경험적 분석(그림 1)에 따르면 궤적의 성공률은 완료가 끝나기 훨씬 전에 안정화되는 경st를 보이는데, 이는 궤적의 접미사(suffix)에 있는 토큰들이 중복된 학습 신호를 운반하고 있음을 시사합니다. 따라서 전체 궤적 업데이트는 이러한 중복된 접미사에 상당한 컴퓨팅 자원과 GPU 메모리를 낭비하게 됩니다.

2. 방법론: PS-PPO

저자들은 역전파를 절단(truncation)하면서도 편향되지 않은 그래디언트 추정치를 유지함으로써 계산 효율성을 높이는, compute-efficient하고 critic-free한 방법인 **Prefix-Sampling Proximal Policy Optimization (PS-PPO)**를 제안합니다.

핵심 메커니즘

전체 완료 문장 $o = [o_1, \dots, o_T]$에 대해 역전파를 수행하는 대신, PS-PPO는 각 궤적에 대해 절단 타임스텝(cutoff timestep) $H$를 샘플링합니다. 정책 업데이트는 오직 접두사 $o_{\le H} = [o_1, \dots, o_H]$에 대해서만 계산됩니다.

이 절단이 편향을 유발하지 않도록 하기 위해, PS-PPO는 재가중치 부여(reweighting) 방식을 채택합니다.

1. 절단 분포(Cutoff Distribution): 프롬프트에 조건화된 분포 $\xi_{1:T}$가 정의되며, 여기서 $\xi_t = \Pr(H \ge t | x)$는 $t$번째 토큰이 유지될 확률을 나타냅니다. 이 분포는 비증가 함수($\xi_1 \ge \xi_2 \ge \dots \ge \xi_T$)입니다.
2. 불편 추정량(Unbiased Estimator): 절단된 궤적에 대한 그래디언트는 역 포함 확률(inverse inclusion probability) $1/\xi_t$에 의해 재가중치됩니다. 결과적인 추정량 $\hat{G}(\theta)$는 다음과 같습니다:
   $$\hat{G}(\theta) := \frac{1}{K} \sum_{k=1}^K \sum_{t=1}^{H^{(k)}} \frac{1}{\xi_t} g_t^{(k)}(\theta)$$
   여기서 $g_t^{(k)}$는 타임스텝별 그래디언트입니다. $H$의 샘플링에 대한 이 추정량의 기댓값은 전체 궤적 업데이트를 복구합니다.

절단 분포 최적화

저자들은 $\xi_{1:T}$의 선택을 컴퓨팅 예산(compute budget) $B$ (역전파되는 기대 토큰 수) 제약 조건 하에서 분산을 최소화하는 볼록 최적화(convex optimization) 문제로 정식화합니다.

• 분산 대용치(Variance Surrogate): 저자들은 중요도 샘플링(importance sampling)에 의해 발생하는 분산의 다루기 쉬운 대용치를 유도하며, 이를 $\sum w_t(x) (1/\xi_t - 1)$로 근사합니다. 여기서 $w_t(x)$는 타임스텝 $t$에서의 업데이트 중요도를 나타냅니다.
• 중요도의 대용치(Proxy for Importance): 정확한 스코어 노름 $\|\nabla_\theta \log \pi_\theta\|^2$을 계산하려면 전체 역전파가 필요하므로(이는 목적에 어긋남), 저자들은 출력 헤드 활성화와 보상 불확실성에 기반한 **포워드 전용 대용치(forward-only proxy)**를 사용합니다.
  • 보상 불확실성 ($u_t(x)$): 접두사 상태에 따른 보상의 분산을 사용하여 추정하며, 성공적인 롤아웃과 실패한 롤아웃 간의 다음 토큰 분포 간 거리를 통해 근사합니다.
  • 그래디언트 노름 대용치 ($\bar{\gamma}_t$): 출력 헤드의 그래디언트 노름으로부터 유도됩니다.
• 최적화: 문제는 단조성(monotonicity)과 예산 제약을 준수하면서 $\sum \frac{\bar{\gamma}_t(x,t) u_t(x)}{\xi_t}$를 최소화하는 것입니다. 해는 단조성 제약을 강제하기 위해 PAV(Pool Adjacent Violators) 알고리즘을 사용하여 찾아집니다.

알고-리즘 흐름

1. 프롬프트에 대해 $K$개의 완료 문장을 생성합니다.
2. 터미널 보상을 계산하고 어드van티지를 브로드캐스트합니다.
3. 롤아웃으로부터 타임스텝별 불확실성과 그래디언트 대용치를 추정합니다.
4. 볼록 최적화를 풀어 최적의 단조 절단 확률 $\xi^*_{1:T}$를 얻습니다.
5. 정의된 $\xi^*$ 분포로부터 각 궤적에 대한 절단 $H^{(k)}$를 샘플링합니다.
6. $t \le H^{(k)}$인 토큰에 대해서만 역전파를 수행하며, 손실을 $1/\xi^*_t$로 재가중합니다.

3. 주요 기여

1. PS-PPO 프레임워크: 역전파의 확률적 절단을 도입하여 계산 및 메모리 비용을 크게 줄이는 새로운 compute-efficient, critic-free RLHF 방법론을 제시합니다.
2. 불편 절단 추정량(Unbiased Truncation Estimator): 추가적인 롤아웃이나 보조 가치 모델 없이도 기대값 측면에서 전체 시퀀스 브로드캐스트 업데이트를 복구하는 역 포함 확률 재가중 추정량을 이론적으로 유도했습니다.
3. 최적화된 절단 전략: 그래디언트 중요도와 보상 불확실성에 대한 포워드 전용 대용치를 사용하여, 분산 감소와 컴퓨팅 예산 사이의 균형을 맞추는 볼록 최적화 문제를 해결함으로써 어디서 절단할지를 결정하는 원칙적인 접근 방식을 제공합니다.
4. 경험적 검증: PS-PPO가 강력한 critic-free 베이스라인(GRPO, DAPO, RLOO)과 대등한 정확도를 달성하면서도 훈련 시간과 피크 GPU 메모리를 실질적으로 줄임을 입증했습니다.

4. 실험 결과

저자들은 Llama-3.1-8B-Instruct 및 Qwen2.5-Math-7B를 사용하여 수학적 추론 벤치마크(MATH500, AMC 2023, Minerva Math, AIME 2024/2025)에서 PS-PPO를 평가했습니다.

• 성능: PS-PPO (Optimized)는 모든 벤치마크에서 DAPO 및 S-GRPO와 같은 강력한 베이스라인과 대등하거나 약간 더 나은 Pass@1 정확도를 달ر성했습니다. 예를 들어, MATH500에서 Llama-3.1-8B를 사용했을 때 PS-PPO는 47.6%를 기록하여 DAPO의 46.8%보다 높은 성능을 보였습니다.
• 효율성:
  • 훈련 시간: PS-PPO는 베이스라인 대비 스텝당 훈련 시간을 33%~45% 단축했습니다. 이는 주로 절단된 시퀀스의 포워드/백워드 비용 감소 덕분이며, 절단 분포를 계산하는 오버헤드를 상쇄하고도 남았습니다.
  • 메모리: 절단된 시퀀스에 대한 활성화 저장 공간이 줄어듦에 따라 피크 GPU 메모리 사용량이 15%~17% 감소했습니다.
  • 스케일링: 효율성 이득은 최대 완료 길이($T_{max}$)가 길어질수록 더욱 두드러졌습니다. $T_{max}=4096$에서 PS-PPO는 업데이트 단계에서 S-GRPO보다 2.8배, DAPO보다 3.3배 빨랐습니다.
• Ablations (절제 실험):
  • 예산 ($B$): $B=128$ 토큰의 예산이 정확도와 속도 사이의 최적의 트레이드오프를 제공했습니다. 더 작은 예산은 분산을 높였고, 더 큰 예산은 높은 계산 비용과 함께 완만한 정확도 향상만을 보였습니다.
  • 절단 전략: (볼록 최적화를 푸는) "Optimized" 전략은 uniform, time-prior, heuristic 절단 전략보다 성능이 월등히 뛰어났으며, 이는 정보가 풍부한 접두사를 적응적으로 선택하는 것이 매우 중요하다는 것을 확인시켜 줍니다.
  • 롤아웃 ($K$): $K$를 늘리면 안정성과 성능이 향상되지만 스텝당 계산량이 증가합니다. $K=8$이 실용적인 트레이드오프 지점으로 식별되었습니다.

5. 의의 및 주장

본 논문은 PS-PPO가 훨씬 적은 자원으로 대규모 언어 모델을 정렬할 수 있는 확장 가능한 경로를 제공한다고 주장합니다. 핵심적인 의의는 추론 과정의 중간 접두사가 최종 결과를 결정하기에 충분한 정보를 담고 있는 경우가 많다는 관찰에 있으며, 이를 통해 전체 궤적 업데이트의 계산적 중복성을 제거합니다.

확률적 접두사 샘플링과 중요도 가중치를 통해 학습 신호를 전체 시퀀스 길이로부터 분리함으로써, PS-P로 다음과 같은 효과를 얻을 수 있습니다:

• 하드웨어 장벽 완화: 낮은 피크 메모리 및 계산 비용은 RLHF를 학계 및 자원이 제한된 그룹에서도 더 쉽게 접근 가능하게 만듭니다.
• 효율적인 롱 컨텍스트(Long-Context) 훈련: 이 방법은 전체 역전파 비용이 과도한 긴 추론 궤적에 특히 효과적입니다.

저자들은 PS-PPO가 효율성을 개선하는 최적화 방법이지, 모델의 안전성이나 행동을 본질적으로 결정하는 것은 아니라는 점을 명시하며 신중한 태도를 유지합니다. 모델의 안전성은 여전히 보상 규격, 학습 데이터 및 배포 안전장치에 달려 있습니다. 이 방법은 기대값 측면에서 전체 시퀀스 목적 함수의 이론적 특성을 보존하므로, 효율성 이득이 편향된 학습을 초래하지 않음을 보장합니다.