RL with KL penalties is better viewed as B ayesian inference (2022)

Abstract: 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.

## introduction large language models (lms), such as @xcite , tend to generate outputs that reflect undesirable features of their training data such as offensiveness @xcite , social bias @xcite , harmfulness @xcite or dishonesty @xcite . addressing these biases and constraining lms to be honest, helpful and harmless is an essential part of the problem of aligning lms with human preferences @xcite . one intuitive approach to aligning lms is reinforcement learning (rl): capturing human preferences as a reward function and fine-tuning the lm to maximise the reward expected under lm distribution. a practical recipe for implementing this idea is rl from human feedback @xcite : first, a reward model is trained to predict which of two texts a human prefers and then a pretrained lm is fine-tuned to maximise reward given by the reward model while also being penalized for kullback-leibler (kl) divergence from its initial distribution. however, despite immense popularity of rl from human feedback @xcite @xcite , the motivation for kl penalty is not widely understood. in this paper, we discuss an underappreciated perspective on kl-regularised rl -the objective employed by rl from human feedback for fine-tuning lms -which explains its empirical success. we start with describing a problem that arises from naively applying the standard rl objective: distribution collapse. the optimal policy under the rl objective would be a minimalentropy lm generating a small set of sequences that obtain the highest reward. then, we discuss how kl-regularised rl avoids distribution collapse due to its kl penalty. this constraint, we argue, transforms the problem from rl to bayesian inference: updating a prior to conform with evidence provided by the reward. the bayesian perspective moves kl-regularised rl closer to other divergence-minimisation-based approaches to fine-tuning lms @xcite and, more broadly, to other divergence-minimisation-based accounts of control @xcite . these divergence minimisation approaches naturally avoid the distribution collapse problem because they formalize the agent as a generative model. in contrast, rl avoids distribution collapse only with reward functions that make it equivalent to divergence minimisation. therefore, we conclude, rl is not an adequate formal framework for problems such as fine-tuning lms. let x be the set of sequences of tokens from some vocabulary. an lm π can be seen as a probability distribution over x . while most modern lms are autoregressive, for simplicity we will only talk about full sequences, e.g. π(x) denotes the probability of a sequence @xmath0. similarly, a reward function r assigns sequences @xmath1. in practice, r(x) could represent human preferences we would like π to be aligned with, e.g. a non-offensiveness reward would assign low values to sequences that are offensive. if π θ is our parametric lm (with parameters θ), the rl objective for fine-tuning it with our reward function r is just the reward expected under lm distribution: intuitively, maximising j rl (θ) means sampling a number of sequences from the lm and rewarding the lm for good sequences and penalising for bad ones (e.g. offensive sentences). the problem with the rl objective is that it treats the lm as a policy, not as a generative model. while a generative model is supposed to capture a diverse distribution of samples, a policy is supposed to chose the optimal action. since we don't have a notion of state for lms, the rl objective reduces to searching for x * , the sequence with highest reward. if there is one, the optimal policy π * is a degenerate, deterministic generative model that puts entire probability mass on that single sequence: while the distribution collapse problem is exacerbated by rl failure modes such as insufficient exploration or reward hacking, it is distinct from exploration-exploitation trade-off or reward misspecification. even with perfect exploration (if we sampled sequences uniformly from x as opposed to sampling from π θ ), the optimal policy will still put all probability mass on x * . similarly, even if r perfectly captures human preferences across the whole space of possible sequences x and if x * is truly the best thing, we still would not want the lm to generate only x * . 1 essentially, the distribution collapse problem arises from the fact that 1 there is a case to be made that in conditional generation (e.g. translation or summarisation) one really cares only about the single best output for a given context (e.g. summary of a document). there are still, however, substantial benefits of caring about distributional aspects in conditional generation. first, when the lm produces a full distribution, we are able to measure its uncertainty. for larger models, these uncertainty estimates happen to be well-calibrated and allow for safer deployment in high-stakes scenarios @xcite . second, map estimates of the output distribution (the single, most likely output) are frequently of poor quality and can be substantially improved upon with decoding procedures taking into account the entire distribution, e.g. minimum bayes risk in translation @xcite or self-consistency chain-of-thought in question answering @xcite . @xcite provide a unifying perspective on multistep generation as latent variable modelling. the rl objective for lm alignment is flawed: it does not care about preserving distributional properties of an lm and will always penalise the lm for putting any probability mass on non-optimal sequences until the lm collapses into a degenerate distribution. ## fine-tuning language models via kl-regularised rl there is an obvious solution to the distribution collapse problem: including preserving distributional properties of an lm as part of the reward function. the notion of preserving distributional properties of an lm π θ can be formalised as penalising for kullback-leibler (kl) divergence between π θ and some other, pretrained lm π 0 . typically, π θ is initialised to π 0 and then fine-tuned to maximise the following objective: the first term in the right-hand side of (3) is equivalent to j rl (θ) in (1) while the second additionally constrains π θ to stay close (in terms of kl) to π 0 . almost always some reward needs to be sacrificed for that; the coefficient β determines the trade-off of how much reward is needed to justify departing from π 0 by a certain distance. this objective is commonly used as part of a popular recipe for finetuning lms termed "rl from human feedback" (rlhf) and works surprisingly well in practice @xcite @xcite . earlier approaches to fine-tuning lms employing this objective used the called it "conservative fine-tuning" @xcite or kl-control @xcite . here, we focus only on the policy optimisation part of this setup, which we term "kl-regularised rl". the kl-regularised rl objective (3) can easily be reformulated as just expected reward as in (1). we only have to define a new reward function r ′ θ (x) which incorporates the original reward r and the kl penalty, using the definition of kl divergence: ## kl-regularised rl as variational inference fine-tuning a pretrained lm π 0 to align with preferences encoded by a reward function r is essentially a bayesian inference problem. intuitively, bayesian inference is the problem of updating a distribution to conform with new evidence. in our setting, we're updating π θ , which is initially equal to a prior π 0 to conform with evidence provided by the assumption that π θ is optimal in terms of r. a reward function can be represented as a distribution over x that makes high-reward sequences more likely that low-reward sequences. a simple way of doing that is exponentiating reward r and renormalizing it. then, the posterior is given by: moreover, the kl-regularised rl objective can be cast as minimising the kl divergence between the lm π θ and this target distribution π * kl-rl : this divergence is different from the kl penalty term d kl (π θ , π 0 ) in (3). minimising this new divergence coincides with a variational inference @xcite , a well-known approach to approximating bayesian inference. more formally, j kl-rl (θ) is the evidence lower bound (elbo) on the log likelihood of π θ being optimal under r, assuming a prior π 0 . minimising this bound makes π θ approximate the true posterior π * kl-rl . a derivation of these equalities can be found in the appendix below. why is this picture insightful? it explains where the kl penalty term βd kl (π θ , π 0 ) in klregularised rl's original objective comes from. it is necessary to transform the problem from rl to minimising a divergence from a target distribution π * rlkl . this in turn makes the distributional character of an lm a first-class citizen which explains why kl-regularised rl is able to maintain the fluency and diversity of the original lm π 0 . ## separation of modelling and inference the bayesian perspective suggests that that aligning an lm with task preferences is a two-step process. it consists of, first, defining a distribution specifying the desired behaviour of your lm, and second, solving the problem of sampling from that posterior. these two steps roughly correspond to modelling and inference in probabilistic programming @xcite . modelling is encoding knowledge in probabilistic terms (usually by defining a probabilistic graphical model) while inference corresponds to using this model to answer queries. it is hard to overstate how useful -theoretically and practically -separating these two concerns could be. let us discuss these two steps, separately, below. ## is rl a good framework for fine-tuning language models? there is a family of other divergence minimisation approaches to fine-tuning lms which are not equivalent to rl. take generative distributional control (gdc) @xcite , an approach to fine-tuning lms that obtains results comparable with kl-regularised rl but minimises a slightly different divergence (forward as opposed to reverse kl). however, j gdc (θ) is no longer equivalent to rl @xcite because the expectation in forward kl divergence is with respect to a π * kl-rl , not π θ . similarly, standard supervised training objective can be seen as minimising d kl (π * mle , π θ ), a divergence from the empirical distribution π * mle provided by the training set. one can therefore mount a double dissociation argument in favour of the divergence minimisation perspective on kl-regularised rl: rl without kl divergence minimisation leads to degenera-tion while kl divergence minimisation without rl works well. therefore, it is the kl divergence minimisation aspect of kl-regularised rl that seems to account for its success, not the reward maximisation aspect. in consequence, calling it rl is just a redescription of it that happens to be correct under a particular choice of reward function r ′ θ . however, this redescription does not provide motivation for this choice of r ′ θ and does not hold for alternative divergence minimisation approaches to fine-tuning lms such as gdc @xcite . the divergence minimisation perspective on klregularised rl we presented stems from a general framework known as control as inference @xcite . control as inference provides a formalisation of intelligent decision making as inference on a probabilistic graphical model representing the agent, its preferences and environmental dynamics. while control as inference is typically considered with graphical models parameterised to make it equivalent to rl, it does not have to be. moreover, there are frameworks such as active inference @xcite or action and perception as divergence minimisation @xcite ) that further generalise control as inference to a principle of minimising the kl divergence from a probability distribution representing desired behaviour of the agent. in contrast with rl, they conceptualise the agent as a generative model, not as a decision rule represented as a probability distribution out of convenience. therefore, they naturally avoid the distribution collapse problem and preserve the distributional properties of the agent. what if rl simply isn't an adequate formal framework for problems such as aligning lms?