PaperlayerView exact revision on arXiv

Source-linked AI summary

STAIR: Improving Safety Alignment with Introspective Reasoning

Yichi Zhang, Siyuan Zhang, Yao Huang, Zeyu Xia, Zhengwei Fang, Xiao Yang, Ranjie Duan, Dong Yan, Yinpeng Dong, Jun Zhu

arXiv:2502.02384v2cs.CL

TL;DR

Existing safety alignment can sacrifice general performance and remain vulnerable to jailbreaks because it relies heavily on direct refusals. STAIR trains safety-aware introspective reasoning through structured CoT, SI-MCTS-generated stepwise optimization, and process-reward-guided test-time search. Experiments report stronger safety with preserved or improved helpfulness, including jailbreak resistance comparable to Claude-3.5 under test-time scaling.

  • Problem

    Existing safety alignment methods face safety-performance trade-offs and remain susceptible to jailbreak attacks that bypass direct refusals.

  • Method

    STAIR combines structured CoT alignment, iterative step-level DPO on SI-MCTS-generated data, and PRM-guided test-time search using safety and helpfulness signals.

  • Results

    STAIR consistently improves resistance to harmful queries while improving multiple general-performance dimensions, reaching a 0.94 StrongReject goodness score comparable to Claude-3.5 with test-time scaling.

  • Takeaways & Limitations

    Safety-aware reasoning enables more balanced safety and general performance than instinctive alignment strategies across the reported benchmarks.

  • Takeaways & Limitations

    Test-time search increases inference costs approximately in proportion to the search budget, although deployment can adjust or omit this optional practice.

Abstract

from arXiv · show

Ensuring the safety and harmlessness of Large Language Models (LLMs) has become equally critical as their performance in applications. However, existing safety alignment methods typically suffer from safety-performance trade-offs and the susceptibility to jailbreak attacks, primarily due to their reliance on direct refusals for malicious queries. In this paper, we propose STAIR, a novel framework that integrates SafeTy Alignment with Itrospective Reasoning. We enable LLMs to identify safety risks through step-by-step analysis by self-improving chain-of-thought (CoT) reasoning with safety awareness. STAIR first equips the model with a structured reasoning capability and then advances safety alignment via iterative preference optimization on step-level reasoning data generated using our newly proposed Safety-Informed Monte Carlo Tree Search (SI-MCTS). We further train a process reward model on this data to guide test-time searches for improved responses. Extensive experiments show that STAIR effectively mitigates harmful outputs while better preserving helpfulness, compared to instinctive alignment strategies. With test-time scaling, STAIR achieves a safety performance comparable to Claude-3.5 against popular jailbreak attacks. Relevant resources in this work are available at https://github.com/thu-ml/STAIR.

1. Introduction

LLMs are increasingly used in high-stakes applications, making harmful outputs a serious concern. Existing safety alignment often relies on direct refusals, which can trade off general performance and remain vulnerable to jailbreaks; STAIR addresses this with introspective reasoning.

  • LLMs now support applications including medical diagnostics, education, and legal consulting, where harmful content creates serious safety concerns.
  • Safety alignment methods such as SFT, RLHF, and DPO can compromise general performance because their objectives conflict.
  • Direct-refusal training remains limited in complex scenarios because jailbreaks can conceal threats and bypass shallow alignment.
  • STAIR uses CoT analysis and safety-aware System 2 reasoning, then applies SI-MCTS-generated stepwise preference optimization and PRM-guided test-time search.
  • 0.88 goodness on StrongReject for LLaMA exceeded the best baseline by 0.15, while AlpacaEval winning rates increased by 13.11 for LLaMA and 6.25 for Qwen.

2. Safety Alignment with Introspective Reasoning

STAIR improves safety alignment by having models reason about risks step by step rather than relying on direct refusals. It combines structured chain-of-thought training with Safety-Informed MCTS, safety-aware rewards, and iterative step-level preference optimization.

  • Motivation: STAIR targets safety-aware reasoning to help models analyze risky queries instead of directly refusing them.The motivation is to preserve general performance while improving safety awareness.
  • Structured reasoning: The framework first fine-tunes models on structured chain-of-thought data containing titled reasoning steps and final answers.GPT-4o rewrites responses into n-step reasoning paths, which are used for supervised fine-tuning.
  • Safety-Informed MCTS: Safety-Informed MCTS treats reasoning steps as search nodes and assigns rewards based on both helpfulness and safety.Safety information is backpropagated through the search tree to select safety-aware reasoning data.
  • Safety-Informed MCTS: The reward design prioritizes safe responses while retaining helpfulness-sensitive behavior and reducing to single-objective rewards when appropriate.Its stated properties include safety priority, dual monotonicity of helpfulness, and degeneration to a single objective.
  • Iterative optimization: Search trees yield threshold-filtered stepwise preference pairs for step-level DPO, and the model repeatedly regenerates data for iterative self-improvement.The process uses self-rewarding rather than external evaluators and can be repeated across iterations.

3. Experimental Results

STAIR is evaluated across safety, helpfulness, reasoning, robustness, and privacy-oriented benchmarks, with main results and targeted ablations examining its design and computational costs.

  • Experimental setup: STAIR is evaluated on 10 benchmarks covering harmlessness and general performance using LLaMA-3.1-8B-Instruct and Qwen-2-7B-Instruct.The evaluation includes safety datasets such as StrongReject, XsTest, WildChat, and Do-Not-Answer, alongside general-capability benchmarks.
  • Main results: Baseline SFT and DPO exhibit safety-performance trade-offs, while SACPO mitigates trade-offs but degrades factual knowledge and benign privacy-query performance.The comparison also finds that CoT prompting improves reasoning-dependent tasks without enhancing safety, while Self-Rewarding benefits apparently harmful-query refusal behavior.
  • Main results: STAIR produces more balanced improvements across benchmarks, including stronger harmful-input resilience while maintaining or improving general performance.LLaMA-3.1’s AlpacaEval winning rate reaches 38.66%, versus 25.55% for the base model, after self-improvement.
  • Main results: 0.8798 and 0.8486 are the eventual StrongReject goodness scores for the two evaluated LLMs after STAIR training.The reported results attribute these gains to integrating reasoning with safety alignment.
  • Test-time scaling: Test-time scaling with Best-of-N or Beam Search improves both safety and helpfulness, adding 0.06 StrongReject goodness and more than 3.0% AlpacaEval winning rate.The extra computation is measured relative to one-time greedy decoding using logarithmic generated-step cost.
  • Ablation studies: Varying the safety-data ratio reveals a safety-helpfulness trade-off, yet both dimensions consistently outperform the base model with structured CoT training.This study is conducted during CoT format alignment and reports the trend in Figure 5.
  • Ablation studies: Ablations support step-level optimization, iterative training, and flexible reward functions as contributors to safety-aware reasoning improvements.Stepwise data provides finer supervision, later iterations yield higher-quality data, and alternative reward forms improve both safety and general capabilities with insignificant differences.
  • Computation costs: STAIR’s training is annotation-free and takes approximately 30 hours on 8 A800 GPUs, while SI-MCTS tree construction averages 15 seconds per prompt.The average valid preference-pair time is about 0.47 seconds, compared with 0.40 seconds for self-rewarding full-trajectory sampling; test-time search adds further inference cost.

4. Discussions

STAIR is compared with open-source reasoning models, deliberative-alignment models, and commercial LLMs on safety benchmarks. Its safety-aware iterative reasoning reaches a level comparable to Claude-3.5 and can surpass it with test-time scaling.

  • Unlike Deliberative Alignment, STAIR does not require a powerful reasoning foundation model and instead integrates safety-aware reasoning into normal instruction-tuned LLMs.
  • STAIR is compared with open-source reasoning LLMs and models trained with Deliberative Alignment across multiple benchmarks.
  • Most compared commercial LLMs exceed 0.95 goodness on straightforward harmful questions, but several remain vulnerable to modern jailbreak attacks.
  • Claude-3.5 provides the strongest commercial defense against the evaluated jailbreak attacks.
  • STAIR consolidates LLaMA to a safety level comparable to Claude-3.5 and surpasses it with test-time scaling.

5. Related Work

Prior safety-alignment methods use SFT, DPO, or RLHF but can trade general performance for harmlessness. STAIR incorporates reasoning into safety alignment without requiring access to a large reasoning model.

  • SFT, DPO, and RLHF are established safety-alignment techniques, but harmlessness and helpfulness objectives can conflict.
  • STAIR uses fine-tuning to incorporate reasoning about potential risks into safety alignment.
  • Deliberative Alignment also uses reasoning for safety but assumes access to a large reasoning model, whereas STAIR does not.
  • Chain-of-thought and tree-of-thought are prompting-based techniques that preceded later work on LLM reasoning and self-improvement.

6. Conclusion

STAIR introduces System 2 thinking into safety alignment so models can analyze complex risks while maintaining general performance. It combines structured CoT warm-up, SI-MCTS-based self-improvement, process rewards, and test-time scaling.

  • STAIR introduces System 2 thinking to distinguish potential safety risks in complex scenarios such as jailbreaks while maintaining general performance.
  • The framework begins with structured CoT data, then iteratively self-improves on SI-MCTS-generated stepwise data using safety and helpfulness signals.
  • STAIR also trains a process reward model from the same search trees and evaluates test-time scaling for safety alignment.

Impact Statement

The paper addresses ethically sensitive jailbreak material while stating that its objective is to improve LLM safety alignment and mitigate societal and ethical risks.

  • The work discusses jailbreak attacks as ethically sensitive material but frames its objective as mitigating societal and ethical risks through safer LLM alignment.

A.1. Dataset Summary

The dataset combines safety and helpfulness examples and supports structured chain-of-thought training for safe, logically organized responses.

  • Dataset composition: The seed dataset contains 50k pairwise samples combining 25k UltraFeedback helpfulness examples with 22k PKU-SafeRLHF safety examples and 3k jailbreak prompts.Helpfulness responses are selected by ratings, while safety prompts are filtered for genuinely toxic content; jailbreak prompts are decontaminated and annotated with GPT-4o-generated refusals.
  • Training use: Baseline training uses chosen responses or pairwise data, whereas STAIR and self-rewarding use only dataset prompts rather than directly training on GPT-generated jailbreak refusals.
  • Structured CoT synthesis: Structured CoT data are synthesized from 10k prompts each from PKU-SafeRLHF and UltraFeedback using GPT-4o-generated reasoning steps.Incorrectly formatted or unsafe responses are filtered before training.
  • Safety-aware reasoning: The structured instructions require risk identification and mitigation for sensitive queries while emphasizing logical, solution-oriented responses for harmless queries.
  • Structured CoT format: The response format separates Problem Analysis, multiple Reasoning sections, and a Final Answer using dedicated reasoning-step and output tokens.This format is intended to organize reasoning for both sensitive and harmless queries.

B.1. Derivation of Safety-Informed Reward

The safety-informed reward formalizes safety priority, safety-sensitive helpfulness, and preservation of standard single-objective MCTS behavior. Its admissible family is R(H, S) = F(S) · H + S − C2 · F(S) + c.

  • Desired properties: The reward design requires safe responses to outrank unsafe responses regardless of helpfulness.
  • Desired properties: For safe responses, helpfulness increases reward; for unsafe responses, more helpful harmful instructions decrease reward.
  • Desired properties: When either safety or helpfulness is held constant, reward differences preserve the corresponding single-objective differences, reducing selection to standard MCTS behavior.
  • Derivation: R(H, S) = F(S) · H + S − C2 · F(S) + c is necessary and sufficient for the stated dual-monotonicity and single-objective properties.The function satisfies F(0) = 0, F(C1) = 1, F(S) > 0 for S > 0, and F(S) < 0 for S < 0.
  • Corollary: A simple instantiation, F(S) = S, C1 = 1, C2 = −1, and c = 0, yields R(H, S) = S · H + 2S and guarantees any safe answer outranks any unsafe answer.
  • Safety-Informed MCTS: SI-MCTS applies safety-informed rewards within standard MCTS, using UCB1 values that balance exploration and exploitation.The implementation uses c = 1.5, search budget n = 200, and m = 4 children by default.
  • Self-rewarding: Self-rewarding judges helpfulness on a 1–5 scale and safety as safe or unsafe, normalizing ratings to [−1, 1] and sometimes adding rule-based refusal rewards.The trained model evaluates its own rolled-out answers, with reasoning included to improve rating reliability.
  • Preference data construction: Stepwise preference pairs are selected when their value gap exceeds v0 and the higher-valued response exceeds v1.These thresholds ensure a meaningful quality difference and an adequately good chosen response.

C.2. Evaluation Details

Evaluation uses reproducible generation and benchmarks safety, trustworthiness, robustness, privacy awareness, and helpfulness, including no-jailbreak and jailbreak settings.

  • Generation: Default evaluation uses greedy decoding, while test-time scaling uses temperature 0.6, top-p 0.9, and top-k 50 for response diversity.
  • Safety evaluation: StrongReject reports a goodness score of 1 − rubric score, evaluated on ordinary prompts and PAIR and PAP-Misrepresentation jailbreaks.GPT-4o applies the official rubric-based evaluation protocol.
  • General evaluation: General evaluation covers truthfulness with SimpleQA, adversarial robustness with AdvGLUE, privacy awareness with InfoFlow, and helpfulness with GSM8k, AlpacaEval, and BIG-bench HHH.InfoFlow uses correlation coefficient, while AlpacaEval reports winning rate against GPT-4.

D. Examples

Qualitative cases compare baseline DPO responses with STAIR responses under original and jailbroken harmful prompts. STAIR’s reasoning examines hidden risks and ends in appropriate refusals.

  • Case format: Each case presents an original harmful prompt, a jailbroken version with the baseline answer, and STAIR’s reasoning process and answer.
  • Presentation: The examples redact harmful content and mark reasoning steps and final answers with dedicated tokens; STAIR may use one or multiple reasoning steps.
  • Observed behavior: Baseline models can refuse direct harmful prompts but remain vulnerable to jailbreaks that fabricate scenarios, whereas STAIR-DPO-3 analyzes underlying risks and appropriately refuses.
Loading 2502.02384v2…