K-choose-2 single-batch comparison training

Training language models to follow instructions with human feedbackintroduced

A reward-model training trick that places all C(K,2) pairwise comparisons from one labeled ranking of K outputs into a single batch element, avoiding the overfitting that results from treating correlated comparisons as independent shuffled data points.

Ranking K outputs at once yields every pairwise comparison among them, and shuffling those correlated pairs into training as if they were independent makes a reward model overfit. The fix in InstructGPT (2022) is small: keep all the comparisons from one ranking inside a single batch element.

How the trick fits, what it enabled, and what it actually is, in that order, make up what follows: a small fix nested under the reward model it protects from overfitting.

How it fits

K-choose-2 batching is a training detail nested one level under the Reward model (RM), inside the broader The mechanics that make PPO fine-tuning work theme that documents the specific engineering choices that make InstructGPT's PPO stage actually work rather than merely describing it at a narrative level. It exists to solve a concrete problem created by how the RM (comparison) dataset is collected: labelers rank $K=4$ to $9$ completions per prompt rather than judging isolated pairs, and a natural-seeming way to turn that ranking into training data — enumerate every pairwise comparison and shuffle them all into one flat dataset — turned out to break the reward model.

What it enabled

The comparisons drawn from a single ranked set are highly correlated with each other, since they all describe the same $K$ completions to the same prompt. Training on them as independently shuffled data points meant each completion could appear in up to $K-1$ separate gradient updates scattered across an epoch, and this caused the reward model to overfit after a single pass over the data (instructgpt, §"3.5 Models", p. 8, n. 5). InstructGPT's fix is to place every comparison from one ranking into a single batch element instead: all $\binom{K}{2}$ comparisons from one prompt are processed together. This is not only a fix for overfitting but a compute win, since it requires only one reward-model forward pass per completion rather than one pass per comparison — a $K$-fold reduction, given that $\binom{K}{2}$ grows quadratically in $K$ while the number of completions grows only linearly (instructgpt, §"3.5 Models", p. 8). The paper reports that batching this way, rather than shuffling, achieves substantially improved validation accuracy and log loss for the reward model that everything downstream, PPO training included, depends on.

What it is

Concretely, for a ranking of $K$ labeler-ordered completions to a single prompt, there are $\binom{K}{2} = \frac{K(K-1)}{2}$ distinct pairwise comparisons obtainable from that ranking (for $K=4$, six comparisons; for $K=9$, thirty-six). K-choose-2 batching groups all of them, for a given prompt, into one batch element for a single gradient step, rather than distributing them across the training set as if they were independent, identically-distributed examples. This is why the RM (comparison) dataset's comparison count runs an order of magnitude beyond its roughly 33,000 prompts even though each gradient update still only touches one prompt's worth of correlated structure at a time.