OR Else: A Differentiable Trust Region for Policy Optimization

기술 요약: OR ELSE: 정책 최적화를 위한 미분 가능한 신뢰 영역 (A DIFFERENTIABLE TRUST REGION FOR POLICY OPTIMIZATION)

문제 정의
대규모 언어 모델(LLM)을 위한 인간 피드백 기반 강화 학습(RLHF)은 흔히 근사 정책 최적화(PPO)와 그 변형인 그룹 상대적 정책 최적화(GRPO)에 의존한다. 두 방법 모두 정책 드리프트(policy drift)를 방지하기 위해 클리핑된 대리 목적 함수(clipped surrogate objective)를 사용하여 정책 업데이트를 제한한다. 클리핑 메커니즘은 "유리한" 방향으로의 과도한 정책 변화를 막기 위해 한쪽 방향으로 그래디언트가 0이 되는 영역을 생성하지만, 이 영역으로의 전이는 기하학적으로 급격하다. 즉, 클리핑 경계에서 목적 함수의 미분값이 확률 비율(probability ratio)에 대해 불연속적으로 변한다. 저자들은 원래 분류기 학습을 위해 설계된 매끄러운 일방향 포화 규칙(one-sided saturation rule)인 **출력 리셋(Output Reset, OR)**이 LLM 사후 학습(post-training)에 더 우수한 대안이 될 수 있는지 조사한다. 구체적으로, 본 논문은 클리핑된 대리 함수를 미분 가능한 OR 기반 손실 함수로 교체하는 것이 새로운 병리 현상을 도입하지 않으면서 최적화 안정성과 최종 보상 점수를 개선할 수 있는지 묻는다.

방법론
본 논문은 클리핑된 정책 항을 로그-비율(log-ratio) 공간에서의 OR 유도 제곱 마진 손실(squared-margin loss)로 대체하는 두 가지 변형인 PPO-OR과 GRPO-OR을 소개한다.

• 핵심 메커니로: 확률 비율 $r_t(\theta)$를 사용하는 대신, 토큰별 로그 확률 비율 $\rho_t = \log \pi_\theta(a_t|s_t) - \log \pi_{\theta_{old}}(a_t|s_t)$를 기반으로 작동한다.
• OR 손실: 비제로(non-zero) 추정 어드밴티지 $\hat{A}_t$를 가진 토큰에 대해, 타겟 마진 $\sigma_t \alpha_{OR}$이 설정된다 (여기서 $\sigma_t = \text{sign}(\hat{A}_t)$). 이는 제곱 일방향 마진 형태를 띤다: $\ell_{OR,t} = [\max(0, \alpha_{OR} - \sigma_t \rho_t)]^2$.
• 출력 리셋 동작: $\rho_t$가 어드밴티지의 부호가 선호하는 방향으로 마진을 넘어서면, 타겟은 현재의 $\rho_t$로 "리셋"된다. 이는 해당 토큰에 대해 제로 잔차(zero residual)와 제로 직접 그래디언트 기여를 초래한다. 미분값은 경계에서 연속적으로 0에 수렴하며, 이는 PPO 클립의 $C^0$ 연속성(미분 불가능)과 달리 $C^1$ 연속성을 갖는다.
• 어드밴티지 추정:
  • PPO-OR: 학습된 가치 헤드(value head)를 갖는 일반화 어드밴티지 추정(GAE)을 유지한다.
  • GRPRO-OR: 크리틱(critic) 없이 그룹 상대적 어드밴티지(프롬프트당 $G=2$개의 응답으로부터 계산됨)를 유지한다.
• 실험 설정: 본 연구는 Anthropic hh-rlhf 데이터셋을 사용하여 Llama-3.2-1B-Instruct를 사용한다. 모든 방법은 단일 고정된 보상 모델을 공유하며, 세 개의 시드(seed)에 대해 평가된다. 실험은 두 가지 매칭된 패밀리 내 비교로 설계되었다: PPO-clip vs. PPO-OR (GAE 하에서), 그리고 GRPO vs. GRPO-OR (그룹 상대적 어드밴티지 하에서). 패밀리 간 비교는 어드밴티지 추정, 크리틱 사용, 응답 길이의 차이로 인해 통제되지 않은 것으로 간주된다.

주요 기여

1. 알고리즘 제안: 클리핑된 대리 함수를 로그-비율 공간에서의 연속 미분 가능한 일방향 제곱 마진 손실로 대체하는 PPO-OR 및 GRPO-OR을 도입한다.
2. 이론적 규명: OR 목적 함수가 포화 경계에서 $C^1$ 연속이며, 샘플 수준의 업데이트 방향이 어드밴티지 부호에 의해 엄격하게 결정되고 마진을 넘어서면 연속적으로 소멸함을 공식적으로 증명한다.
3. 경험적 평가: 클리핑과 OR를 두 가지 서로 다른 어드밴티지 추정 패밀리 내에서 엄격하게 비교하여, OR의 이점이 균일하지 않으며 주변 최적화 파이프라인에 따라 달라짐을 강조한다.
4. 진단적 통찰: 클리핑을 OR로 교체하는 것이 그룹 상대적 설정에서 관찰되는 거대한 누적 정책 드리프트 문제를 본질적으로 해결하지 못한다는 점을 식별하여, 국소적 그래디언트 포화와 전역적 정책 제어를 구분한다.

결과
논문은 세 개의 시드에 대한 훈련 시간 보상 모델 점수와 최적화 진단 결과를 보고한다:

• GAE 패밀리 (PPO-clip vs. PPO-OR):

  • 보상 점수: PPO-OR은 PPO-clip (+0.195) 대비 높은 평균 최종 점수 (+0.500)를 달성하였으며, 그 차이는 +0.305이다.
  • 안정성: PPO-OR은 PPO-clip (0.110)보다 더 큰 시드 간 표준 편차 (0.158)를 보였다.
  • 드리프트: 두 방법 모두 유사하고 낮은 롤아웃-현재 로그-비율 변위(약 0.3–0.5)를 유지했다.

• 그룹 상대적 패밀리 (GRPO vs. GRPO-OR):

  • 보상 점수: GRPO-OR은 평균 점수를 개선하지 못했으며, GRPO (+0.488)보다 약간 낮았다 (+0.457).
  • 안정성: GRPO-OR은 GRPO (0.080)에 비해 현저히 작은 시드 간 표준 편차 (0.027)를 보였다.
  • 진단: GRPO-OR은 거의 제로에 가까운 터미널 OR 잔차와 감소하는 오버슈트 분율을 달성한 반면, 매칭된 GRPO 트레이스는 여전히 크고 변동성이 높았다.
  • 드리프트: 결정적으로, GRPO-OR은 GRPO에서 관찰된 큰 롤아웃-현재 변위를 줄이지 못했다. 두 그룹 상대적 방법 모두 GAE 방법들에 비해 실질적으로 높은 변위(약 5–13)를 나타냈다.

의의 및 주장
논문은 관찰된 OR의 효과가 문맥 의존적임을 강조하며 다음과 같이 겸허하게 결론짓는다:

• 보편적 신뢰 영역이 아님: 저자들은 OR의 매끄러운 포화가 공식적인 신뢰 영역을 구성하지 않는다는 점을 명시적으로 밝힌다. 이는 개별 토큰이 마진을 넘었을 때 직접적인 그래디언트 기여를 억제하지만, 공유된 파라미터, 보조 손실 또는 반복 간의 변화하는 어드밴티지 부호로 인해 발생하는 누적 정책 드리프트를 제한하지는 못한다.
• 파이프라인 의존성: OR의 효능은 어드밴티지 추정 방식과 결부되어 있다. GAE 설정에서 OR은 더 높은 보상 점수와 상관관계가 있었으나 분산을 증가시켰다. 그룹 상대적 설정에서 OR은 분산을 줄이고 손실 트레이스를 안정화했으나, 평균 보상 점수를 개선하거나 정책 드리프트를 줄이지는 못했다.
• 한계점: 연구는 결과가 경량 1B 모델, 단일 고정 보상 모델, 그리고 제한된 그룹 크기($G=2$)에 특화되어 있음을 인정한다. 본 연구는 PPO나 GRPO의 일반적인 우월성을 확립하기보다는 이 특정 실험 환경을 규명하는 데 중점을 둔다. 논문은 그룹 상대적 방법에서의 거대한 정책 드리프트가 여-전히 해결되지 않은 문제임을 시사하며, OR만으로는 이를 해결할 수 없음을 지적하고, 향-후 연구가 더 강력한 참조 정규화(reference regularization)나 명시적인 변위 페널티를 통해 누적 변위를 다루어야 한다고 제언한다.

요약하자면, 본 논문은 출력 리셋(Output Reset)이 클리핑에 대한 미분 가능한 대안을 제공하여 최적화 동작을 변화시키고 특정 문맥에서 보상 점수를 개선할 수 있음을 보여주지만, 이것이 그룹 상대적 RLHF의 안정성이나 드리프트 문제를 본질적으로 해결하는 것은 아님을 입증한다.