Google

The policy learning stage in RLHF is an iterative fine-tuning process. For each step, a prompt, $x$, is sampled from a dataset, $D$. The current language model, acting as the policy, then generates a corresponding output, $y$, by sampling from its probability distribution, $P_r(y|x)$. This input-output pair, \{x, y\}, is evaluated by the trained reward model, which assigns it a numerical reward score, $r(x, y)$. This score serves as the feedback signal for a reinforcement learning algorithm, which updates the policy's parameters to favor outputs that receive higher rewards.

Policy Learning in RLHF

The objective in the policy learning phase of Reinforcement Learning from Human Feedback (RLHF) is to find the optimal policy parameters, denoted as $\tilde{\theta}$, that maximize the expected reward. The optimization process starts with the parameters of a pre-trained model, $\hat{\theta}^{+}$, and seeks to maximize the reward assigned by a learned reward model, $R_{\hat{\omega}}$. The formal expression is: 

$$ \tilde{\theta} = \arg \max_{\hat{\theta}^{+}} \mathbb{E}_{(\mathbf{x}, \mathbf{y}_{\hat{\theta}^{+}}) \sim \mathcal{D}_{\text{rlft}}} R_{\hat{\omega}}(\mathbf{x}, \mathbf{y}_{\hat{\theta}^{+}}) $$ 

Here:
- $\tilde{\theta}$ are the optimized policy parameters.
- $\arg \max_{\hat{\theta}^{+}}$ indicates that we are searching for the parameters that maximize the objective, starting from the initial parameters $\hat{\theta}^{+}$.
- $\mathbb{E}_{(\mathbf{x}, \mathbf{y}_{\hat{\theta}^{+}}) \sim \mathcal{D}_{\text{rlft}}}$ represents the expected value over the dataset $\mathcal{D}_{\text{rlft}}$. For each input $\mathbf{x}$ from the dataset, a response $\mathbf{y}_{\hat{\theta}^{+}}$ is generated by the current policy.
- $R_{\hat{\omega}}(\mathbf{x}, \mathbf{y}_{\hat{\theta}^{+}})$ is the score assigned by the reward model (with parameters $\hat{\omega}$) to the generated response for the given input.

Objective Function for Policy Learning in RLHF

In practical applications of Reinforcement Learning from Human Feedback (RLHF), advanced algorithms like Proximal Policy Optimization (PPO) are frequently employed during the policy learning phase. The use of PPO helps to achieve more stable training and leads to better overall performance of the language model.

Use of Proximal Policy Optimization (PPO) in RLHF

The Advantage Actor-Critic (A2C) method is a specific reinforcement learning algorithm that can be utilized within the Reinforcement Learning from Human Feedback (RLHF) framework. Its application is aimed at fine-tuning Large Language Models to better align their outputs with human preferences.

Application of A2C in RLHF for LLM Alignment

In RLHF, the reference model, with parameters denoted by $\theta_{ref}$, serves as the baseline Large Language Model that provides the starting point for policy training. This model is typically a prior version of the LLM being trained or a model fine-tuned without human feedback, such as an SFT model. During the policy training phase, the reference model has two key functions: it is used to perform sampling across the range of possible outputs, and it is a component in the loss calculation, helping to regulate the policy updates.

Role and Definition of the Reference Model in RLHF

In the final stage of the RLHF process, the policy and value models undergo simultaneous training, guided by the previously trained reward model. This iterative update process occurs at each token position within a generated sequence. The value function's parameters are adjusted by minimizing the Mean Squared Error (MSE) of its predictions, while the policy is refined by minimizing the Proximal Policy Optimization (PPO) loss to encourage the generation of outputs that receive higher rewards.

Joint Optimization of Policy and Value Functions in RLHF

The goal of the policy training stage in Reinforcement Learning from Human Feedback (RLHF) is to find the optimal policy parameters $$\tilde{\theta}$$ that maximize expected reward without deviating too far from a reference policy. The training objective evaluates the quality of an output $$\mathbf{y}$$ given an input $$\mathbf{x}$$ using a reward model $$r(\mathbf{x},\mathbf{y})$$. The objective minimizes the negative reward (loss) and includes a penalty for policy divergence:

$$\tilde{\theta} = \arg \min_{\theta} \mathbb{E}_{\mathbf{x} \sim \mathcal{D}} \mathbb{E}_{\mathbf{y} \sim \pi_{\theta}(\cdot|\mathbf{x})} \big[ \underbrace{-r(\mathbf{x}, \mathbf{y})}_{\text{loss}} + \beta \underbrace{(\log \pi_{\theta}(\mathbf{y}|\mathbf{x}) - \log \pi_{\theta_{\mathrm{ref}}}(\mathbf{y}|\mathbf{x}))}_{\text{penalty}} \big]$$

Here, the penalty regularizes the current policy $$\pi_{\theta}$$ against the reference policy $$\pi_{\theta_{\mathrm{ref}}}$$ using a coefficient $$\beta$$.

RLHF Policy Optimization Objective

In Reinforcement Learning from Human Feedback (RLHF), the reference policy, denoted as $\pi_{\theta_{\text{ref}}}(\mathbf{y}|\mathbf{x})$, is a fixed policy used as a baseline during the optimization of the active policy $\pi_{\theta}$. It is typically a copy of the supervised fine-tuned (SFT) model before the RLHF stage begins. The reference policy's role is to prevent the active policy from deviating too far from the original language style and safety constraints, which is enforced by a penalty term (e.g., KL-divergence) that measures the difference between the two policies.

Reference Policy in RLHF

The objective of the reinforcement learning phase in RLHF is to minimize a loss function, formally expressed as `min L(x, {y1, y2}, r)`. This function is designed to optimize the language model's policy. The loss `L` is calculated using the input prompt `x`, a set of sampled outputs like `{y1, y2}`, and a reward model `r`. This reward model, which is pre-trained on human preference data, provides the critical feedback signal within the loss function, guiding the policy towards generating responses that align with human preferences.

RLHF Policy Optimization as Loss Minimization

A language model is being fine-tuned using an iterative feedback process. In each step, the model generates a response to a prompt. A separate, pre-trained scoring model then assigns a numerical score to this response based on its quality. What is the most direct and immediate use of this numerical score within a single step of this training loop?

Arrange the following events into the correct chronological order as they would occur within a single iterative step of the policy learning phase for a language model.

Based on the described iterative training loop, which component is the most likely source of this unintended behavior, and why?

Diagnosing a Training Failure in an Iterative Fine-Tuning Process

Direct Preference Optimization (DPO) is an alignment method that simplifies the training framework by eliminating the need to explicitly model rewards. Instead of developing a separate reward model—which can be difficult to train reliably and negatively impact policy learning if poorly trained—DPO directly optimizes the language model's policy based on human preferences. By doing so, it achieves human preference alignment in a straightforward, supervised learning-like fashion.

Learn Before

Related