From Pre-training to Post-training

Large language models are usually built in two major stages:

1. Pre-training: The model learns to predict the next token given a prefix, acquiring broad linguistic competence.
2. Post-training: After the model can "speak", we teach it to "think". This stage splits into
   • Supervised Fine-Tuning (SFT)—imitating curated demonstrations.
   • Reinforcement Learning from Human Feedback (RLHF)—using a reward model and an RL algorithm to align outputs with nuanced human preferences.

SFT fits the model to demonstrations, while RL lets the model explore and discover better answers guided by rewards.

Direct Preference Optimization (DPO)

DPO bypasses explicit reward modeling by turning preference pairs into a differentiable loss. Given a prompt x, a preferred response yw, and a dis-preferred response yl, the objective is

\mathcal{L}_{\text{DPO}}(\pi_\theta)=-\mathbb{E}_{(x,y_w,y_l)\sim\mathcal{D}}\Bigl[\log\sigma\!\left(\beta\log\frac{\pi_\theta(y_w|x)}{\pi_{\text{ref}}(y_w|x)}-\beta\log\frac{\pi_\theta(y_l|x)}{\pi_{\text{ref}}(y_l|x)}\right)\Bigr]

where πref is a frozen reference (often the SFT checkpoint), and β controls the strength of the preference signal.

Implementation Sketch

1. Concatenate prompt and each response, tokenize, and feed to both πθ and πref.
2. Compute per-token log-probabilities for the response tokens only:

logp_w  = gather_logprobs(model, ids_w, len_prompt)
logp_l  = gather_logprobs(model, ids_l, len_prompt)
ref_w   = gather_logprobs(ref_model, ids_w, len_prompt)
ref_l   = gather_logprobs(ref_model, ids_l, len_prompt)

diff = beta * ((logp_w - ref_w) - (logp_l - ref_l))
loss = -torch.log(torch.sigmoid(diff)).mean()

Proximal Policy Optimization (PPO)

PPO keeps policy updates close to the previous policy via a clipped surrogate objective. Let rt(θ)=πθ(at|st)/πθold(at|st). The policy loss is

L^{\text{CLIP}}(\theta)=\hat{\mathbb{E}}_t\!\left[\min\!\bigl(r_t(\theta)\hat{A}_t,\ \text{clip}(r_t(\theta),1-\epsilon,1+\epsilon)\hat{A}_t\bigr)\right]

where Ât is the generalized advantage estimate (GAE). The full loss adds a value-functon MSE and an entropy bonus:

\mathcal{L}(\theta,\phi)=-\mathbb{E}_t[L^{\text{CLIP}}(\theta)] + c_1\mathbb{E}_t[(V_\phi(s_t)-\hat{R}_t)^2] - c_2\mathbb{E}_t[H[\pi_\theta](s_t)]

Training Loop

1. Rollout: For each prompt, sample K completions from πθold.
2. Reward & Advantage: Compute sequence-level reward r with a reward model, add KL penalty against πref, then estimate advantages via GAE using a critic Vψ.
3. Update: Minimize the composite loss for several epochs on the collected batch.

ratio = torch.exp(logp - old_logp)
clipped = torch.clamp(ratio, 1-eps, 1+eps)
policy_loss = -torch.min(ratio*adv, clipped*adv).mean()

value_loss  = F.mse_loss(values, returns)
entropy     = -(probs*logp).sum(-1).mean()

total = policy_loss + 0.5*value_loss - 0.01*entropy + beta*kl_penalty

Group Relative Policy Optimization (GRPO)

GRPO extends PPO to optimize over a group of responses per prompt. For a prompt q, sample G answers {oi}, compute their rewards {ri}, and normalize within the group:

\hat{A}_i=\frac{r_i-\mu_r}{\sigma_r+\epsilon},\quad \mu_r=\frac{1}{G}\sum_i r_i,\ \sigma_r=\sqrt{\frac{1}{G}\sum_i(r_i-\mu_r)^2}

The per-tokan GRPO loss is

\mathcal{J}_{\text{GRPO}}(\theta)=\mathbb{E}_{q}\!\left[\frac{1}{G}\sum_{i=1}^{G}\frac{1}{|o_i|}\sum_{t=1}^{|o_i|}\min\!\bigl(r_{i,t}(\theta)\hat{A}_i,\ \text{clip}(r_{i,t}(\theta),1-\epsilon,1+\epsilon)\hat{A}_i\bigr)\right]-\beta\,D_{\text{KL}}(\pi_\theta\|\pi_{\text{ref}})

Practical Tweaks

• Remove the 1/|oi| term to avoid length bias.
• Skip standardization when σr is near zero by injecting a synthetic perfect response (reward = 1.0).

rewards = torch.tensor([...])          # shape [G]
rewards = torch.cat([rewards, torch.tensor([1.0])])
mean = rewards.mean(); std = rewards.std()
adv = (rewards[:-1] - mean) / (std + 1e-8)

Quick Comparison

Choose DPO for fast, reward-free alignment; PPO for fine-grained token-level control; GRPO when you have verifiable rewards and want to leverage response diversity.