Reasoning as an Adaptive Defense for Safety

TL;DR

This paper introduces TARS, a reinforcement learning method that trains large language models to reason about safety, improving their robustness against harmful prompts and jailbreak attacks by adaptively allocating compute during inference.

Contribution

The paper presents TARS, a novel RL-based training recipe that enhances LLM safety and robustness through adaptive reasoning and carefully designed training strategies.

Findings

Models trained with TARS better distinguish safe and unsafe prompts.

TARS-trained models show increased robustness to white-box and black-box attacks.

Adaptive reasoning leads to improved safety-refusal trade-offs.

Abstract

Reasoning methods that adaptively allocate test-time compute have advanced LLM performance on easy to verify domains such as math and code. In this work, we study how to utilize this approach to train models that exhibit a degree of robustness to safety vulnerabilities, and show that doing so can provide benefits. We build a recipe called $\textit{TARS}$ (Training Adaptive Reasoners for Safety), a reinforcement learning (RL) approach that trains models to reason about safety using chain-of-thought traces and a reward signal that balances safety with task completion. To build TARS, we identify three critical design choices: (1) a ``lightweight'' warmstart SFT stage, (2) a mix of harmful, harmless, and ambiguous prompts to prevent shortcut behaviors such as too many refusals, and (3) a reward function to prevent degeneration of reasoning capabilities during training. Models trained with TARS exhibit adaptive behaviors by spending more compute on ambiguous queries, leading to better safety-refusal trade-offs. They also internally learn to better distinguish between safe and unsafe prompts and attain greater robustness to both white-box (e.g., GCG) and black-box attacks (e.g., PAIR). Overall, our work provides an effective, open recipe for training LLMs against jailbreaks and harmful requests by reasoning per prompt.