Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

RAPO: Risk-Aware Preference Optimization for Generalizable Safe Reasoning

Zeming Wei, Qiaosheng Zhang, Xia Hu, Xingcheng Xu

TL;DR

This paper addresses the generalization gap in safe reasoning for Large Reasoning Models under sophisticated jailbreaks. It introduces Risk-Aware Preference Optimization (RAPO), a two-stage framework combining an SFT warm-up for format-aligned safe reasoning with RL that uses a risk-aware judgment and a general utility reward to adapt the depth of safe reasoning to prompt risk. Theoretical analysis from in-context learning shows the safe reasoning budget must scale with attack complexity ($t = \Omega(k)$), and empirical results across multiple models and benchmarks demonstrate RAPO achieves strong safety (low ASR on JailbreakBench, HarmBench, and WildJailbreak) while maintaining reasoning utility (MMLU performance). The work provides a principled, generalizable alignment technique for LRM safety with practical implications for deploying LRMs in safety-critical contexts.

Abstract

Large Reasoning Models (LRMs) have achieved tremendous success with their chain-of-thought (CoT) reasoning, yet also face safety issues similar to those of basic language models. In particular, while algorithms are designed to guide them to deliberately refuse harmful prompts with safe reasoning, this process often fails to generalize against diverse and complex jailbreak attacks. In this work, we attribute these failures to the generalization of the safe reasoning process, particularly their insufficiency against complex attack prompts. We provide both theoretical and empirical evidence to show the necessity of a more sufficient safe reasoning process to defend against advanced attack prompts. Building on this insight, we propose a Risk-Aware Preference Optimization (RAPO) framework that enables LRM to adaptively identify and address the safety risks with appropriate granularity in its thinking content. Extensive experiments demonstrate that RAPO successfully generalizes multiple LRMs' safe reasoning adaptively across diverse attack prompts whilst preserving general utility, contributing a robust alignment technique for LRM safety. Our code is available at https://github.com/weizeming/RAPO.

RAPO: Risk-Aware Preference Optimization for Generalizable Safe Reasoning

TL;DR

This paper addresses the generalization gap in safe reasoning for Large Reasoning Models under sophisticated jailbreaks. It introduces Risk-Aware Preference Optimization (RAPO), a two-stage framework combining an SFT warm-up for format-aligned safe reasoning with RL that uses a risk-aware judgment and a general utility reward to adapt the depth of safe reasoning to prompt risk. Theoretical analysis from in-context learning shows the safe reasoning budget must scale with attack complexity (), and empirical results across multiple models and benchmarks demonstrate RAPO achieves strong safety (low ASR on JailbreakBench, HarmBench, and WildJailbreak) while maintaining reasoning utility (MMLU performance). The work provides a principled, generalizable alignment technique for LRM safety with practical implications for deploying LRMs in safety-critical contexts.

Abstract

Large Reasoning Models (LRMs) have achieved tremendous success with their chain-of-thought (CoT) reasoning, yet also face safety issues similar to those of basic language models. In particular, while algorithms are designed to guide them to deliberately refuse harmful prompts with safe reasoning, this process often fails to generalize against diverse and complex jailbreak attacks. In this work, we attribute these failures to the generalization of the safe reasoning process, particularly their insufficiency against complex attack prompts. We provide both theoretical and empirical evidence to show the necessity of a more sufficient safe reasoning process to defend against advanced attack prompts. Building on this insight, we propose a Risk-Aware Preference Optimization (RAPO) framework that enables LRM to adaptively identify and address the safety risks with appropriate granularity in its thinking content. Extensive experiments demonstrate that RAPO successfully generalizes multiple LRMs' safe reasoning adaptively across diverse attack prompts whilst preserving general utility, contributing a robust alignment technique for LRM safety. Our code is available at https://github.com/weizeming/RAPO.
Paper Structure (26 sections, 2 theorems, 17 equations, 7 figures, 8 tables, 1 algorithm)

This paper contains 26 sections, 2 theorems, 17 equations, 7 figures, 8 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Large Language Model Safety
  4. Safe Reasoning for LRMs
  5. Motivation
  6. A Unified View through In-Context Alignment
  7. Theoretical Illustration
  8. Empirical Observations
  9. Risk-Aware Preference Optimization
  10. Overview
  11. SFT Warm-up for Format Alignment
  12. Reward Design: Risk-Aware and General
  13. Implementation Details and Discussion
  14. Experiment
  15. Experiment Set-up
  16. ...and 11 more sections

Key Result

Theorem 3.1

Given prompt $x_0=\frac{1}{k+1}(c_0+c_1+\cdots +c_k)$ and reasoning traces $\{z_i\}_{i=1}^T$, the number of safe reasoning trace should be at least $t=\Omega(k)$ to refuse the prompt.

Figures (7)

  • Figure 1: Illustrations on effectiveness and failure modes of safe reasoning.
  • Figure 2: An outline of our RAPO framework.
  • Figure 3: Reward with simple moving average during training of Qwen (1.7B) and DeepSeek models.
  • Figure 4: System prompt for the Safe Reasoning Generator.
  • Figure 5: System prompt for the Risk-Aware Reward Judge.
  • ...and 2 more figures

Theorems & Definitions (2)

  • Theorem 3.1
  • Proposition A.1: Restated from von2023transformers