RLHF Scaling Math

Three budgets and one frontier

Strip RLHF down and it optimizes a single objective under a leash:

max_π  E_{x, y~π(·|x)} [ r̂(x, y) ]  −  β · KL( π || π_ref )

Here r̂ is the proxy reward — a learned reward model — not the gold reward r a careful human would assign. Every scaling question reduces to how three quantities interact: the number of preference comparisons used to fit r̂, the reward model’s parameter count, and the KL distance you let the policy travel. Notice what is missing: the size of the policy — not an oversight, and we come back to it. The unifying object is the frontier: gold reward against KL.

Why sqrt(KL) is the natural coordinate

Plot gold reward against raw KL and you get an ugly, sharply bending curve; plot it against d = sqrt(KL) and it straightens into a two-parameter fit. That is no accident — it falls out of a second-order expansion.

KL(π_θ || π_ref)  ≈  (1/2) · Δθ^T F Δθ      F = Fisher information matrix
  →  sqrt(KL)  ∝  ||Δθ||_F        (a genuine distance)

E[r]  ≈  r_0  +  ∇r · Δθ             (reward gain is FIRST-order in Δθ)
  →  E[r]  −  r_0  ∝  d                 with d ≡ sqrt(KL)

KL is quadratic in displacement; reward gain is linear in it. So reward gain is linear in sqrt(KL), and d is the Fisher–Rao distance travelled, in sqrt(nats).

Overoptimization: the shapes, and where the peak sits

Gao, Schulman and Hilton fit the gold-reward frontier empirically and found two clean shapes, one per optimization method:

best-of-n:   R_gold(d)  =  d · ( α_bon  −  β_bon · d )
RL (PPO):    R_gold(d)  =  d · ( α_rl   −  β_rl  · log d )

proxy:       R̂(d)      same family, larger α, smaller β → monotonically increasing

Read the best-of-n form as a linear gain αd minus a Goodhart correction βd²: the linear term is the reward model being right, the quadratic term is the policy finding regions where it is wrong, an error that compounds with distance. The decisive asymmetry is the last line — the proxy curve never turns over, so rising training reward is zero evidence that anything is improving. Now take α = 1.0, β = 0.2 and differentiate:

dR/dd  =  α − 2βd  =  0     →   d* = α / (2β) = 2.5 sqrt(nats)
KL*    =  d*²  =  6.25 nats
R_max  =  α² / (4β)  =  1.00 / 0.80  =  1.25

at d = 4.0 (KL = 16 nats):   R = 4.0 · (1.0 − 0.8)  = 0.80   → 64% of peak
at d = 1.25 (KL = 1.56):     R = 1.25 · (1.0 − 0.25) = 0.94   → 75% of peak

Spending 2.6× the optimal KL gives back a third of everything you gained, while a quarter of the optimal KL already banks three quarters of the available reward. The frontier is brutally concave: the cheap drift is where the value is.

Best-of-n gives you an exact KL ruler

Best-of-n is the workhorse of overoptimization studies because its KL is known in closed form, with no estimator noise:

KL_bon(n)  =  log n  −  (n − 1)/n

n = 16    →  2.773 − 0.938  =  1.835 nats   →  d = 1.355
n = 1408  →  7.250 − 0.999  =  6.251 nats   →  d = 2.500  = d*

Feed d = 1.355 back into the frontier: R = 1.355 · (1.0 − 0.271) = 0.99, or 79% of the peak from best-of-16. Reaching the optimum needs n ≈ 1400 — roughly ninety times the sampling cost for the last 21%. Because KL grows like log n while d grows like sqrt(log n), every further step costs exponentially more samples. That is why a PPO run is judged against a best-of-n baseline at matched KL, not at matched compute.

What a bigger reward model buys

The reward model’s parameter count enters almost entirely through β, the Goodhart coefficient: a larger reward model is a more robust proxy, staying accurate further off-distribution, so its error grows more slowly with distance. Suppose scaling it drops β from 0.20 to 0.15, leaving α alone:

β = 0.20:   d* = 2.50,  KL* =  6.25 nats,  R_max = 1.25
β = 0.15:   d* = 3.33,  KL* = 11.11 nats,  R_max = 1.67   (+33% reward, +78% KL)

Because d* = α/(2β) and R_max = α²/(4β), shrinking β raises the ceiling and extends the runway. The double dividend: a better reward model does not merely score more accurately, it licenses more optimization. The corollary bites the other way — if your reward model is small, the correct response is a tighter KL leash, not more PPO steps.

Preference data: bits per comparison

Preference data scales differently, and its ceiling is informational rather than computational. A single pairwise comparison carries at most one bit, and human labelers disagree. If two labelers agree with probability p, the information a label carries about the true preference is:

I  =  1 − H(p)  bits,     H(p) = −p·log₂p − (1−p)·log₂(1−p)

p = 0.75  →  H = 0.811  →  I = 0.189 bits per comparison
100,000 comparisons  →  ≈ 18,900 bits  ≈ 2.4 kB of signal

Two kilobytes — the entire budget of new information a large preference dataset conveys. That is why a reward model is a fine-tune and never trained from scratch: the pretrained prior supplies the structure, and the comparisons only select among hypotheses it already represents. It also explains the log-linear accuracy curve and the hard ceiling near inter-annotator agreement. A reward model scoring far above the human agreement rate is fitting annotator idiosyncrasy, not preference.

Policy size: the near-invariance

Now the omission from the first section. Empirically, α and β are roughly independent of policy size in the sqrt(KL) coordinate: a 1B and a 70B policy start at different absolute reward levels, but the shape of the rise and fall with distance travelled is close to the same curve.

This is practical: the KL budget you tune on a small, cheap policy transfers to the expensive one, so the search over β, PPO epochs and the early-stopping point runs at 1/50th the cost. What does not transfer is the reward model — β is a property of the proxy, so a reward model adequate for a small policy will be over-optimized faster by a stronger one that explores harder and finds its blind spots sooner.

FLOP accounting: the loop versus pretraining

Pretraining costs about 6ND FLOPs for N parameters over D tokens (2 forward, 4 backward, per parameter per token). An RLHF iteration is dearer per token because four models touch every rollout. Per generated token, with K PPO epochs and a reward and value model each of size N:

So RLHF costs roughly 5× more FLOPs per token than pretraining, but runs on four to five orders of magnitude fewer tokens. Take a 7B policy pretrained on 2T tokens, then aligned with 400k rollouts of 512 generated tokens each:

pretrain:  C = 6ND       = 6 × 7e9 × 2e12   = 8.4e22 FLOPs
RLHF:      T = 4e5 × 512 = 2.05e8 generated tokens
           C = 32NT      = 32 × 7e9 × 2.05e8 = 4.6e19 FLOPs

ratio  =  4.6e19 / 8.4e22  ≈  5.5e−4   →  about 0.05% of pretraining

Alignment is a rounding error in FLOPs — but the ratio flatters it. Two thirds of that cost is autoregressive decoding, memory-bandwidth-bound at roughly one FLOP per byte moved, typically achieving 5–15% model-FLOP utilization against 40–50% for pretraining. In wall-clock terms the gap narrows by 3–10×.

Why PPO does not fit a CPU-class model

Run the memory accounting for a 1.3B policy with bf16 weights and Adam in fp32 (4-byte master copy plus two 4-byte moments = 12 bytes per trained parameter):

PPO: weights  policy+ref+RM+value = 4 × 2.6 GB          = 10.4 GB
     optim    policy+value trained = 2 × 12B × 1.3e9  = 31.2 GB
     total                                                 ≈ 42 GB

DPO: weights  policy only          = 2.6 GB
     optim    12B × 1.3e9         = 15.6 GB
     ref      log-probs PRECOMPUTED = 0 GB resident
     total                                                 ≈ 18 GB

DPO more than halves resident memory, but the bigger win is structural: it deletes the rollout loop entirely, converting the memory-bandwidth-bound decode phase into a batched supervised pass. On CPU, where decode throughput is the binding constraint, that is the difference between a job that finishes overnight and one that does not. The frontier math still applies — DPO’s β plays the same leash role — you simply cannot afford to measure the curve by sampling.

Measuring KL honestly

All of the above is worthless if your x-axis is wrong, and it usually is. Three failures dominate. First, the naive estimator log(π/π_ref) is unbiased but so high-variance that single-batch estimates routinely come out negative; use Schulman’s k3 estimator instead, with r = π_ref(y)/π(y):

KL̂_k3  =  (r − 1) − log r     ≥ 0 always, and unbiased for KL(π || π_ref)

Second, sequence-level KL scales with generation length, so a per-token KL of 0.01 over 512 tokens is 5.1 nats of sequence KL — nearly the whole budget in our worked example. Always state which one you are quoting. Third, KL is measured in the model’s own token space, so budgets are not comparable across tokenizers or reference models. Fix the reference, fix the length normalization, and only then trust a number like d* = 2.5.