Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Reducing the Probability of Undesirable Outputs in Language Models Using Probabilistic Inference

Stephen Zhao, Aidan Li, Rob Brekelmans, Roger Grosse

TL;DR

This work addresses the challenge of minimizing the probability of undesirable outputs in language models without sacrificing average performance. It introduces RePULSe, which augments standard RL training with a low-reward-focused sampling mechanism guided by probabilistic inference, including a learned proposal $q_{\xi}$ and a target distribution $oldsymbol{ackslash}\sigma_{ heta}$. The approach yields improved tradeoffs between expected reward and the likelihood of bad outputs and demonstrates enhanced robustness to adversarial prompts in both toy and realistic LM settings. By combining RL with mass-covering sampling and end-to-end proposal learning, RePULSe offers a scalable path toward safer, more robust LM alignment in practice.

Abstract

Reinforcement learning (RL) has become a predominant technique to align language models (LMs) with human preferences or promote outputs which are deemed to be desirable by a given reward function. Standard RL approaches optimize average reward, while methods explicitly focused on reducing the probability of undesired outputs typically come at a cost to average-case performance. To improve this tradeoff, we introduce RePULSe, a new training method that augments the standard RL loss with an additional loss that uses learned proposals to guide sampling low-reward outputs, and then reduces those outputs' probability. We run experiments demonstrating that RePULSe produces a better tradeoff of expected reward versus the probability of undesired outputs and is more adversarially robust, compared to standard RL alignment approaches and alternatives.

Reducing the Probability of Undesirable Outputs in Language Models Using Probabilistic Inference

TL;DR

This work addresses the challenge of minimizing the probability of undesirable outputs in language models without sacrificing average performance. It introduces RePULSe, which augments standard RL training with a low-reward-focused sampling mechanism guided by probabilistic inference, including a learned proposal and a target distribution . The approach yields improved tradeoffs between expected reward and the likelihood of bad outputs and demonstrates enhanced robustness to adversarial prompts in both toy and realistic LM settings. By combining RL with mass-covering sampling and end-to-end proposal learning, RePULSe offers a scalable path toward safer, more robust LM alignment in practice.

Abstract

Reinforcement learning (RL) has become a predominant technique to align language models (LMs) with human preferences or promote outputs which are deemed to be desirable by a given reward function. Standard RL approaches optimize average reward, while methods explicitly focused on reducing the probability of undesired outputs typically come at a cost to average-case performance. To improve this tradeoff, we introduce RePULSe, a new training method that augments the standard RL loss with an additional loss that uses learned proposals to guide sampling low-reward outputs, and then reduces those outputs' probability. We run experiments demonstrating that RePULSe produces a better tradeoff of expected reward versus the probability of undesired outputs and is more adversarially robust, compared to standard RL alignment approaches and alternatives.
Paper Structure (57 sections, 22 equations, 12 figures, 1 table, 1 algorithm)

This paper contains 57 sections, 22 equations, 12 figures, 1 table, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background and preliminaries
  3. Language models
  4. Reinforcement learning
  5. Probabilistic inference
  6. Methodology
  7. $\text{RePULSe}$ gradient
  8. Low-reward target distributions
  9. Learning the proposal $q_\xi$ for approximate $\sigma_\theta$ sampling
  10. Experiments
  11. Experimental setup for all experiments
  12. Toy experiment -- illustrating motivation
  13. More realistic experiments -- investigating findings in practice
  14. Tradeoffs between average return and undesirable outputs
  15. Robustness to adversarial attack
  16. ...and 42 more sections

Figures (12)

  • Figure 1: Toy experiment: number of samples drawn vs. log total probability of bad outputs based on analytic calculation on a list of bad words. $\text{RePULSe}$ (blue dashed line) achieves much lower total probability of a bad output compared to baselines.
  • Figure 2: Toy experiment: Average reward over 500 samples as training progresses. We use the number of samples drawn during training on the x-axis to measure training progress. All methods achieve similar average reward, although RePULSe produces fewer low reward outputs (\ref{['fig:toy']}).
  • Figure 3: \ref{['setting1']}. Plot shows average return (including KL divergence) vs. total probability of bad outputs ($r(\bm{s} _{0:T}) < -5$) estimated from $p_{\theta}$ samples, evaluated on 10,000 held-out prompts with 5 samples each. Each point is an average over 10 seeds with 95% confidence intervals for both axes. $\text{RePULSe}$ improves the Pareto frontier at lower probabilities of bad outputs.
  • Figure 4: \ref{['setting2']}. Plot shows average return (including KL divergence) vs. total probability of bad outputs ($r(\bm{s} _{0:T}) < -7$) estimated from $p_{\theta}$ samples, evaluated on 2,500 held-out prompts with 4 samples each. Each point is an average over 5 seeds with 95% confidence intervals for both axes. $\text{RePULSe}$ improves the Pareto frontier.
  • Figure 5: \ref{['setting1']}. Average return vs. GCG attack success rate, evaluated on 10 held-out prompts, measured by any of 1000 $p_{\theta}$ samples having $r(\bm{s} _{0:T}) < -5$. Each point is an average over 10 seeds with 95% confidence intervals for both axes. $\text{RePULSe}$ achieves an improved frontier with better robustness over baselines.
  • ...and 7 more figures