Hugging Face가 RLOO 알고리즘을 도입해 PPO 대비 GPU 메모리 50-70% 감소 및 2-3배 학습 속도 개선

Context

PPO는 RLHF 학습의 표준 알고리즘이지만 정책 모델, 참조 정책 모델, 보상 모델, 가치 모델 등 4개의 모델 복사본을 메모리에 로드해야 해 GPU 메모리 부담이 크고 구현 세부사항이 복잡합니다. 이로 인해 대규모 배치 크기 사용이 어렵고 학습 수렴 시간이 길어집니다.

Technical Solution

• 모델 로드 개수 감소: 4개(정책, 참조 정책, 보상, 가치) → 3개(정책, 참조 정책, 보상)로 축소
• 보상 할당 방식 변경: 토큰 단위 개별 행동 모델링에서 전체 완성도를 하나의 행동으로 모델링으로 전환
• REINFORCE 손실 함수 적용: 토큰별 advantage 계산 대신 전체 생성 완성도의 합산 reward에서 baseline을 뺀 값에 로그 확률을 곱하는 방식 도입
• 가치 모델 제거: GAE(Generalized Advantage Estimation) 기반 advantage 계산 제거로 구현 복잡도 감소
• 클리핑 메커니즘 개선: 임계값 통과 시 그래디언트 널링 이슈 완화

Impact

• GPU vRAM 사용량: 모델 크기별 50-70% 감소
• 학습 속도: 1B 모델 대비 2배, 6.9B 모델 대비 3배 개선
• 응답 품질: GPT-4 판정 기준 win rate에서 PPO와 경쟁 수준 달성
• DPO 대비 성능: 오프라인 방식인 인기 있는 DPO 방식을 일관되게 초과 달성

Key Takeaway

전체 완성도를 단일 행동으로 모델링하고 가치 모델을 제거함으로써 온라인 RLHF 훈련의 메모리 효율성과 속도를 극적으로 개선할 수 있으며, 이는 더 많은 연구가 오프라인 방식보다 온라인 RL의 효과성을 입증함에 따라 커뮤니티의 온라인 RL 탐색을 가능하게 합니다.