STAIR: Improving Safety Alignment with Introspective Reasoning

• LLMs now support applications including medical diagnostics, education, and legal consulting, where harmful content creates serious safety concerns.Source: §1 · p. 1
• Safety alignment methods such as SFT, RLHF, and DPO can compromise general performance because their objectives conflict.Source: §1 · pp. 1–2
• Direct-refusal training remains limited in complex scenarios because jailbreaks can conceal threats and bypass shallow alignment.Source: §1 · p. 2
• STAIR uses CoT analysis and safety-aware System 2 reasoning, then applies SI-MCTS-generated stepwise preference optimization and PRM-guided test-time search.Source: §1 · p. 2
• 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.Source: §1 · p. 2

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.Source: §2.1 · pp. 2–3, §2 · p. 2

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.Source: §2.1 · p. 3, Figure 2 · p. 3

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.Source: §2.2 · pp. 3–4
• 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.Source: §2.2 · p. 4

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.Source: §2.2 · p. 4

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.Source: §3.1 · p. 5

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.Source: §3.2 · p. 5, §3.1 · p. 5
• 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.Source: §3.2 · p. 6
• 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.Source: §3.2 · p. 6

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.Source: §3.3 · p. 6

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.Source: §3.4 · pp. 6–7
• 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.Source: §3.4 · p. 7

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.Source: §3.4 · p. 8

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

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

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

• The work discusses jailbreak attacks as ethically sensitive material but frames its objective as mitigating societal and ethical risks through safer LLM alignment.Source: Impact Statement · p. 10

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.Source: §A.1 · p. 15

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.Source: §A.1 · p. 15

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.Source: §A.2 · p. 15

Safety-aware reasoning

• The structured instructions require risk identification and mitigation for sensitive queries while emphasizing logical, solution-oriented responses for harmless queries.Source: §A.2 · pp. 15–16

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.Source: §A.2 · pp. 15–16

Desired properties

• The reward design requires safe responses to outrank unsafe responses regardless of helpfulness.Source: §B.1 · p. 16
• For safe responses, helpfulness increases reward; for unsafe responses, more helpful harmful instructions decrease reward.Source: §B.1 · p. 16
• When either safety or helpfulness is held constant, reward differences preserve the corresponding single-objective differences, reducing selection to standard MCTS behavior.Source: §B.1 · p. 16

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.Source: §B.1 · p. 17

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.Source: Necessity · pp. 17–18

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.Source: §B.2.1 · p. 18

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.Source: §B.2.2 · pp. 18–19

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.Source: §B.2.3 · p. 19

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.Source: §C.2 · p. 20

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.Source: §C.2 · p. 20

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.Source: §C.2 · pp. 20–21

Case format

• Each case presents an original harmful prompt, a jailbroken version with the baseline answer, and STAIR’s reasoning process and answer.Source: D. Examples · p. 21

Presentation

• The examples redact harmful content and mark reasoning steps and final answers with dedicated tokens; STAIR may use one or multiple reasoning steps.Source: D. Examples · p. 21

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.Source: D. Examples · p. 21