Reducing the Probability of Bad Outputs in Language Models Using Probabilistic Inference
Stephen Zhao, Aidan Li, Rob Brekelmans, Roger Grosse
To avoid bad language model (LM) outputs, there are many alignment approaches (e.g., RLHF, DPO). Ideally, we would like our LM to have zero probability of undesirable outputs. Standard reinforcement learning (RL) would achieve this at optimality (if unregularized). However, in practice, there may be a tradeoff between methods focusing on the expected reward (standard RL) and methods explicitly focused on reducing the probability of undesired outputs. Our goal is to improve this tradeoff, reducing the probability of bad outputs as much as possible, while maintaining performance on expected reward. To do this, 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 to test whether our method provides better reduction of the probability of bad outputs and adversarial robustness, at minimal cost to expected reward, compared to standard RL alignment approaches and other alternatives.