Research map / Large language models / Reasoning and chain of thought
STAIR: Improving Safety Alignment with Introspective Reasoning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes STAIR, which teaches LLMs structured chain-of-thought safety reasoning, then iteratively improves safety via step-level DPO on data from a safety-informed MCTS with a theoretically derived reward, and trains a process reward model for test-time search.
Problem
Safety alignment via direct refusals is shallow and easily bypassed by jailbreak attacks, and existing methods suffer safety-helpfulness trade-offs.
Approach
Integrates introspective step-by-step safety reasoning into alignment: a safety-informed reward (R = S*H + 2S, derived from three properties) guides MCTS over reasoning steps, producing stepwise preference pairs for iterative DPO without external annotation, plus a PRM for best-of-N and beam search at inference.
Key results
- On StrongReject, STAIR-DPO-3 on LLaMA-3.1-8B reaches goodness 0.8798 vs 0.4054 base, and Qwen-2-7B reaches 0.8486.
- With Best-of-8 / beam search test-time scaling, LLaMA-based STAIR reaches average StrongReject goodness 0.9391, comparable to Claude-3.5 (0.9359).
- AlpacaEval win rate vs GPT-4 for LLaMA rises from 25.55% (base) to 38.66% after three iterations.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add structured CoT alignment and three SI-MCTS preference-optimization iterations | LLaMA-3.1-8B-Instruct · StrongReject | goodness score: 0.4054 -> 0.8798 | Helped |
| Add structured CoT alignment and three SI-MCTS preference-optimization iterations | Qwen-2-7B-Instruct · StrongReject | goodness score: 0.3808 -> 0.8486 | Helped |
| Add structured CoT alignment and three SI-MCTS preference-optimization iterations | LLaMA-3.1-8B-Instruct · AlpacaEval | winning rate against GPT-4: 25.55% -> 38.66% | Helped |
| Add structured CoT alignment and three SI-MCTS preference-optimization iterations | LLaMA-3.1-8B-Instruct · GSM8k | accuracy: 85.60% -> 87.64% | Helped |
| Add structured CoT alignment and three SI-MCTS preference-optimization iterations | LLaMA-3.1-8B-Instruct · GSM8k | accuracy: 85.60% -> 87.64% | Hurt |
| Use stepwise rather than full-trajectory preference data (ablation) | STAIR-SFT on LLaMA-3.1-8B-Instruct · StrongReject | goodness score: 0.6663 -> 0.6955 | Helped |
| Use stepwise rather than full-trajectory preference data (ablation) | STAIR-SFT on LLaMA-3.1-8B-Instruct · AlpacaEval | winning rate against GPT-4: 32.87% -> 34.48% | Helped |
| Train iteratively across three self-improvement rounds rather than on the full dataset in one iteration (ablation) | STAIR on LLaMA-3.1-8B-Instruct · StrongReject | goodness score: 0.7342 -> 0.8798 | Helped |
| Train iteratively across three self-improvement rounds rather than repeat first-split data for 3× epochs (ablation) | STAIR on LLaMA-3.1-8B-Instruct · StrongReject | goodness score: 0.6745 -> 0.8798 | Helped |
| Use PRM Best-of-4 inference (ablation) | STAIR-SFT on LLaMA-3.1-8B-Instruct · AlpacaEval | winning rate against GPT-4: 30.02% -> 30.95% | Helped |
| Use self-rewarding Best-of-4 inference instead of PRM Best-of-4 (ablation) | STAIR-DPO-1 on LLaMA-3.1-8B-Instruct · AlpacaEval | winning rate against GPT-4: 33.11% -> 32.03% | Hurt |
| Change the safety-informed reward form among the tested alternatives (ablation) | STAIR-SFT on LLaMA-3.1-8B-Instruct · StrongReject | goodness score: 0.6712; 0.6633; 0.6753; 0.6688 | No clear effect |
Surprising findings
- Plain CoT prompting does not improve safety (StrongReject 0.3790 vs base 0.4054 on LLaMA).
- SACPO degrades SimpleQA and InfoFlow sharply despite safety gains, reflecting knowledge loss and over-refusal.
Limitations
- Test-time search multiplies inference cost roughly with search budget.
- Reliance on self-rewarding by the same model may be noisy; judge is GPT-4o/self-rating.
- Gains on jailbreak evaluation depend on a fixed set of attack methods (PAIR, PAP).
Experimental setup
Datasets: PKU-SafeRLHF JailbreakV-28k UltraFeedback StrongReject XSTest WildChat Do-Not-Answer SimpleQA InfoFlow AdvGLUE GSM8K AlpacaEval 2.0
Benchmarks: StrongReject XSTest WildChat Do-Not-Answer Stereotype SimpleQA InfoFlow AdvGLUE GSM8K AlpacaEval 2.0 BIG-bench HHH
Metrics: goodness score refusal rate accuracy winning rate correlation coefficient
Baselines: CoT prompting SFT DPO SACPO Self-Rewarding GPT-4o Claude-3 Claude-3.5 Qwen-Max Gemini-1.5 DeepSeek-R1 LLaMA-o1
Builds on: DPO Self-Rewarding SACPO Monte Carlo Tree Search chain-of-thought
Closest papers by meaning
- SAFEPATH: Preventing Harmful Reasoning in Chain-of-Thought via Early Alignment (NeurIPS 2025 · similarity 0.92)
- Safety Alignment Should be Made More Than Just a Few Tokens Deep (ICLR 2025 · similarity 0.92)
- Towards Safe Reasoning in Large Reasoning Models via Corrective Intervention (ICLR 2026 · similarity 0.91)
- Jailbreak Antidote: Runtime Safety-Utility Balance via Sparse Representation Adjustment in Large Language Models (ICLR 2025 · similarity 0.91)
- RSafe: Incentivizing proactive reasoning to build robust and adaptive LLM safeguards (NeurIPS 2025 · similarity 0.91)
- CodeMimicry: Exploiting Safety Generalization Lag in Large Language Models via Structured Code Completion (NeurIPS 2026 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.