ReMax: A Simple, Effective, and Efficient Reinforcement Learning Method for Aligning Large Language Models

Published 16 Oct 2023 in cs.LG | (2310.10505v4)

Abstract: Reinforcement Learning from Human Feedback (RLHF) is key to aligning 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.

Abstract PDF HTML Upgrade to Chat

References (46)

Citations (21)

View on Semantic Scholar

Summary

The paper introduces ReMax, a reinforcement learning method that eliminates the value model to simplify LLM alignment and reduce computational overhead.
It employs variance reduction techniques from REINFORCE to stabilize training and cut GPU memory usage by nearly 46% on 7B models.
Empirical tests on Mistral-7B show that ReMax achieves competitive, state-of-the-art performance on benchmarks while improving training efficiency.

ReMax: A Simple, Effective, and Efficient Reinforcement Learning Method for Aligning LLMs

Introduction

"ReMax: A Simple, Effective, and Efficient Reinforcement Learning Method for Aligning LLMs" (2310.10505) introduces ReMax, a method designed to align LLMs more efficiently compared to the widely used Proximal Policy Optimization (PPO). The core idea exploits the intrinsic properties of Reinforcement Learning from Human Feedback (RLHF) for LLM training: fast simulation, deterministic transitions, and trajectory-level rewards.

Architectural Innovations

ReMax differentiates itself from PPO by eliminating the value model and instead relies on variance reduction techniques derived from the REINFORCE algorithm. This approach simplifies the implementation, substantially reduces computational overhead, and maintains the alignment efficacy in large-scale models.

ReMax Design:

Reference Model: Utilizes a reference model for KL-divergence penalties, similar to PPO.
Without Value Model: ReMax eliminates the need for a separate value model, reducing GPU memory requirements.
Figure 1: Building blocks of PPO and ReMax. ReMax keeps the reference model (for calculating KL penalty) and removes all the components related to the value model in PPO.

Computational Efficiency

ReMax is designed to address the demanding GPU resources and computation time traditionally associated with PPO. For instance, ReMax cuts down GPU memory consumption by approximately 46% when training a 7B model, and avoids the need for memory-saving offloading techniques.

Figure 2: GPU memory consumption and training time by PPO and ReMax, respectively. These measurements are conducted on a Llama-2-7B model using A800-80GB GPUs.

Reinforcement Learning Adaptation

ReMax employs a reward-weighted likelihood maximization approach, rooted in the REINFORCE algorithm. This avoids the complexity inherent in PPO's advantage estimation and off-policy corrections. Instead, ReMax introduces a novel variance reduction technique that leverages the deterministic environment of LLMs under RLHF.

Key Properties:

Fast Simulation: LL models quickly generate complete responses, optimizing turnaround time.
Deterministic Transitions: LLM environments are predictable, reducing noise in training.
Trajectory-level Rewards: Focus on complete output evaluation simplifies the reward structure.
Figure 3: Illustration of StarCraft II (a general RL task example) and RLHF in LLMs.

Variance Reduction and Stability

The variance resulting from stochastic gradients in REINFORCE is mitigated in ReMax by incorporating a greedy baseline reward. This method maintains unbiased gradient estimation while reducing variability, which is pivotal in stable training of large-scale models.

Algorithm Implementation:

def remax_align(language_model, reward_model, dataset):
    for prompt in dataset:
        seq = language_model.sample(prompt, greedy=False)
        seq_max = language_model.sample(prompt, greedy=True)
        rewd_diff = reward_model.evaluate(seq) - reward_model.evaluate(seq_max)
        log_prob = language_model.logits(prompt, seq)
        loss = -torch.mean(torch.sum(log_prob, dim=-1) * rewd_diff)
        language_model.optimize(loss)
    return language_model

Practical Implications

In empirical tests, ReMax demonstrated competitive or superior performance compared to PPO, thanks to more consistent hyperparameter tuning and training stability. When applied to Mistral-7B, ReMax set a new state-of-the-art (SOTA) for open-source 7B models on the AlpacaEval and MT-bench leaderboards.

Conclusion

ReMax represents a pragmatic evolution in aligning LLMs through RLHF, reducing resource burden and improving accessibility to training large models. This methodological approach may encourage widespread adoption and refinement, given its balance of simplicity, efficiency, and effectiveness in real-world applications.