Designing Reward Functions from AI Preferences

The process of designing a reward function in AI alignment involves translating pairwise AI preferences, typically represented by a model like $p_\theta(y_w \succ y_l | x)$ that discerns preferred responses from AI-generated labels, into a usable scalar reward signal, $r(x, y)$. This reward function serves as a foundation for the subsequent reinforcement learning phase, which guides the LLM policy towards generating outputs that align with the learned preferences.

From Pairwise Preferences to Scalar Rewards

The fundamental assumption, inherited from RLHF, is that the preference model $p_\theta$ approximates an underlying latent reward function $r_\phi(x, y)$. A common modeling choice is the Bradley-Terry model, which posits that the probability of preferring response $y_w$ over $y_l$ given prompt $x$ can be expressed as:

$$p(y_w \succ y_l | x) = \frac{\exp(r_\phi(x, y_w))}{\exp(r_\phi(x, y_w)) + \exp(r_\phi(x, y_l))} = \sigma(r_\phi(x, y_w) - r_\phi(x, y_l))$$

Here, $\sigma(\cdot)$ is the sigmoid function. During preference model training, we fit $p_\theta$ to the AI-generated preference data $(x, y_w, y_l)$, effectively learning an approximation of the difference in rewards between pairs of responses.

Our goal now is to extract a scalar reward function $r(x, y)$ that reflects the quality of a single response $y$ given prompt $x$, based on the trained preference model $p_\theta$. Since $p_\theta$ models the difference in rewards, the absolute scale of the extracted reward function $r(x, y)$ is somewhat arbitrary, but its relative values should reflect the preferences learned by $p_\theta$.

Extracting the Reward Signal

A prevalent method to define the reward signal $r(x, y)$ involves leveraging the internal computations of the trained preference model $p_\theta$. Many preference model architectures compute an internal scalar score, let's call it $s_\theta(x, y)$, for each input response $y$. The preference probability is then calculated based on the difference between the scores of the winning and losing responses:

$$p_\theta(y_w \succ y_l | x) = \sigma(s_\theta(x, y_w) - s_\theta(x, y_l))$$

Given this structure, a natural choice for the reward function is to directly use this internal score:

$$r(x, y) = s_\theta(x, y)$$

This score $s_\theta(x, y)$ represents the learned quality of response $y$ according to the AI preference labeler. It directly captures the information optimized during the preference model training.

Alternative Formulation:

In some implementations, the reward function might be defined using the logarithm of the sigmoid applied to the score:

$$r(x, y) = \beta \log(\sigma(s_\theta(x, y)))$$

The scaling factor $\beta$ is often introduced as a hyperparameter to control the magnitude of the rewards, which can significantly impact the RL training dynamics. This formulation emphasizes deviations from a neutral score (often zero).

Practical Notes for Reward Function Design

Simply extracting a score is often insufficient for stable and effective RL training. Several practical considerations arise:

1. Reward Scaling and Normalization: The raw scores $s_\theta(x, y)$ might have an arbitrary scale or offset. Large reward values can lead to excessively large policy updates and instability in algorithms like PPO. Conversely, very small rewards might result in slow learning. Common techniques include:

   • Whitening: Normalize rewards within each batch to have zero mean and unit variance. This helps stabilize PPO updates by ensuring the reward scale is consistent across different batches and prompts.
   • Clipping: Limit reward values to a predefined range to prevent extreme updates.
   • Scaling Factor ($\beta$): Explicitly scaling the reward signal, as seen in the alternative formulation, provides direct control over its influence relative to other components of the RL objective, such as the KL divergence penalty.

2. KL Divergence Penalty: To prevent the RL policy $\pi_{RL}$ from deviating too drastically from the original supervised fine-tuned (SFT) policy $\pi_{SFT}$ (or the policy trained after the CAI supervised phase), PPO typically incorporates a KL divergence penalty. The final reward signal used in the PPO update often takes the form:

   $$r_{final}(x, y) = r(x, y) - \beta_{KL} \text{KL}(\pi_{RL}(\cdot|x) || \pi_{SFT}(\cdot|x))$$

   Here, $r(x, y)$ is the reward derived from the AI preference model, and $\beta_{KL}$ controls the strength of the penalty for deviating from the reference policy. Designing the reward function $r(x, y)$ must consider its interaction with this KL term. The scale of $r(x, y)$ relative to the KL penalty is a sensitive hyperparameter.

3. Reward Model Drift: If the AI preference model itself is updated or changed during the RL training process (less common, but possible in advanced setups), the reward function becomes non-stationary. This can complicate RL convergence and requires careful handling.

4. Bias Amplification: The AI preference model might inherit or even amplify biases present in the model used to generate the labels (e.g., the constitution-based critiquer or another LLM). The resulting reward function will encode these biases. It's important to be aware that the reward function optimizes for what the AI perceives as good, which might not perfectly align with the intended constitutional principles or broader safety goals.

Integrating the Reward into the RL Algorithm

The designed reward function $r(x,y)$ serves as the primary signal guiding the optimization of the LLM policy $\pi_{RL}$ using PPO (or another suitable RL algorithm). During the PPO rollout phase, for each prompt $x$ sampled from a dataset, the current policy $\pi_{RL}$ generates a response $y$. This response $y$ is then fed into the frozen, trained AI preference model to compute the reward $r(x,y)$. This reward, potentially combined with the KL penalty, is used to calculate the advantages and update the policy parameters.

Flow diagram illustrating the generation and use of the reward signal within the RLAIF PPO loop. The frozen AI preference model computes a score for the generated response, which is then transformed into the reward signal used for the policy update.

Designing an effective reward function from AI preferences is a careful process. It requires not only understanding the theoretical link between preference probabilities and scalar rewards but also careful implementation regarding scaling, normalization, and integration with the chosen RL algorithm's objective function to ensure stable and efficient policy optimization. The quality of the reward function directly depends on the quality and consistency of the upstream AI preference labeler and the preference model trained upon its outputs.