RLHF: Reinforcement Learning from Human Feedback

RLHF proceeds in three sequential stages: supervised fine-tuning (SFT) to create a well-behaved starting policy, reward model (RM) training on human preference data, and RL fine-tuning of the policy against the RM. Each stage builds on the last, and each has its own loss function and optimization dynamics.

Stage 1: Supervised Fine-Tuning

Starting from a pretrained base model $\pi_\text{base}$, we collect a dataset of (prompt, high-quality response) pairs — typically written or curated by human contractors — and fine-tune with standard cross-entropy:

Explanation: This is the standard next-token prediction loss applied only to the response tokens $y$, conditioning on prompt $x$, pushing the model to imitate demonstrated behavior.

Symbols:

• $\theta$: policy (LLM) parameters
• $x$: input prompt
• $y$: target response tokens, $y_t$ is the $t$-th token
• $\mathcal{D}_\text{SFT}$: supervised fine-tuning dataset of curated (prompt, response) pairs

This produces a policy $\pi_\text{SFT}$ that is better-behaved than the base model and serves as the initialization for both the reward model and the RL stage.

Stage 2: Reward Model Training

Human raters are shown a prompt $x$ and two candidate completions $(y_w, y_l)$ — one preferred ("winner") and one not ("loser") — and select which is better. The reward model $r_\phi$ is trained to assign higher scalar scores to preferred completions using the Bradley-Terry preference model:

Explanation: This is a binary cross-entropy loss that encourages the reward model to assign a higher score to the preferred response $y_w$ than the dispreferred response $y_l$, modeled as a logistic probability over the score difference.

Symbols:

• $\phi$: reward model parameters
• $r_\phi(x, y) \in \mathbb{R}$: scalar reward assigned to response $y$ given prompt $x$
• $\sigma(\cdot)$: sigmoid function
• $y_w$: human-preferred ("winning") response
• $y_l$: human-dispreferred ("losing") response
• $\mathcal{D}_\text{pref}$: dataset of human preference comparisons

In practice, the reward model is initialized from $\pi_\text{SFT}$ with the final language model head replaced by a linear layer outputting a single scalar. This weight initialization means the RM starts with strong language understanding.

Stage 3: RL Fine-Tuning with PPO

The policy $\pi_\theta$ is now fine-tuned to maximize the expected reward from $r_\phi$, subject to a KL-divergence penalty that prevents the policy from drifting too far from $\pi_\text{SFT}$ (which would cause reward hacking — producing outputs the RM scores highly but humans would not):

Explanation: The policy is rewarded for producing high-scoring responses according to $r_\phi$, but penalized for diverging from the SFT policy — the $\beta$ coefficient controls this alignment-capability trade-off.

Symbols:

• $\mathcal{J}_\text{RLHF}$: the RL objective to maximize
• $r_\phi(x, y)$: scalar reward from the frozen reward model
• $\beta > 0$: KL penalty coefficient (typically 0.01–0.1)
• $D_\text{KL}(\pi_\theta \| \pi_\text{SFT})$: per-token KL divergence between the updated policy and the SFT reference
• $\mathcal{D}$: prompt distribution (often separate from SFT/pref data)

The KL term can be decomposed per token, making it compatible with autoregressive generation. The full per-token reward used in PPO is:

Explanation: The terminal reward from the reward model is applied only at the last token $T$, while the KL penalty against the reference policy is applied at every token — this is the standard "reward shaping" used in InstructGPT and subsequent systems. The PPO clipped surrogate objective is then applied to update $\theta$:

Explanation: PPO clips the probability ratio $\rho_t = \pi_\theta / \pi_\text{old}$ so that policy updates are not too large in a single step, where $\hat{A}_t$ is the estimated advantage (how much better this action was than the baseline value function predicted).

Symbols:

• $\rho_t(\theta) = \pi_\theta(y_t \mid x, y_{<t}) / \pi_{\theta_\text{old}}(y_t \mid x, y_{<t})$: probability ratio (importance weight)
• $\hat{A}_t$: advantage estimate from a learned value function (critic), typically computed via GAE
• $\epsilon$: PPO clip parameter (typically 0.1–0.2)

```python import torch import torch.nn.functional as F def reward_model_loss( r_preferred: torch.Tensor, # shape: (B,) r_rejected: torch.Tensor, # shape: (B,) ) -> torch.Tensor: """Bradley-Terry preference loss for reward model training.""" return -F.logsigmoid(r_preferred - r_rejected).mean()

def kl_penalty_per_token( log_probs_policy: torch.Tensor, # (B, T) log_probs_ref: torch.Tensor, # (B, T) beta: float = 0.05, ) -> torch.Tensor: """Per-token KL penalty: beta (log π_θ - log π_SFT).""" return beta (log_probs_policy - log_probs_ref)

def compute_rlhf_rewards( scalar_reward: torch.Tensor, # (B,) — from reward model, at EOS log_probs_policy: torch.Tensor, # (B, T) log_probs_ref: torch.Tensor, # (B, T) beta: float = 0.05, ) -> torch.Tensor: """Combine terminal reward with per-token KL penalties.""" B, T = log_probs_policy.shape rewards = -kl_penalty_per_token(log_probs_policy, log_probs_ref, beta)

# Add scalar reward at the final token position rewards[:, -1] += scalar_reward return rewards # (B, T) — feed into PPO as token-level rewards ```