PaperlayerView exact revision on arXiv

Source-linked AI summary

RL with KL penalties is better viewed as Bayesian inference

Tomasz Korbak, Ethan Perez, Christopher L Buckley

arXiv:2205.11275v2cs.LGstat.ML

TL;DR

Standard RL alignment can collapse a language model’s distribution, while KL-regularised RL constrains fine-tuning to remain near the pretrained model. The paper interprets that regularized objective as variational inference over a reward-informed Bayesian posterior, explaining how it preserves fluency and diversity; it also notes that its Bayesian account is preliminary and misses some empirical regularities.

  • Problem

    Standard RL fine-tuning can turn an LM into a degenerate distribution concentrated on high-reward sequences, despite the goal of aligning outputs with human preferences.

    Source: §1 · p. 1, §2 · p. 2

  • Method

    The paper frames KL-regularised RL as variational inference that approximates a Bayesian posterior updating a prior LM with evidence from the reward function.

    Source: §4 · p. 3

  • Results

    The KL-regularised RL solution coincides with the reward-updated version of the prior LM, and this Bayesian view explains how it maintains the original LM’s fluency and diversity.

    Source: §4 · pp. 3–4

  • Takeaways & Limitations

    The authors conclude that LM fine-tuning is better framed as Bayesian inference that approximates a target distribution than as RL alone.

    Source: §1 · p. 2

    The Bayesian proposal is preliminary and does not account for some empirical regularities found in RLHF experiments.

    Source: Limitations · p. 5

Abstract

from arXiv · show

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.

1 Introduction

Language models can reflect undesirable features of their training data, motivating alignment with human preferences. The paper argues that standard RL can collapse the model’s distribution, while KL-regularised RL is better understood as Bayesian inference.

1 Introduction

  • Language models can reflect undesirable features of training data, making alignment with human preferences an important goal.Source: §1 · p. 1
  • RL from human feedback trains a reward model from human preferences, then fine-tunes a pretrained LM to maximize reward while penalizing KL divergence from its initial distribution.Source: §1 · p. 1
  • The paper argues that standard RL can collapse an LM into a small set of high-reward sequences, whereas KL-regularised RL updates a prior to fit reward-provided evidence.Source: §1 · pp. 1–2
  • The authors conclude that RL is not an adequate formal framework for LM fine-tuning because it avoids collapse only when equivalent to divergence minimization.Source: §1 · p. 2

2 Fine-tuning language models using standard RL and distribution collapse

With no state notion, standard RL treats an LM as a policy seeking the highest-reward sequence rather than as a model of a diverse distribution. The resulting collapse is observed in generation tasks and remains distinct from exploration failure or reward misspecification.

2 Fine-tuning language models using standard RL and distribution collapse

  • Without a state notion, standard RL searches for the highest-reward sequence and yields a deterministic policy concentrated on it, rather than a diverse generative model.Source: §2 · p. 2
  • Reward-maximization-induced collapse manifests as decreased fluency and diversity, measured through perplexity, entropy, and repetition frequency, across multiple generation tasks.Source: §2 · p. 2
  • Collapse is distinct from exploration-exploitation trade-offs and reward misspecification: even perfect exploration or perfectly specified preferences would not make generating only the best sequence desirable.Source: §2 · p. 2
  • Even in conditional generation, retaining distributional information supports uncertainty estimates and decoding methods that can improve on the most likely output.Source: §2 · p. 2
  • The standard RL objective penalizes probability mass on non-optimal sequences until the LM collapses into a degenerate distribution.Source: §2 · p. 3

3 Fine-tuning language models via KL-regularised RL

KL-regularised RL adds a penalty for divergence from a pretrained LM to the expected-reward objective. The paper then questions whether this objective needs an RL framing, noting that its penalty can also be incorporated into a modified reward.

3 Fine-tuning language models via KL-regularised RL

  • KL-regularised fine-tuning maximizes expected reward while penalizing divergence from a pretrained LM, preserving distributional properties of the original model.Source: §3 · p. 3
  • The KL penalty trades off reward against staying close to the initial LM, and the objective is commonly used in RLHF.Source: §3 · p. 3
  • The KL-regularised objective can be rewritten as expected reward under a modified reward that incorporates the original reward and KL penalty.Source: §3 · p. 3
  • The paper asks whether framing this maximization as RL is necessary, introducing an alternative view based on approximate Bayesian inference.Source: §3 · p. 3

4 KL-regularised RL as variational inference

The paper recasts KL-regularised RL as Bayesian inference: a reward defines evidence for updating a pretrained LM prior. Its objective approximates the resulting posterior, explaining the role of the KL penalty and how the method preserves the original LM’s fluency and diversity.

4 KL-regularised RL as variational inference

  • The target posterior updates prior π0 using reward-derived evidence exp(r(x)/β), with Z normalizing the distribution.The reward supplies evidence that makes high-reward sequences more likely.Source: §4 · p. 3
  • The Bayesian posterior π*KL-RL coincides with the optimal policy for the KL-regularised RL objective.Source: §4 · p. 3
  • KL-regularised RL can be framed as minimizing divergence from the target posterior, and its objective is an ELBO that supports variational inference.This target divergence differs from the original KL penalty against π0.Source: §4 · p. 3
  • Treating the LM as a distribution to approximate explains how KL-regularised RL maintains the original LM’s fluency and diversity.The Bayesian view explains the KL penalty as necessary for transforming the problem into divergence minimization against the target.Source: §4 · pp. 3–4

5 Separation of modelling and inference

The Bayesian view separates alignment into defining a desired distribution and approximating it. It presents reward interpretation as a modeling choice and variational or sampling-based methods as inference approaches.

5 Separation of modelling and inference

  • Alignment has two steps: define a distribution for desired behavior, then solve the problem of sampling from that posterior.The paper relates these steps to modeling and inference in probabilistic programming.Source: §5 · p. 4
  • Exponentiating reward is one way to turn preferences into a distribution; binary constraints can instead assign zero probability to curse-containing sequences.Under that constraint, other sequences retain their original probabilities up to normalization.Source: §5 · p. 4
  • Inference can approximate a posterior with a parametric LM through variational inference or generate resembling samples through sampling-based methods.The paper connects KL-regularised RL to variational inference and decoding-time filtering, weighted decoding, and PPLM to sampling-based approaches.Source: §5 · p. 4
  • Separating target definition from approximation distinguishes misspecified preferences from failure to approximate the model well.This Bayesian framing unifies fine-tuning and decoding-time alignment approaches.Source: §5 · p. 4

6 Is RL a good framework for fine-tuning language models?

The paper argues that divergence minimization, rather than reward maximization, accounts for KL-regularised RL’s success. Other divergence-based fine-tuning approaches and broader inference frameworks show why RL is not an adequate general framework for LM alignment.

6 Is RL a good framework for fine-tuning language models?

  • Generative Distributional Control obtains results comparable to KL-regularised RL using forward rather than reverse KL, and its objective is not equivalent to RL.Source: §6 · p. 4
  • Standard supervised training can also be framed as minimizing KL divergence from a target distribution provided by the training set.Source: §6 · p. 4
  • Reward-only RL leads to degeneration, while KL divergence minimization can work without RL, implicating divergence minimization rather than reward maximization in KL-RL’s success.Source: §6 · pp. 4–5
  • The RL description is equivalent only under a particular reward choice and does not justify that choice or extend to alternatives such as GDC.Source: §6 · p. 5
  • Control-as-inference and related divergence-minimization frameworks treat the agent as a generative model, naturally avoiding distribution collapse.Source: §6 · p. 5

Limitations

The discussion identifies limits in both RLHF’s reward-model pipeline and the proposed Bayesian framing: reward models are imperfect proxies, while the Bayesian account remains preliminary and offers limited experimental guidance.

Other limitations of RLHF

  • The analysis focuses on fine-tuning against a reward model, leaving other RLHF shortcomings outside its scope.Reward-model training is treated as given, although the model is a proxy for underlying preferences and may be exploited through adversarial examples.Source: Limitations · p. 5
  • Reward-model design can introduce unintended bias through data curation, annotation guidelines, and annotator selection or compensation.Crowdworker preferences may not represent the general population, and annotators may disagree with one another and with researchers.Source: Limitations · p. 5

Limitations of the Bayesian perspective

  • The Bayesian account remains preliminary: it does not explain some empirical regularities in RLHF or guide choices such as hyperparameter selection.The authors specifically flag an observed relationship between expected reward and KL distance as unexplained.Source: Limitations · pp. 5–6

Other limitations of RLHF

Ethics statement

The ethics statement situates the paper in work on social bias and human-aligned AI, while noting that RLHF can also be diverted to harmful uses.

Ethics statement

  • Biased language-model outputs may amplify and perpetuate sexist, racist, or homophobic viewpoints present in training data.The paper frames research on reducing these risks as one motivation for studying RLHF.Source: Ethics statement · p. 6
  • The paper also frames language-model preference alignment as part of the broader effort to align AI systems with human values.The authors describe capable language models as a testbed for alignment techniques intended for future, more powerful systems.Source: Ethics statement · p. 6
  • Although the paper’s theoretical contribution is not expected to pose significant risk, RLHF is dual use and could support misinformation or harmful content.Source: Ethics statement · p. 6

Appendix

The appendix derives KL-regularised RL as approximate Bayesian inference by treating optimality as evidence and the pretrained language model as a prior. It connects the resulting variational objective to both the posterior target and the KL-RL objective.

Appendix

  • The derivation assumes a sequence prior and a reward that is always negative, then represents optimality as evidence about the language model.For a non-offensiveness reward, the optimality likelihood corresponds to the probability that a sequence is non-offensive.Source: Appendix · p. 9
  • With the prior set to the pretrained model, conditioning on optimality yields a posterior proportional to the prior times exponentiated reward.The normalising constant is denoted Z, and the posterior is equivalent to the KL-RL optimal policy, up to the choice of β.Source: Appendix · p. 9
  • Variational inference introduces a parametric distribution over sequences and uses Jensen’s inequality to form an evidence lower bound on the log likelihood of optimality.The derivation proceeds by introducing a latent sequence, factorizing the joint distribution, and then bounding the likelihood.Source: Appendix · p. 9
  • The resulting KL-RL objective is equivalent to minimizing divergence from the target distribution formed by weighting the prior with exponentiated reward.The paper also expresses the objective as proportional to negative KL divergence from that target, up to a constant.Source: Appendix · p. 9
Loading 2205.11275v2…