Language models learn to safely avoid risky actions when unsure

Training Language Models to Fail Gracefully

Machine Learning

Summary

Language models often try to solve tasks by guessing answers, but sometimes guessing can cause big problems, like breaking software or corrupting data. The authors propose a way to train these models so that if they think they might fail, they choose a safe fallback option instead of taking risky guesses. They created a method called GracefulRL that helps models learn this behavior just from general success or failure signals, without needing extra instructions. This helps the models keep doing well on tasks while avoiding dangerous mistakes when uncertain.

What this means in practice

  • For software developers: Improve code generation tools by enabling models to safely defer when unsure, reducing the chance of shipping faulty code.
  • For database administrators: Use language models that cautiously avoid actions that could corrupt databases when reliability is low, protecting data integrity.

Authors

Daman Arora, Xuecheng Liu, Himanshu Gaurav Singh, Fahim Tajwar, Abitha Thankaraj, Ruslan Salakhutdinov, Haiwen Feng, Yuda Song, Andrea Zanette

Abstract

Reinforcement learning makes language models better at solving tasks, but the same mechanism also rewards guessing. When success is the only way to earn reward, opting for a safe fallback is always worse than a long-shot attempt. When a wrong action can corrupt a database or ship broken code, that incentive can be costly. Therefore, we propose training models to fail gracefully: to preserve successful behavior whenever possible while learning to execute a safe fallback policy when acting would be unreliable. However, naively using policy gradients with intermediate reward for fallback behavior is brittle due to unfavorable optimization dynamics. We address this challenge with GracefulRL, which reshapes the optimization landscape to promote graceful failure. Importantly, this behavior emerges directly from outcome-level reward signals, without annotations specifying when to fall back or additional inference-time machinery such as a verifier. Across a range of tasks and scales, we train models that largely preserve their task-solving capability while selectively falling back when acting would be unreliable. More broadly, our results show that reinforcement learning can not only incentivize models to succeed, but also to reason about the limits of their capabilities and change course when they determine that success is out of reach.