Solving Non-Differentiability of Human Feedback with Proximal Policy Optimization

Introduction

The advent of Reinforcement Learning with Human Feedback (RLHF) marks a significant development in machine learning, particularly for generative models such as large language models (LLMs) like GPT-3. This approach is guided by the principle of Proximal Policy Optimization (PPO), a potent policy optimization method that has demonstrated its effectiveness in training these models.

This article explores PPO in-depth, unraveling its mathematical foundations and elucidating its central role in the RLHF framework. Moreover, we delve into the application of PPO in the training and fine-tuning of large language models, highlighting its impact on shaping these advanced AI systems.

One of the critical aspects we tackle is the notion of Responsible AI, a theme that permeates every aspect of these technologies. As we delve into the intricacies of PPO and RLHF, we also underscore the importance of building AI systems that are not only intelligent and capable but also ethical and responsible. This means ensuring our models adhere to fairness, transparency, interpretability, and robustness standards.

Know Thy Data!

As we continue to push the boundaries of what AI can achieve, we must not lose sight of these values. Therefore, this exploration is not just a deep dive into the nuts and bolts of AI systems but also examines the principles that guide their development, ensuring that we are building beneficial AI for all.

Differentiability in Transformer Models

Transformers, which form the foundation of large language models like GPT-3, depend heavily on the principle of differentiability. Each layer within a transformer model comprises operations such as matrix multiplication and the calculation of attention, all of which are differentiable. This differentiability is critical for backpropagation during model training, which is a process that adjusts the model's parameters to minimize the discrepancy between the model's output and the actual target.

The architecture of the transformer model is engineered to account for dependencies in input data, regardless of the distance between the individual data elements in the sequence. This trait makes them particularly effective for processing natural language data. The self-attention mechanism, a key feature of transformers, computes a weighted sum of all input elements, expressed as:

Self-Attention(Q,K,V) = softmax(QK^T/√d_k)V

Here, Q, K, and V are the query, key, and value vectors, respectively, which are computed from the input. d_k is the dimensionality of the key vectors. The softmax function is applied to the dot product of the query and key vectors, scaled by the square root of d_k. This operation results in attention scores used to weight the value vectors. These operations are all differentiable, which allows the model to adjust its attention weights during training through gradient descent (Vaswani et al., 2017).

Differentiability and Non-Differentiability in Reinforcement Learning with Human Feedback

In RLHF, the idea is to use human feedback as the reward signal to guide the learning process. However, human feedback is inherently non-differentiable, posing a unique challenge to standard RL algorithms (Hu et al., 2017).

To tackle this, RLHF uses a differentiable reward model. This reward model is trained to predict human feedback and provides an approximation of the human feedback that can be differentiated, allowing the RL algorithm, such as PPO, to compute gradients and update the policy.

Differentiability: The Foundation of Learning

Differentiability is a fundamental concept in machine learning, particularly in training neural networks. The requirement for differentiability stems from the need to compute gradients, or derivatives, of the loss function with respect to the model parameters. This process uses the backpropagation algorithm, which leverages the chain rule from calculus to compute these gradients efficiently.

In mathematical terms, the gradient of a function f at a point x (denoted by ?f(x)) is a vector that points in the direction of the greatest increase of f, and its magnitude is the rate of increase in that direction. For a neural network with a loss function L and parameters θ, we compute the gradient of L with respect to θ, denoted by ?θL, to know how to change θ to decrease L.

The backpropagation algorithm can be broken down into two steps. First, we perform a forward pass to compute the output and loss. Then we perform a backward pass to compute the gradients. The backward pass applies the chain rule of calculus, which in its simplest form states that if y = f(g(x)) for some functions f and g, then the derivative of y with respect to x is dy/dx = (dy/dg) * (dg/dx).

In the context of reinforcement learning and policy optimization, differentiability plays a crucial role in policy gradient methods, including PPO. The policy, typically parameterized by a neural network, needs to be differentiable with respect to its parameters so that gradients can be computed for policy updates.

In PPO, the policy π is represented as a neural network with parameters θ, and the objective function is a function of π. The gradient of the objective function with respect to θ, denoted by ?θJ(π), is computed to update the policy.

This leads us to the policy gradient theorem, which provides a way to estimate the gradient of the expected cumulative reward with respect to the policy parameters without needing a model of the environment. The theorem states that:

?θJ(π) = E[?θlogπ(a_t|s_t) * A_t],

Where A_t is the advantage function, a measure of how much better an action a_t is compared to the average action at state s_t.

The advantage function itself can be estimated in several ways. One common method is by subtracting a baseline (e.g., the value function V(s)) from the Q-value (the expected return of taking action a, in state s):

A_t = Q(s_t, a_t) - V(s_t).

The formula for the policy gradient derived from the policy gradient theorem tells us how to adjust the parameters θ to increase the expected cumulative reward. It essentially weighs the log probability of each action by its advantage and then averages this over many timesteps. Doing so encourages actions that yield a higher-than-average reward and discourages actions that yield a lower-than-average reward.

Policy in the Context of Reinforcement Learning with Human Feedback (RLHF)

In the context of Reinforcement Learning with Human Feedback (RLHF), a policy is a strategy the learning agent employs to determine the next action based on its current state. It is a mapping from the state of the environment to actions that the agent can take.

In more technical terms, a policy is a function π(a|s) that defines the probability distribution over actions 'a' given a state 's'. In the case of deterministic policies, the function directly provides the action to be taken for a given state. However, in most practical cases, policies are stochastic, providing a probability distribution over actions.

In RLHF, the policy is typically represented by a parametrized function, such as a neural network, where the parameters θ are learned through interaction with the environment and feedback. RLHF aims to find the optimal parameters that maximize the expected cumulative reward over time. This involves taking actions, receiving feedback, and updating the policy based on the feedback to improve future actions.

Proximal Policy Optimization: A Primer

Policy optimization methods aim to find the best policy—a mapping from states to actions—that maximizes the expected cumulative reward in a reinforcement learning setting. PPO, introduced by Schulman et al. (2017), is a policy optimization method that introduces a novel objective function designed to limit the policy update at each step to ensure learning stability.?

Mathematically, this can be expressed as an optimization problem where the objective is to maximize the expectation of a certain function of the state-action pairs sampled from trajectories generated under the current policy.

This can be represented as the following equation:

Maximize π E[Σ(γ^t * r(s_t, a_t))],

In this equation, E denotes the expected value, Σ denotes summation over the trajectory τ (which is (s_0, a_0, s_1, a_1, ..., s_T, a_T)), π is the policy, γ is the discount factor, and r(s_t, a_t) is the reward function.

While traditional policy gradients methods such as Vanilla Policy Gradient (VPG) or Actor-Critic methods (A2C, A3C) suffer from issues like high variance and potential instability, PPO addresses these by constraining the updates to the policy to be "close" to the previous policy. In mathematical terms, PPO modifies the objective function to include a penalty term that discourages large changes in the policy.?

This penalty term is a function of the Kullback-Leibler (KL) divergence between the new and old policies, measuring the difference between two probability distributions. The objective function of PPO can be written as follows:

Maximize π E[Σ(γ^t * r(s_t, a_t)) - β * KL(π_old || π)],

In this equation, β is a hyperparameter that controls the trade-off between maximizing the expected cumulative reward and minimizing the policy change.

This approach makes PPO more sample-efficient and stable than traditional policy gradient methods and more straightforward to implement than second-order methods like Natural Policy Gradient (NPG) or Trust Region Policy Optimization (TRPO). NPG and TRPO also aim to constrain policy updates. However, they do so in a more complex way that requires second-order optimization, which can be computationally expensive and difficult to implement. In contrast, PPO achieves similar performance benefits with a simpler first-order optimization method, making it more accessible for various applications.

Kullback-Leibler (KL) Divergence In Detail

The Kullback-Leibler (KL) divergence quantifies the difference between two probability distributions (Kullback & Leibler, 1951). Used in machine learning and statistics, it aids model selection and information retrieval (Cover & Thomas, 2006).

For discrete distributions, it has defined as?D_KL(P||Q) = sum[P(i) * log(P(i) / Q(i))]?over all i, and for continuous distributions, it is?D_KL(P||Q) = integral[P(x) * log(P(x) / Q(x)) dx]?over all x (Bishop, 2006).

Critical properties of KL divergence are non-symmetry and non-negativity. It is not a valid metric (Cover & Thomas, 2006). It measures information loss when Q approximates P (Kullback & Leibler, 1951).

Policy Optimization Methods: An Evolution

Before PPO, policy optimization methods like Vanilla Policy Gradient (VPG), Trust Region Policy Optimization (TRPO), and Actor-Critic methods were used to update policies (Schulman et al., 2015; Lillicrap et al., 2015; Sutton et al., 1999). While these methods have been successful, they have some limitations.

VPG suffers from high variance in gradient estimates, leading to unstable learning (Schulman et al., 2015). TRPO addresses this by constraining the policy update step to be within a 'trust region' to ensure stability, but it is complicated to implement due to its second-order nature (Schulman et al., 2015). Actor-Critic methods, like Advantage Actor-Critic (A2C), use a critic to reduce variance in gradient estimates but still can suffer from unstable learning (Sutton et al., 1999).

PPO emerged as a solution to these problems, bringing the stability of TRPO and the simplicity of VPG (Schulman et al., 2015). It constrains the policy update to ensure that the new policy does not deviate drastically from the old policy, thereby maintaining stable learning (Schulman et al., 2015).

Proximal Policy Optimization: How It Works

Proximal Policy Optimization (PPO) introduces a novel objective function, often referred to as the surrogate objective, to constrain policy updates (Schulman et al., 2015). The idea behind PPO is to avoid excessively large policy updates that could potentially harm the learning process (Schulman et al., 2015). This is achieved by adding a penalty term to the objective function, discouraging a large deviation from the old policy during an update (Schulman et al., 2015).

Mathematically, the PPO objective function is expressed as follows:

L(θ) = E_t[min(r_t(θ)A_t, clip(r_t(θ), 1-ε, 1+ε)A_t)].

Here, θ represents the parameters of the policy, A_t is the advantage function at time t, and r_t(θ) is the ratio of the new policy probability to the old policy probability for action a_t at state s_t, given by π_θ(a_t|s_t) / π_θ_old(a_t|s_t) (Schulman et al., 2015). The advantage function, A_t, measures how much better the chosen action is compared to the average action for that state (Schulman et al., 2015).

The objective function includes two terms inside the min function: the first term, r_t(θ)A_t, encourages improving the policy by increasing the probability of actions with positive advantage and decreasing the probability of actions with negative advantage (Schulman et al., 2015). The second term, clip(r_t(θ), 1-ε, 1+ε)A_t, is designed to discourage large policy updates, where ε is a hyperparameter usually set to a small value like 0.1 or 0.2 (Schulman et al., 2015). The clip function limits the r_t(θ) value to the range [1-ε, 1+ε]. This means that even if an action has a very high advantage, the policy update is constrained to prevent it from becoming too large (Schulman et al., 2015).

The expectation E_t[.] is taken over the timesteps t, and the objective is to maximize this function with respect to the policy parameters θ. The resulting update rule encourages improving the policy while keeping the policy updates relatively small to ensure stable learning (Schulman et al., 2015).

Comparing PPO with Other Policy Optimization Methods

PPO presents several advantages over other policy optimization methods. Compared to vanilla policy gradient methods, PPO reduces the variance of policy updates, leading to more stable learning. It performs similarly to TRPO and NPG but with less computational complexity and easier implementation (Schulman et al., 2017).

Adapting PPO for RLHF in Fine-Tuning LLMs

Applying PPO to RLHF in fine-tuning LLMs presents an innovative approach to the non-differentiability of human feedback. It employs a reward model that furnishes a differentiable approximation of human feedback, bridging the gap between non-differentiable real-world feedback and the requirements of gradient-based learning methods (Christiano et al., 2017).

In the RLHF framework, a reward model is trained to predict the feedback a human demonstrator provides. This reward model is typically a neural network that takes as input the state-action pair and outputs a scalar reward. The training is performed by collecting a dataset of state-action pairs and corresponding human feedback, which serves as the ground truth during the training of the reward model.

The PPO algorithm then uses this reward model to fine-tune the LLM. The process follows these general steps:

1. Given a context, the LLM predicts the log-likelihood probabilities of the next token in the sequence.
2. The reward model assigns a reward to each state-action pair in the generated sequence, where an "action" is the model's choice of the next token based on the predicted probabilities.
3. PPO uses these rewards to update the policy, aiming to produce sequences that would get higher rewards from the reward model.

The integration of PPO and RLHF in fine-tuning LLMs is an active field of research. Ongoing challenges include

• developing more accurate reward models,
• handling the non-differentiability inherent in human feedback, and
• the exploration-exploitation trade-off in training with human feedback.

The resolution of these challenges through future research is anticipated to further enhance the quality and capabilities of large language models, leading to more contextually accurate AI systems that are user-aligned and effective in diverse applications (OpenAI, 2021).

Concluding Remarks

Through this in-depth analysis of PPO, differentiability, and RLHF, we have highlighted the intricate dynamics among these elements in training and fine-tuning large language models. The confluence of these techniques forms the backbone of current state-of-the-art AI systems, underlining their importance and influence in the field of generative AI.

As our knowledge of these methodologies continues to expand and evolve, we gain a deeper understanding of the principles that guide them and a more nuanced appreciation of their complexities and challenges. This insight is essential in our quest to harness the full potential of these models and to contribute meaningfully to the development of AI systems.

Our goal extends beyond simply improving the capabilities of these systems. We also strive to create more contextually accurate AI, better aligned with human values, and more effectively able to interact with and understand the complexities of the world around it. As we move forward, this dual focus on capability and alignment will continue to shape the evolution of AI, driving us toward creating systems that are more powerful and more attuned to the needs and values of the societies they serve.

References

• Bishop, C. M. (2006).?Pattern Recognition and Machine Learning. Springer.
• Christiano, P. F., Leike, J., Brown, T., Martic, M., Legg, S., & Amodei, D. (2017). Deep reinforcement learning from human preferences. In Advances in Neural Information Processing Systems (pp. 4302-4310).
• Cover, T. M., & Thomas, J. A. (2006). Elements of Information Theory (2nd ed.). Wiley-Interscience.
• Hu, Z., Yang, Z., Liang, X., Salakhutdinov, R., & Xing, E. P. (2017, July). Toward controlled generation of text. In International Conference on Machine Learning (pp. 1587-1596). PMLR.
• Kullback, S., & Leibler, R. A. (1951). On Information and Sufficiency. Annals of Mathematical Statistics, 22(1), 79-86.
• Lillicrap, T. P., Hunt, J. J., Pritzel, A., Heess, N., Erez, T., Tassa, Y., ... & Wierstra, D. (2015). Continuous control with deep reinforcement learning. arXiv preprint arXiv:1509.02971.
• OpenAI. (2021). Fine-Tuning Large Language Models: Challenges and Directions.?https://openai.com/blog/fine-tuning/
• Schulman, J., Levine, S., Moritz, P., Jordan, M. I., & Abbeel, P. (2015). Trust Region Policy Optimization. In Proceedings of the 32nd International Conference on Machine Learning (ICML) (pp. 1889-1897).
• Schulman, J., Wolski, F., Dhariwal, P., Radford, A., & Klimov, O. (2017). Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347.
• Sutton, R. S., McAllester, D. A., Singh, S. P., & Mansour, Y. (1999). Policy gradient methods for reinforcement learning with function approximation. In Advances in Neural Information Processing Systems (NIPS) (pp. 1057–1063).
• Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., ... & Polosukhin, I. (2017). Attention is all you need. In Advances in neural information processing systems (pp. 5998-6008).