LLM Post-training

For SFT collecting a high quality instruction-response dataset can be a bottleneck. So, a bigger "teacher" model can be used to generate synthetic dataset of tasks specific instruction-response pairs based on a few human-written "seed" tasks. This involves:

1. Task Generation: Sample tasks from the pool and generate new diverse instructions/tasks
2. Input/Output Generation: Take a task and generate the output
3. Filtering: To remove similar or low quality pairs. Alternatively, instead of removing similar pairs and wasting generated pairs, loosening the similarity threshold and focusing of local diversity within a training batch gives better & cost effective results.
4. Poll update: Add synthetic tasks back to the pool for iterative improvement

RLHF is done to align the model output to human preference. It is an important and distinct phase than SFT because:

1. SFT is a form of imitation learning, and the model can't do better than the examples that humans provide.
2. SFT also promotes copying the structure of the response instead of understanding the rules behind them. This is because SFT loss is based on Maximum Likelihood Estimation.
3. It is difficult to specify high quality answer to nuanced objectives like helpfullness, and harmlessness. Instead it is easier to judge whether one answer is better than the other.

RLHF solves the issues with SFT by optimizing for the human preference. There are three stages:

1. SFT: To establish a baseline

2. Reward Model training: Based on human ranking of multiple completions for same prompt, a reward model is trained to predict human choice.

   Based on ranking data, reward model \(r_{\phi}\) can be trained to assing reward value to each answer by optmizing for following objective (called the Bradley-Terry preference model):

   \begin{align*} \mathcal L(\phi) = - \mathbb E_{(x, y_a, y_b) \sim D} \left[ \log P(y_{a} \succ y_{b}| x) \right] \\ \textrm{where, } P(y_{a} \succ y_{b} | x) = \sigma(r_{\phi}(x, y_{a}) - r_{\phi}(x, y_{b})) \ . \end{align*}
3. Policy Optimization: SFT model is now fine-tuned using PPO with reward from Reward Model.

Using PPO requires multiple models: a) policy model being trained, b) a reference model to compute the clipped reward, c) reward model and d) a critic model. This is computationally expensive.

Instead of training Reward Model, we can do direct policy optimization. The KL constrained, RLHF objective can be mathematically manipulated to arrive at an alternative but equivalent loss function that just uses the reference model and the policy model:

\begin{align*} \mathcal L_{DPO}(\theta) = - \mathbb E_{(x, y_a, y_b) \sim D} \left[ \log \sigma \left( \beta \log {\frac{\pi_{\theta}(y_{a}|x)} {\pi_{ref}(y_{a}|x)}} - \beta \log {\frac{\pi_{\theta}(y_{b}|x)} {\pi_{ref}(y_{b}|x)}} \right) \right]\ . \end{align*}

But the lack of Reward model means this is "offline" in nature and can only learn from a pre-collected dataset of peferences, i.e. it cannot explore new responses that might be better than those in the training set.

Note that, DPO learns the preferences but it is not RLHF because DPO update is not Reinforcement Learning.

GRPO removes the memory bottleneck of PPO by getting rid of the critic model. For policy optimization, critic gives the advantage value, which is the amount by which a candidate response is better than the average response. In GRPO, it is computed not by a critic model, but by using the average reward for the group of responses generated for the same prompt. Specifically, the advantage is the given by normalizing with respect to the mean & standard deviation of group's rewards:

\begin{align*} A_{i} = \frac {r_{i} - \mu_{r}} {\sigma_{r}} \ . \end{align*}

In short what GRPO does is, generate a set of responses for the same prompt, evaluate the responses using a reward model and tune the model to generate responses that get better than average reward within the generated set.