8 Regularization
Throughout the RLHF optimization, many regularization steps are used to prevent over-optimization of the reward model. Over-optimization in these contexts looks like models that output nonsensical text. Some examples of optimization "off the rails" are that models can output followable math reasoning with extremely incorrect answers, repeated text, switching languages, or excessive special characters.
The most popular variant, used in most RLHF implementations at the time of writing, is a KL Distance from the current policy to a reference policy across the generated samples. Many other regularization techniques have emerged in the literature to then disappear in the next model iteration in that line of research. That is to say that regularization outside the core KL distance from generations is often used to stabilize experimental setups that can then be simplified in the next generations. Still, it is important to understand tools to constrain optimization in RLHF.
The general formulation, when used in an RLHF framework with a reward model, \(r_\theta\) is as follows:
\[r = r_\theta - \lambda r_{\text{reg.}}\]
With the reference implementation being:
\[r = r_\theta - \lambda_{\text{KL}} \mathcal{D}_{\text{KL}} \left( \pi^{\text{RL}}(y \mid x) \| \pi^{\text{Ref.}}(y \mid x) \right)\]
8.1 KL Distances in RL Optimization
For mathematical definitions, see Chapter 5 on Problem Setup. Recall that KL distance is defined as follows: