PPO-ptx (pretraining gradient mixing)

Training language models to follow instructions with human feedbackintroduced

A modification of PPO fine-tuning that mixes gradients from the language-modeling loss on the pretraining distribution into the PPO updates, to counteract performance regressions on public NLP tasks.

Fine-tuning a language model with reinforcement learning tends to make it worse at things it already knew how to do. PPO-ptx is InstructGPT's answer: while the model chases reward, it keeps practicing plain language modeling on its original pretraining data, so the two pulls balance.

After the modification itself comes the problem it was built to fix, the so-called alignment tax, and then the more interesting question: how it stacks up against its two rivals.

The modification

PPO-ptx augments ordinary PPO fine-tuning by mixing gradients from the language-modeling loss on the pretraining distribution into the RL updates. The combined objective InstructGPT maximizes is $$\text{objective}(\phi) = \mathbb{E}_{(x,y)\sim D_{\pi^{RL}_\phi}}\Big[r_\theta(x,y) - \beta \log\big(\pi^{RL}_\phi(y|x)/\pi^{SFT}(y|x)\big)\Big] + \gamma\,\mathbb{E}_{x\sim D_{pretrain}}\big[\log(\pi^{RL}_\phi(x))\big],$$ where the first bracket is the standard reward-minus-KL-penalty term and the second injects the ordinary Language modeling (next-token prediction) objective on pretraining text, scaled by a coefficient $\gamma$ (instructgpt, §"3.5 Models", p. 9; Equation 2). In the final models $\gamma = 27.8$, and the pretraining data is drawn eight times more frequently than the number of RL episodes, with PPO and pretraining gradients accumulated together for each minibatch (instructgpt, §"C.4 Details of RLHF training", p. 42). Introduced by this paper, PPO-ptx is the variant meant by "InstructGPT" unless otherwise specified (instructgpt, §"3.5 Models", p. 9).

Why it exists: fixing the alignment tax

Plain PPO fine-tuning produced regressions on public NLP datasets, notably SQuAD v2 and DROP; PPO-ptx exists specifically to counteract them. Adding pretraining updates "greatly reduce[s] the performance regressions" on these datasets and "even surpasses GPT-3 on HellaSwag" (instructgpt, §"4.2 Results on public NLP datasets", p. 15), while, per the introduction, doing so "without compromising labeler preference scores" (instructgpt, §"1 Introduction", p. 4).

Against its two rivals

Two edges place PPO-ptx in context. Against tightening the Per-token KL penalty from the SFT model instead, the PPO-ptx (pretraining gradient mixing) — outperforms tightening of → Per-token KL penalty from the SFT model edge reports that raising $\beta$ sacrifices validation reward and never fully recovers DROP and SQuADv2 performance, while a well-chosen $\gamma$ reverses the regressions at minimal reward cost, because pretraining gradients restore the missing training distribution directly rather than merely constraining policy movement (instructgpt, §"4.2 Results on public NLP datasets", p. 15). Against plain PPO (without pretraining mix), the PPO-ptx (pretraining gradient mixing) — diverges only off distribution from → PPO (without pretraining mix) edge shows the two variants are nearly indistinguishable in human preference, differing almost entirely in capability retention, which is what makes PPO-ptx, not PPO, the model the paper crowns as InstructGPT. This concept is a child of PPO fine-tuning stage (RLHF policy optimization) within the The mechanics that make PPO fine-tuning work theme.