Google

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

The general objective function of Proximal Policy Optimization (PPO) can be specifically adapted for the training of Large Language Models. This involves formulating the optimization problem for LLMs within the PPO framework, which is a widely adopted approach in the field.

PPO Objective for LLM Training

The RLHF policy optimization objective can be reformulated by substituting the definition of the optimal reward-weighted policy, $$\pi^{*}(\mathbf{y}|\mathbf{x})$$. This transforms the objective into one that minimizes the Kullback-Leibler (KL) divergence between the learned policy, $$\pi_{\theta}(\cdot|\mathbf{x})$$, and the optimal policy. The substitution yields: $$ \tilde{\theta} = \arg \min_{\theta} \mathbb{E}_{\mathbf{x} \sim \mathcal{D}} \mathbb{E}_{\mathbf{y} \sim \pi_{\theta}(\cdot|\mathbf{x})} \left[ \log \pi_{\theta}(\mathbf{y}|\mathbf{x}) - \log \pi^{*}(\mathbf{y}|\mathbf{x}) - \log Z(\mathbf{x}) \right] $$ Because the expectation of the difference in log-probabilities over $$\pi_{\theta}$$ corresponds to the definition of KL divergence, this simplifies to: $$ \tilde{\theta} = \arg \min_{\theta} \mathbb{E}_{\mathbf{x} \sim \mathcal{D}} \left[ \text{KL}(\pi_{\theta}(\cdot|\mathbf{x}) || \pi^{*}(\cdot|\mathbf{x})) - \log Z(\mathbf{x}) \right] $$ This derivation shows that the objective is equivalent to minimizing the KL divergence from the optimal policy, offset by the normalization term $$\log Z(\mathbf{x})$$.

Derivation of the KL Divergence Objective for Policy Optimization

Consider the objective function used for policy optimization in reinforcement learning with human feedback. This function includes a term to maximize rewards and a penalty term scaled by a coefficient, β, to regulate divergence from a reference policy. Analyze the distinct consequences for the language model's generated outputs if the coefficient β is set to a very large value versus if it is set to zero. Explain the reasoning behind each outcome.

Analyzing the Trade-off in Policy Optimization

Analyze the most likely outcome of the experiment described in the case study. Explain why this outcome would occur by referencing the roles of the different components in the policy optimization objective.

Analyzing a Modified Policy Optimization Objective

During the policy optimization stage of training a large language model, an engineer observes that the model's outputs are coherent and safe, but they show very little improvement over the initial supervised fine-tuned version and consistently receive mediocre scores from the reward model. Which of the following is the most likely cause of this issue, based on the policy optimization objective function that balances maximizing rewards with a penalty for policy divergence?

Learn Before

Related