Nathan Lambert on KL Regularization, RL and Supervised Fine-Tuning

This lecture connects regularization, policy updates and the question of why reinforcement learning may generalize differently from supervised fine-tuning. The immediate question is how an RL update changes a model's distribution and how we control that change.

In a conventional RLHF formulation, a learned reward is balanced against a KL-divergence penalty between the trained policy and a reference policy, often the model before RL. The penalty weight controls how strongly the optimization resists moving away from that reference. Related control ideas predate modern language-model RLHF.

A fixed reference-model KL penalty has become less common in some RLVR and agentic recipes, but the broader need to control updates remains. Different generations of training systems address related stability and over-optimization problems in different ways.

The reward objective subtracts the weighted divergence from the reward-model score. The reference is the policy before the RL updates. The same coefficient may be written as lambda or beta in different literature, so the notation must be read in context.

The penalty commonly uses reverse KL: samples come from the policy being trained and are evaluated relative to the reference. KL is often casually called a distance, but it is not symmetric and therefore is not a true distance metric. Direction matters because it determines which distribution supplies the samples.

Calling KL a distance can be helpful when watching how much a model changes, but swapping its arguments changes its value. Lambert returns to this asymmetry later when comparing distributions and the behavior of forward and reverse objectives.

In practice, teams estimate KL from sampled token log probabilities and monitor the curve for spikes or unusual movement. A run can track divergence from its starting model even when its explicit penalty is zero. Sequence-level estimates depend on implementation and sampling choices.

The implementation uses log probabilities of sampled tokens and can aggregate them over a sequence. Different estimators trade off practical properties; one should inspect the actual library rather than assume every code snippet computes the same quantity.

A plotted OLMo math-training run starts at zero divergence because the trained policy initially equals its reference. Even without an active penalty, the trajectory can show how the model moves. The shape and scale of the curve are diagnostic rather than a stand-alone measure of quality.

Older RLHF work sometimes adjusted the penalty coefficient over time to maintain a target divergence, using a controller. Modern recipes often use a fixed coefficient or omit the explicit reference penalty. Lambert points to earlier PPO and human-preference fine-tuning work as examples of adaptive control.

Lambert distinguishes an explicit penalty in the reward objective from a separate description of the shape of an optimization problem. His term implicit regularization refers to the latter: analyzing how the objective and sampling distribution influence the policy, rather than adding a separate term by hand.

Supervised fine-tuning samples from a fixed dataset or teacher; minimizing its negative log likelihood corresponds, up to a constant independent of the model, to minimizing forward KL from the data distribution to the learned policy. RL samples from its own policy. Under a positive KL penalty, the regularized reward objective can be rewritten as reverse-KL minimization toward a reward-tilted reference.

That exact reverse-KL equivalence depends on the penalty being present. If its coefficient is zero, this particular derivation does not apply. Lambert stresses the condition because many current verifiable-reward recipes omit the explicit term.

The reward-tilted optimal policy reweights the reference by an exponential function of reward, with a normalization factor. A smaller penalty weight allows reward to pull the policy further; a larger one keeps the optimum closer to the reference. Substituting that optimum into the reverse-KL expression recovers the regularized reward objective up to constants and scale.

The normalization factor in the reward-tilted policy makes it a valid distribution. Lambert walks through substituting this optimum into reverse KL, separating terms and dropping constants. Multiplying by the penalty coefficient relates the minimization back to the original reward-maximization form.

For supervised fine-tuning, expanding forward KL yields a data entropy term independent of the model plus the supervised loss. The constant drops out when taking gradients. The two views describe different problems—learning from data and maximizing reward—but help compare their optimization shapes.

Forward KL is often described as mass-covering: if the target data assigns probability where the learned policy approaches zero, the loss becomes large. The learned distribution is pushed to cover target mass. Reverse KL is often described as mode-seeking, concentrating on high-reward regions under the stated setup.

If a target example has probability under the data distribution but nearly none under the trained policy, forward KL produces a large loss. That pressure encourages coverage of the target distribution. It should not be confused with an instruction to retain every mode of the original model.

A similar distinction appears in distillation: learning from a fixed teacher resembles the forward direction, while on-policy methods connect naturally to the reverse direction. These are useful intuitions, but the exact distributions and assumptions still matter.

Lambert reviews work titled SFT Memorizes, RL Generalizes, which compares training on simple tasks with tests that vary task properties. The reported experiments show RL generalizing better in those settings while narrow supervised training can harm generalization. He presents this as a useful controlled example, not a universal guarantee.

Why might a mass-covering supervised objective still forget? The target data distribution is not identical to everything the original model knew. In the illustrated multimodal example, fitting a new target can spread a learned mode and take probability from an older mode. Covering the new data does not ensure preservation of every prior capability.

In the diagram, the original and target policies contain more than one mode. The forward-KL update can broaden a new mode and reduce mass in an old one; the reverse-KL illustration moves the new mode more locally. Lambert cautions that the figure is an intuition, not a complete account of neural-network training.

Lambert discusses another paper, RL Razors, whose experiments compare RL- and SFT-trained models at similar task performance. The reported RL policies move less in KL terms and forget less out-of-domain knowledge in that setting. He treats these as empirical findings with limits, not proof that every RL run preserves every capability.

Other regularizers have appeared in training recipes. InstructGPT included a pretraining-data next-token objective alongside the RL reward to preserve language-model behavior. Some preference methods add a likelihood term for chosen responses rather than relying solely on the margin between chosen and rejected answers.

Llama 2 also used information about the strength of preferences in reward-model training. Such additions can solve a particular problem at a particular time, but Lambert recommends revisiting complexity in later recipes: simpler methods can be easier to scale and understand.

Recent agentic RL introduces a different stability concern from the original reference-model penalty. Besides controlling the trained policy relative to a fixed reference, systems must manage mismatch between the sampler that generates rollouts and the trainer that updates weights. These are distinct distributions and distinct problems.

The role of the reference changed across eras. Early RLHF emphasized keeping a model from exploiting a learned reward while drifting from its starting policy. In tool-use RL, numerical differences between the rollout engine and the trainer can destabilize updates even when no reference KL penalty is used.

Lambert discusses divergence proximal policy optimization, or DPPO, used in T-MAX, alongside related token-level approaches. Rather than relying only on a clipped probability ratio, these methods can mask tokens with problematic discrepancies. He notes that he is still learning the new algorithm and may have missed details.

The newer methods he names include token-level masking approaches associated with GLM and Kimi as well as DPPO. The aim is to handle extreme or unreliable discrepancies in sampled tokens. Lambert explicitly describes this section as preliminary and invites correction from practitioners.

The main distinction is between a penalty anchored to the original reference and controls on drift relative to the sampling policy for an update. Lambert concludes that forward and reverse KL provide helpful lenses for supervised and regularized RL objectives, while empirical studies of generalization and forgetting require careful interpretation.