10.2 PPO for Language-Model Alignment

Proximal Policy Optimization (PPO)은 고전적인 RLHF system에서 사용된 online reinforcement-learning 방법입니다 [1] [2]. Language-model PPO에서는 policy가 response를 생성하고, 보상 모델(reward model)이 점수를 주고, value 헤드(value head)가 return을 추정하며, frozen reference policy가 drift를 제한합니다. 이 중 하나를 빼면 다른 알고리즘이 됩니다.

네 모델 관점

실무 RLHF PPO 실행은 다음을 구분합니다.

policy: SFT checkpoint에서 시작해 update합니다.
reference policy: token-level KL 제어에 쓰는 frozen snapshot입니다.
보상 모델(RM): prompt/response를 terminal 또는 shaped reward로 바꿉니다.
value head/critic: 생성한 response 각 위치의 return을 예측합니다.

Policy와 value가 backbone을 공유하거나 reference/RM을 offload할 수 있습니다. Revision, tokenizer/chat template, precision, quantization, 공유 parameter를 기록합니다. 평가 artifact에는 decoding, tool, reward preprocessing도 포함됩니다.

Rollout과 Reward Shaping

Tokenizer-native chat template과 add_generation_prompt=True로 prompt를 rendering합니다. Variable-length response를 생성하고 EOS를 보존하며, padding 전 실제 생성 token에만 1인 응답 마스크(response mask) 를 만듭니다. Prompt, padding, terminal EOS 뒤 token에는 policy/value loss를 주지 않습니다.

Response token $t$ 의 일반적인 shaped reward는

r_t = -\beta\left(\log\pi_{\text{old}}(a_t\mid s_t)- \log\pi_{\text{ref}}(a_t\mid s_t)\right)

이며 마지막 유효 response token에 RM score를 더합니다. 이 토큰별 KL(per-token KL) sample penalty가 모든 analytic KL estimator와 같은 것은 아니므로 구현을 문서화합니다. Reward hacking을 볼 수 있도록 raw RM reward, KL penalty, final shaped reward를 분리해 기록합니다.

RM은 자체적으로 고정한 tokenizer/template과 보정한 domain을 사용해야 합니다. Scalar reward에는 length, position, style bias가 있을 수 있습니다. 최적화를 믿기 전에 slice별 reward distribution과 human agreement를 평가합니다.

Return, Value Head, GAE

Value prediction $V_t$ 와 temporal-difference residual

\delta_t = r_t + \gamma V_{t+1} - V_t,

Generalized Advantage Estimation (GAE) 은

A_t = \delta_t + \gamma\lambda(1-d_t)A_{t+1},

입니다. $d_t$ 는 terminal/padded 위치를 표시합니다. Recursion은 response mask 안에서만 계산합니다. Return target은 $R_t=A_t+V_t$ 입니다. 많은 system이 rollout batch에 masked 어드밴티지 정규화(advantage whitening) 를 적용합니다. Distributed 구현은 일관된 global statistic을 사용하고 padding zero를 포함하지 않아야 합니다.

Masked value loss와 선택적인 value clipping으로 value head를 학습하고 PPO clipped surrogate로 policy를 학습합니다.

L^{\text{clip}}_t=\min\left(r_t(\theta)A_t, \operatorname{clip}(r_t(\theta),1-\epsilon,1+\epsilon)A_t\right),

여기서 $r_t(\theta)=\exp(\log\pi_\theta-\log\pi_{\text{old}})$ 는 probability ratio이며 위 reward와 다른 기호입니다.

Advantage (\hat{A}_t): 1.0

Positive Advantage: The action was better than expected. Objective is clipped when r_t > 1.2 to prevent over-updating.

Unclipped Value

Final Objective

Canonical Update Skeleton

아래 pseudocode는 tensor 의미론을 보여 줍니다. Model wrapper, generation backend, distributed rollout store가 다르므로 runnable이라고 표시하지 않습니다.

rollout = generate_with_old_policy(prompts)
response_mask = rollout.response_mask

with no_grad():
    old_logp = completion_logp(old_policy, rollout, response_mask)
    ref_logp = completion_logp(reference, rollout, response_mask)
    rm_score = reward_model_score(rollout)
    old_values = value_head(rollout)

token_rewards = -kl_beta * (old_logp - ref_logp)
token_rewards[rollout.last_valid_index] += rm_score
advantages, returns = masked_gae(
    token_rewards, old_values, response_mask, gamma, gae_lambda
)
advantages = masked_global_whiten(advantages, response_mask)

for _ in range(ppo_epochs):
    new_logp = completion_logp(policy, rollout, response_mask)
    new_values = value_head(rollout)
    ratio = (new_logp - old_logp).exp()
    policy_loss = masked_clipped_surrogate(ratio, advantages, response_mask)
    value_loss = masked_clipped_value_loss(new_values, old_values, returns, response_mask)
    entropy = masked_entropy(policy, rollout, response_mask)
    optimize(policy_loss + value_coef * value_loss - entropy_coef * entropy)

모든 rollout에 sampling policy revision과 old log probability를 저장합니다. Epoch가 너무 많거나 consumer가 늦으면 정책 지연(policy lag) 이 커집니다. 데이터가 현재 policy에 비해 stale해져 importance ratio를 신뢰하기 어렵습니다. 최대 rollout age를 정하고 stale batch는 버리거나 다시 생성합니다.

안정성과 시스템 계약

SFT policy, reference, RM, value initialization, tokenizer/template, dataset manifest, generation setting, reward normalization, KL controller, $\gamma$ , $\lambda$ , clipping coefficient, response-token 기준 batch, seed, distributed topology를 고정합니다. 네 모델 identity와 optimizer, scheduler, scaler/FP8 state, RNG, rollout cursor/store state, global token을 저장합니다.

다음을 관측합니다.

raw RM reward, non-score reward, per-token KL, shaped reward
policy/value loss, value explained variance, advantage mean/std, entropy
clip fraction, ratio distribution, response length/EOS, truncation, padding
nonfinite/overflow, gradient norm, rollout/train throughput, queue age, policy lag
held-out human preference, capability retention, safety, format/tool, cost slice

Nonfinite가 반복되거나 value가 발산하거나, KL/entropy/length가 범위를 벗어나거나, clip fraction 또는 stale rollout이 과도하거나, held-out 개선 없이 reward만 오르거나, critical safety가 회귀하면 중단합니다. Frozen pre-PPO bundle을 last-known-good로 유지합니다.

신흥 Variant

Decoupled rollout, 대안 trust region, replay, Outer-PPO나 HP3O 같은 hybrid-policy 방법이 연구되고 있습니다. 이는 신흥(emerging) 설계점이며 canonical loop를 이해하는 일을 대신하지 않습니다. Model, RM, rollout budget, data, evaluation이 같은 조건에서 비교하고 maturity와 구현 공개 여부를 명시합니다.

Quizzes

Quiz 1: LLM PPO에 value head가 필요한 이유는 무엇인가요?

Value head는 각 유효 response 위치의 baseline과 return을 추정합니다. GAE는 이를 사용해 variance를 낮추고 생성 trajectory에 credit을 배분합니다.

Quiz 2: Padding에 EOS token을 써도 response mask가 필요한 이유는 무엇인가요?

EOS와 padding ID가 같을 수 있고 sequence 길이가 다릅니다. 명시적 mask가 prompt, padding, terminal 뒤 위치를 KL, advantage, policy, value, entropy reduction에서 제외합니다.

Quiz 3: PPO의 reward hacking을 어떤 신호로 찾을 수 있나요?

Raw RM reward는 오르지만 held-out human preference, correctness, safety는 오르지 않거나, length/style과 KL이 이동하거나, reward gain이 biased slice에 집중됩니다. Reward component와 fixed evaluation을 분리하면 드러납니다.

Quiz 4: Policy lag가 위험한 이유는 무엇인가요?

오래된 policy가 만든 rollout은 점점 off-policy가 됩니다. Importance ratio가 퍼지고 clipping이 지배하며 stale behavior로 update합니다. Rollout age를 제한하고 데이터를 다시 생성해야 합니다.

References

Schulman, J., et al. (2017). Proximal Policy Optimization Algorithms. arXiv:1707.06347.
Ouyang, L., et al. (2022). Training language models to follow instructions with human feedback. arXiv:2203.02155.
Zheng, R., et al. (2023). Secrets of RLHF in Large Language Models Part I: PPO. arXiv:2307.04964.