[논문 리뷰] ReMax: A Simple, Effective, and Efficient Reinforcement Learning Method for Aligning Large Language Models

Ziniu Li, Xu Tian외 5명|arXiv (Cornell University)|2023. 10. 16.

Topic Modeling인용 수 6

한 줄 요약

ReMax는 대규모 언어 모델의 정렬을 위한 새로운 강화학습 방법으로, PPO에서 사용하는 가치 모델을 제거함으로써 메모리 사용량을 50% 감소시키고, 반복당 훈련 속도를 2배로 높이며, 단 6줄의 코드로 구현이 가능해지면서도 성능을 유지하거나 향상시킨다. 이는 빠른 시뮬레이션, 결정론적 전이, 궤적 수준의 보상 등을 활용한다.

ABSTRACT

Reinforcement Learning from Human Feedback (RLHF) is key to aligning Large Language Models (LLMs), typically paired with the Proximal Policy Optimization (PPO) algorithm. While PPO is a powerful method designed for general reinforcement learning tasks, it is overly sophisticated for LLMs, leading to laborious hyper-parameter tuning and significant computation burdens. To make RLHF efficient, we present ReMax, which leverages 3 properties of RLHF: fast simulation, deterministic transitions, and trajectory-level rewards. These properties are not exploited in PPO, making it less suitable for RLHF. Building on the renowned REINFORCE algorithm, ReMax does not require training an additional value model as in PPO and is further enhanced with a new variance reduction technique. ReMax offers several benefits over PPO: it is simpler to implement, eliminates more than 4 hyper-parameters in PPO, reduces GPU memory usage, and shortens training time. ReMax can save about 46% GPU memory than PPO when training a 7B model and enables training on A800-80GB GPUs without the memory-saving offloading technique needed by PPO. Applying ReMax to a Mistral-7B model resulted in a 94.78% win rate on the AlpacaEval leaderboard and a 7.739 score on MT-bench, setting a new SOTA for open-source 7B models. These results show the effectiveness of ReMax while addressing the limitations of PPO in LLMs.

연구 동기 및 목표

• 대규모 언어 모델의 인간 피드백을 통한 강화학습(RLHF)에서 Proximal Policy Optimization(PPO)의 높은 계산 비용과 메모리 오버헤드 문제를 해결하기 위해.
• RLHF 작업에서 미사용된 특성들—빠른 시뮬레이션, 결정론적 전이, 궤적 수준의 보상—을 식별하여 간단하고 효율적인 알고리즘 설계에 활용하기 위해.
• 성능을 훼손하지 않고 가치 모델을 제거함으로써 복잡성과 자원 요구량을 줄이는 새로운 강화학습 알고리즘을 설계하기 위해.
• 단순화된 분산 감소가 있는 REINFORCE 기반 방법이 PPO를 능가하는 효율성과 정렬 품질을 동시에 확보할 수 있는지 증명하기 위해.

제안 방법

• ReMax는 REINFORCE 알고리즘을 기반으로 하지만, 대규모 언어 모델에 특화된 새로운 분산 감소 기법을 통합하여 개선되었다.
• 가치 모델을 완전히 제거함으로써 추가 신경망을 훈련, 저장, 역전파할 필요가 없어진다.
• 각 프롬프트당 두 번의 응답 생성을 사용한다: 하나는 정책 기울기 계산을 위해, 다른 하나는 분산 감소를 위해 활용하며, 토큰 생성의 결정론적 성격을 활용한다.
• LLM 응답이 생성이 빠르며, 과거 토큰에 기반해 전이가 결정론적이고, 보상은 시퀀스의 끝에서만 제공된다는 사실을 활용한다.
• 알고리즘은 단 6줄의 코드로 구현되어 PPO에 비해 초모수 조정이 크게 줄어든다.
• 메모리 절감은 가치 모델의 파라미터, 기울기, 최적화기 상태를 제거함으로써 발생하며, 이로 인해 약 50%의 메모리 사용량 감소를 달성한다.

실험 결과

연구 질문

• RQ1 성능을 유지하거나 향상시키면서도 가치 모델이 없는 강화학습 알고리즘을 LLM 정렬에 적용할 수 있는가?
• RQ2 RLHF에서의 빠른 시뮬레이션, 결정론적 전이, 궤적 수준의 보상 특성을 얼마나 잘 활용할 수 있는가? 이는 훈련을 단순화하고 가속화하는 데 기여하는가?
• RQ3 가치 모델을 제거하면 정렬 품질을 훼손하지 않으면서도 메모리 효율성과 훈련 속도 향상에 측정 가능한 기여를 할 수 있는가?
• RQ4 분산 감소 기법이 적용된 단순화된 REINFORCE 기반 방법이 PPO에 비해 LLM 정렬에서 효율성과 효과성 면에서 뛰어나게 될 수 있는가?

주요 결과

• ReMax는 PPO 대비 약 50%의 메모리 사용량 감소를 기록하여, PPO가 오프로딩 없이 실패하는 상황에서도 Llama2-7B를 A100-80GB GPU에서 훈련 가능하게 한다.
• 오프로딩을 적용한 경우, 4개의 A100-80GB GPU에서 ReMax는 최대 배치 크기 160을 지원하며, PPO는 120에 그친다. 이는 더 높은 처리량을 가능하게 한다.
• 단일 GPU에서 오프로딩을 적용한 경우, ReMax는 배치 크기 42를 지원하지만, PPO는 32에 그치며, 오프로딩 없이선 두 모델 모두 실패한다.
• 4개의 GPU에서 오프로딩이 필요 없어지면서 ReMax는 훈련 시간을 2.8시간으로 단축시켜 총 훈련 시간을 10.2시간에서 2.8시간으로 줄였고, 이는 2.5배의 속도 향상이다.
• 반복당 훈련 시간은 ReMax가 약 2.8시간으로 PPO의 6.0시간(오프로딩 적용 시)에 비해 절반 수준이며, 4개의 GPU에서 동일한 조건에서 측정되었다.
• 간단한 구조임에도 불구하고 ReMax는 PPO와 비교해 성능 저하 없이 강력한 정렬 성능을 유지한다.