RL with KL penalties is better viewed as Bayesian inference

By Tomasz Korbak, Ethan Perez, C. L. Buckley. 2022

138 citations recorded by the source corpus.

Reinforcement learning (RL) is frequently employed in fine-tuning large language models (LMs), such as GPT-3, to penalize them for undesirable features of generated sequences, such as offensiveness, social bias, harmfulness or falsehood. The RL formulation involves treating the LM as a policy and updating it to maximise the expected value of a reward function which captures human preferences, such as non-offensiveness. In this paper, we analyze challenges associated with treating a language model as an RL policy and show how avoiding those challenges requires moving beyond the RL paradigm. We start by observing that the standard RL approach is flawed as an objective for fine-tuning LMs because it leads to distribution collapse: turning the LM into a degenerate distribution. Then, we analyze KL-regularised RL, a widely used recipe for fine-tuning LMs, which additionally constrains the fine-tuned LM to stay close to its original distribution in terms of Kullback-Leibler (KL) divergence. We show that KL-regularised RL is equivalent to variational inference: approximating a Bayesian posterior which specifies how to update a prior LM to conform with evidence provided by the reward function. We argue that this Bayesian inference view of KL-regularised RL is more insightful than the typically employed RL perspective. The Bayesian inference view explains how KL-regularised RL avoids the distribution collapse problem and offers a first-principles derivation for its objective. While this objective happens to be equivalent to RL (with a particular choice of parametric reward), there exist other objectives for fine-tuning LMs which are no longer equivalent to RL. That observation leads to a more general point: RL is not an adequate formal framework for problems such as fine-tuning language models. These problems are best viewed as Bayesian inference: approximating a pre-defined target distribution.

Subfields: Reinforcement Learning

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Reinforcement Learning

Citations in the corpus

Scaling Laws for Reward Model Overoptimization (2022, background): “One possible hypothesis is that, because RLHF can be viewed as Bayesian inference from the prior of the initial policy (Korbak et al., 2022)11, increases in policy size are only improving the modelling accuracy of the human demonstration distribution.”
Aligning Language Models with Preferences through f-divergence Minimization (2023, methodology): “8 is very similar, with a “pseudo-reward” rθ(x) = −f ′ (πθ(x)p(x) ), one difference being that now rθ depends on θ (see (Korbak et al., 2022b) for related remarks).”
On Reinforcement Learning and Distribution Matching for Fine-Tuning Language Models with no Catastrophic Forgetting (2022, methodology): “Consistent with prior work (Khalifa et al., 2021; Korbak et al., 2022a), we observe that Reinforce is able to quickly achieve high levels of constraint satisfaction, but at the cost of large deviations from a, which translates into significantly decreased diversity of generated samples (in terms of…”
OpenChat: Advancing Open-source Language Models with Mixed-Quality Data (2023, methodology): “Following prior works (Peters & Schaal, 2007; Peng et al., 2019; Korbak et al., 2022; Rafailov et al., 2023), it can be shown that the optimal solution to the above KL-regularized reward maximization”
Training Language Models with Language Feedback at Scale (2023, methodology): “This lens highlights the correspondence between ILF and RL with Human Feedback (Ziegler et al., 2019; Stiennon et al., 2020, inter alia ), which was previously demonstrated to be equivalent to Bayesian inference (Korbak et al., 2022).”
Open Problems and Fundamental Limitations of Reinforcement Learning from Human Feedback (2023, methodology): “Although this issue also arises in pretrained language models, RLHF has not been a solution for it and can amplify it in some cases [Perez et al., 2022b].”