Hugging Face TRL RLOO 트레이너 출시

Hugging Face는 TRL 라이브러리에 RLOO (REINFORCE Leave One-Out) 트레이너를 도입했습니다. RLOO는 Proximal Policy Optimization (PPO)에 대한 보다 접근하기 쉽고 효율적인 대안으로 설계된 온라인 RLHF 학습 알고리즘으로, GPU 메모리를 크게 줄이고 수렴 시간을 단축합니다.

성능 및 효율성 향상

RLOO는 PPO에 비해 자원 활용도와 학습 속도에서 상당한 개선을 제공합니다:

Memory Efficiency: RLOO는 모델 크기에 따라 PPO보다 약 50-70% 적은 vRAM을 사용합니다.
Training Speed: RLOO는 1B 파라미터 모델에서는 PPO보다 2배 빠르게, 6.9B 모델에서는 최대 3배 빠르게 실행됩니다.
Model Quality: RLOO는 응답 승률(GPT-4가 평가) 측면에서 PPO와 경쟁력 있게 동작하며, Direct Preference Optimization (DPO)과 같은 오프라인 방법보다 일관되게 우수합니다.

기술 아키텍처 및 동기

전통적인 PPO는 정책 모델, 레퍼런스 정책 모델, 보상 모델, 가치 모델 네 개의 모델 복사본을 메모리에 로드해야 하므로 자원 집약적입니다. RLOO는 가치 모델을 제거하여 아키텍처를 단순화하고, 정책 모델, 레퍼런스 정책 모델, 보상 모델 세 개만 필요합니다.

핵심 알고리즘 차이점

RLOO는 행동과 보상을 처리하는 방식에서 PPO와 차이가 있습니다:

Completion-Level Actions: PPO는 각 완성 토큰을 개별 행동으로 취급하지만, RLOO는 전체 모델 완성을 하나의 행동으로 취급합니다. 이는 일반적으로 EOS(end-of-sentence) 토큰에만 실제 보상이 주어지는 희소 보상 문제를 해결합니다; RLOO는 EOS 보상을 전체 완성에 할당합니다.
REINFORCE Loss: RLOO는 보상과 기준값의 차이를 행동의 로그 확률에 곱하는 REINFORCE 손실을 사용합니다. 이를 통해 PPO에서 필요로 하는 가치 모델과 Generalized Advantage Estimation (GAE)을 제거합니다.
Leave-One-Out Baselines: RLOO는 배치 내 다른 모든 샘플의 보상을 해당 샘플의 기준값으로 사용하여 기준값을 계산합니다. 특정 프롬프트에 대해, 하나의 완성에 대한 기준값은 동일한 프롬프트에 대해 생성된 다른 모든 완성의 평균 보상입니다.

TRL에서의 구현

RLOO 트레이너는 실험적인 PPOv2Trainer를 기반으로 구축되었습니다. RLOO는 개념적으로 REINFORCE 알고리즘에 기반하지만, TRL 구현에서는 REINFORCE 손실이 PPO 손실의 특수 경우이기 때문에 PPO 손실을 사용합니다.

벤치마크 결과

Pythia 1B 및 6.9B 모델에 대한 실험은 RLOO 구현의 효능을 보여줍니다:

Pythia 6.9B: GPT-4를 평가자로 사용하여 78.7% 선호율(k=2)을 달성했으며, 이는 원 논문에서 보고된 77.9% (k=4)와 74.2% (k=2)보다 높습니다.
Pythia 1B: SFT(감독 학습) 체크포인트의 21.3% 승률에 비해 40.1% 승률을 달성했습니다.

bf16에서의 수치 안정성 문제

RLOO는 bf16 정밀도를 사용할 때 상당한 수치 안정성 문제를 보입니다. 생성 중에 얻은 로그 확률은 훈련 전방 패스에서 얻은 로그 확률과 약간 차이가 있을 수 있습니다.

PPO에서는 클리핑 계수(보통 0.2)가 비율이 1.2를 초과하거나 0.8 이하인 개별 토큰에 대해 그래디언트를 0으로 만듭니다. RLOO에서는 전체 완성을 하나의 행동으로 취급하기 때문에 이러한 수치 차이가 누적됩니다. 이로 인해 전체 시퀀스에 대한 그래디언트가 더 자주 0이 됩니다. 실험 관찰에 따르면 PPO는 배치 데이터의 약 3%를 0으로 만들지만, bf16 환경에서 RLOO는 20-40%의 배치 데이터를 0으로 만듭니다.