STAIR: Improving Safety Alignment with Introspective Reasoning

STAIR: Improving Safety Alignment with Introspective Reasoning

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

Conference: ICML 2025

YouTube: https://slideslive.com/39044085

Overview

In an era where large language models (LLMs) are rapidly integrating into critical applications—from medical advice to policy drafting—concerns regarding their safety and trustworthiness have escalated. Despite their powerful capabilities, LLMs can, and often do, generate harmful or illegal content, especially when subjected to sophisticated "jailbreaking" attacks. The talk "STAIR: Improving Safety Alignment with Introspective Reasoning," presented by Yichi Zhang from Tsinghua University and their co-authors, introduces a novel framework designed to address these pressing safety issues without compromising the models' utility.

The STAIR framework proposes a paradigm shift from reactive, System 1-like refusal mechanisms to a more deliberate, System 2-like introspective reasoning for safety alignment. Current alignment methods, largely reliant on refusal datasets, often lead to "shallow alignment" and a problematic safety-performance trade-off. STAIR aims to enable LLMs to self-check, analyze underlying risks, and decide whether to answer, refuse, or reformulate a query, thereby unifying safety with general reasoning capabilities. This innovative approach seeks to build LLMs that are not only robust against advanced adversarial prompts but also maintain their helpfulness across a wide range of tasks.

This article delves into the technical contributions of STAIR, exploring its three-stage training process involving structured chain-of-thought reasoning, safety-informed Monte Carlo Tree Search (MCTS) with self-improvement, and the role of test-time scaling. It also examines the empirical evidence supporting STAIR's efficacy in significantly enhancing safety while preserving or even improving general model capabilities, effectively pushing the safety-helpfulness frontier outward.

Background

▶ Watch: Introduction: LLM safety issues and challenges (0:00)

The rapid advancement and widespread adoption of large language models (LLMs) have brought unprecedented capabilities, but also significant challenges, particularly concerning their safety and trustworthiness. LLMs, despite their potential for societal benefit, have been shown to generate harmful, biased, or illegal content. This problem is exacerbated by jailbreaking attacks, which are carefully crafted prompts designed to bypass safety filters and elicit undesirable outputs. Early examples, such as "grandma roleplay" prompts or more recent sophisticated techniques like "napalm attacks," demonstrate how easily internal safety mechanisms can be subverted, eroding public trust.

Existing safety alignment methods, such as Direct Preference Optimization (DPO) and Reinforcement Learning from Human Feedback (RLHF), predominantly rely on refusal datasets. These datasets typically pair malicious queries with responses that instruct the model to "say sorry" or explicitly refuse the request. While widely adopted, these methods exhibit several critical limitations. Firstly, aligned LLMs remain vulnerable to advanced jailbreaking attacks. Research, highlighted by an ICLR outstanding paper, has identified this as "shallow alignment," where models primarily shift their output distribution only on the initial tokens, still complying with malicious instructions once the first few tokens are bypassed. This suggests a superficial understanding of safety rather than a deep, intrinsic awareness.

Secondly, an overemphasis on refusal often leads to a detrimental safety-performance trade-off. A model that constantly responds with apologies, while theoretically "safe," becomes practically useless. Studies on RLHF for harmlessness have shown that the ratio of harmlessness data can significantly impact the original model's helpfulness. The core problem, then, is how to make LLMs genuinely safer, especially under adversarial conditions, without sacrificing their utility and helpfulness.

To address this, the STAIR framework draws inspiration from human cognition, specifically the distinction between System 1 and System 2 thinking. System 1 is fast, intuitive, and reactive, akin to how current alignment methods instantly provide refusal reactions to suspicious inputs. System 2, in contrast, is slow, deliberate, and analytical, involving introspection and careful reasoning before responding. The intuition behind STAIR is to treat safety not as a reflex, but as a System 2-like deliberation. This approach not only enables more explicit analysis of potential risks but also reframes risk detection as a reasoning task, integrating it seamlessly with other general LLM capabilities like math or coding. The ultimate goal is to empower LLMs with introspection to identify hidden risks and decide whether to answer, refuse, or reformulate a user query, thereby leveraging the power of LLM reasoning for robust safety alignment.

Key Findings

▶ Watch: STAIR's core idea: System 2-like deliberation for safety (2:00)

The STAIR framework delivers compelling results that fundamentally advance the state of LLM safety alignment. Its core contributions and discoveries revolve around significantly enhancing safety under adversarial conditions while simultaneously preserving, and often improving, general model utility.

One of the most striking findings is the substantial improvement in safety scores against advanced jailbreaking techniques. Through its self-improvement loop, STAIR roughly doubles the safety score on StrongReject, a benchmark specifically designed for the most sophisticated jailbreaking attacks. This demonstrates a robust resistance to malicious inputs that far surpasses traditional alignment baselines. Crucially, this enhanced safety does not come at the expense of helpfulness; instead, the utility of STAIR-aligned models remains competitive and often improves across a diverse set of general benchmarks covering truthfulness, conversation, mathematical reasoning, and robustness. This effectively "moves the safety-helpfulness frontier outward," achieving a Pareto improvement in both dimensions.

Another key finding highlights the critical role of test-time scaling in further bolstering safety. The research shows that increasing test-time computations, through techniques like best-of-N sampling or beam search guided by a process reward model (PRM), yields additional safety gains. For instance, an 8-billion parameter model, when subjected to beam search, achieved StrongReject safety scores comparable to, and sometimes even exceeding, larger commercial models like Claude 3.5. This underscores that deliberate reasoning and search at inference time can unlock significant safety improvements, not just accuracy.

The work also establishes that reasoning is a powerful, generalizable mechanism for safety alignment, extending beyond its conventional applications in math or coding. By enabling LLMs to engage in "extra thinking," STAIR allows them to inherently notice and analyze where a request might be risky or malicious. This self-improved safety-aware reasoning fosters a more deliberate and introspective model behavior over time.

Finally, the STAIR framework demonstrates its generalizability and scalability by applying the same principles to larger reasoning models. When fine-tuning DeepSeek-R1 (ranging from 15B up to 32B parameters) with approximately 15,000 safety-aware reasoning trajectories, the models showed significantly enhanced resistance to malicious queries. Importantly, their average scores on typical reasoning benchmarks like AIME or GPQA Diamond remained essentially unchanged, affirming that the safety enhancements are achieved without degrading core reasoning capabilities. This consistent pattern across different model scales and architectures validates STAIR's foundational approach to integrating introspective reasoning for robust and utilitarian safety alignment.

Technical Deep Dive

▶ Watch: Overview of the STAIR three-stage framework (3:15)

The STAIR framework introduces a sophisticated, multi-stage approach to imbue LLMs with System 2-like introspective reasoning for safety alignment. This technical deep dive explores the architecture, algorithms, and training methodologies that underpin STAIR's success.

The core idea is to move beyond instantaneous, reactive refusals by enabling LLMs to perform self-checks and deliberate on potential risks before formulating a response. This allows the model to decide whether to directly answer, refuse, or reformulate a user query, thereby integrating safety consideration directly into its reasoning process. The STAIR framework is structured into three distinct, sequential stages:

1. Structured Chain-of-Thought (CoT) Reasoning Warm-up:

The initial stage focuses on establishing a foundation for structured reasoning within the LLM. The method employs GPT-4 to generate supervised samples of chain-of-thought (CoT) reasoning. This CoT format is predefined, with each reasoning step explicitly wrapped by special tokens. This token-based encapsulation serves two primary purposes:

• Readability and Interpretability: It makes the model's "thinking process" transparent and easier to follow for human inspection.
• Node Separation for Search: Crucially, it cleanly separates individual reasoning steps into distinct "nodes," which are essential for the subsequent safety-aware search process.

The training data for this stage includes a mix of both harmful and normal prompts, ensuring the model learns to be both helpful and safe from the outset, embedding this dual awareness into its reasoning procedures.

2. Self-Improvement with Safety-Informed Monte Carlo Tree Search (MCTS):

This is the central stage where the model iteratively refines its safety-aware reasoning capabilities. Each structured reasoning step from Stage 1 effectively becomes a node within a search tree. Given a user prompt, the model can:

• Expand Alternative Next Steps: Explore different reasoning paths from a given node.
• Roll Out Candidate Final Answers: Generate potential complete responses based on these paths.
• Score Each Path: Critically, each reasoning trajectory is evaluated for both its safety and helpfulness using a novel safety-informed reward function.

The safety-informed reward function is a cornerstone of STAIR. Unlike traditional MCTS that might only evaluate correctness or winning rates at the end of a rollout, STAIR's reward function explicitly rates safety (S) and helpfulness (H) separately. The goal is to enhance LLM safety without any loss in utility. To achieve this, a theoretically grounded function R(S, H) was derived, satisfying three key properties:

• Safety as Priority: A safe answer always receives a higher reward than an unsafe one, regardless of helpfulness.
• Dual Monotonicity of Helpfulness: The expectation of helpfulness and detailed information in a response is contingent upon its safety. If an answer is unsafe, helpfulness is penalized, and vice-versa.
• Degeneration to Single-Objective Condition: The search process should revert to traditional MCTS behavior if only helpfulness is considered, ensuring compatibility with existing optimization paradigms.

Based on these properties, a simple yet effective reward function form was derived. This function is designed to backpropagate safety-related information into the internal nodes of the search tree, enabling a fine-grained optimization of the model's reasoning process.

The model then engages in an iterative self-improvement loop. Based on the safety and helpfulness scores, threshold sampling is applied to identify "winning" and "losing" reasoning steps. This process generates a step-level preference dataset. The model is then updated using Direct Preference Optimization (DPO) on this dataset. This loop is repeated for K iterations (K=3 in the experiments), gradually strengthening the model's ability to resist malicious inputs while preserving its utility.

3. Test-Time Scaling Verification:

The final stage investigates the impact of test-time scaling techniques on safety alignment. From the search paths collected during Stage 2, a process reward model (PRM) is trained. At inference time, this PRM is used to score both intermediate reasoning steps and complete candidate answers. The model generates multiple candidate responses using techniques like best-of-N sampling or beam search. By applying the PRM to select the highest-scoring response, the framework can measure the additional safety gains achieved purely from test-time computation, independent of changes to the underlying model architecture. This stage verifies that deliberate search and evaluation at inference can further enhance safety and accuracy.

In summary, STAIR meticulously combines structured reasoning, reinforcement learning with a custom reward function, and iterative self-improvement, culminating in a robust framework for building LLMs that are not just superficially aligned but deeply introspective about safety.

Experimental Setup & Results

▶ Watch: Detailed explanation of safety-informed MCTS reward function (4:50)

The STAIR framework was rigorously evaluated to demonstrate its effectiveness in improving safety alignment and maintaining general capabilities. The experiments were conducted on two undisclosed open-source LLM backbones, with an 8-billion parameter model being specifically mentioned for certain test-time scaling experiments.

Baselines: The performance of STAIR was compared against common alignment baselines, which typically rely on refusal datasets and standard DPO or RLHF methods.

Benchmarks: The evaluation spanned a diverse set of benchmarks designed to assess both harmlessness (safety) and general capabilities:

• Harmlessness: The primary safety benchmark was StrongReject, a challenging dataset specifically designed to test resistance against the most advanced jailbreaking techniques. Safety scores on StrongReject were a key metric.
• General Capabilities: To ensure helpfulness was preserved, models were evaluated on various general benchmarks, including those for truthfulness, conversation, mathematical reasoning, and robustness. Metrics included Alpaca winning rates (presumably for helpfulness in conversational settings) and performance on specialized reasoning benchmarks like AIME and GPQA Diamond.

Headline Numbers and Key Findings:

The experimental results unequivocally supported the efficacy of the STAIR framework:

• Safety Improvement: As the self-improvement loop progressed through its K=3 iterations, safety scores on StrongReject rose steadily. By the end of the third DPO iteration, STAIR models roughly doubled their safety score on StrongReject compared to initial aligned baselines. This represents a substantial gain in resistance to sophisticated malicious queries.
• Helpfulness Preservation: Crucially, this significant safety improvement did not lead to a collapse in helpfulness. Across all general capability benchmarks (truthfulness, conversation, math, robustness), STAIR-trained models remained competitive and often demonstrated improved performance. This confirmed that STAIR successfully "moves the safety-helpfulness frontier outward," achieving gains in both dimensions simultaneously.
• Impact of Test-Time Scaling: Further experiments verified the contribution of test-time scaling. Increasing test-time computations, through techniques like best-of-N sampling or beam search guided by the process reward model (PRM), consistently improved both StrongReject safety scores and Alpaca winning rates. Notably, an 8-billion parameter model utilizing beam search achieved StrongReject safety scores that were comparable to, and in some cases even surpassed, those of proprietary models like Claude 3.5. This highlights the power of deliberate test-time search in enhancing safety without requiring larger base models.

Ablation Studies: The researchers also conducted ablation studies to understand the contribution of individual components within the STAIR framework. These studies investigated the impact of elements such as the iterative step-level optimization, the specific form of the safety-informed reward function, and the cost associated with test-time computations. While specific results were referred to the full paper, the presence of these studies indicates a thorough analysis of the framework's design choices.

Extension to Reasoning Models: The generalizability of the STAIR approach was demonstrated by applying it to large reasoning models. Approximately 15,000 safety-aware reasoning trajectories were generated and used to fine-tune DeepSeek-R1 models, ranging from 15 billion to 32 billion parameters. The results showed a consistent pattern: average scores across typical reasoning benchmarks (AIME, GPQA Diamond) remained essentially unchanged, while resistance to malicious queries was significantly enhanced. This reinforced the idea that introspective reasoning for safety can be effectively integrated into various LLM architectures without degrading their core functionalities.

In summary, STAIR presents a robust experimental validation, demonstrating its ability to deliver substantial safety gains against advanced jailbreaking, preserve and often enhance general helpfulness, and leverage test-time scaling for further improvements, all while being applicable across different model sizes and architectures.

Practical Implications

▶ Watch: Properties of the proposed safety-informed reward function (5:15)

The STAIR framework presents several significant practical implications for various stakeholders involved in the development, deployment, and utilization of large language models. Its emphasis on introspective reasoning offers a promising path forward for addressing the complex challenges of LLM safety.

For Practitioners and Infra Teams:

• Robust Safety Alignment: STAIR provides a blueprint for building LLMs that are inherently more robust against sophisticated jailbreaking attacks. This means infrastructure teams can deploy models with higher confidence in their safety filters, reducing the risk of generating harmful or inappropriate content in production environments.
• Reduced Safety-Helpfulness Trade-off: The ability of STAIR to improve safety without sacrificing helpfulness is a critical breakthrough. Practitioners no longer need to make severe compromises between model utility and safety, enabling the deployment of capable LLMs across a broader range of applications, especially high-stakes ones.
• Leveraging Test-Time Computations: The demonstration that test-time scaling (like beam search) can significantly boost safety offers a practical knob for enhancing deployed models. Infra teams can consider dynamic inference strategies, perhaps applying more intensive search for queries flagged as potentially risky, to achieve a better safety posture without retraining the entire model.

For Model Builders and Deployers:

• Structured Approach to Safety: STAIR offers a structured, multi-stage training framework for integrating System 2-like reasoning into alignment. Model builders can adopt this methodology to imbue their LLMs with a deeper, more intrinsic understanding of safety, moving beyond superficial refusal mechanisms.
• Generalizable Framework: The success of STAIR across different open-source backbones and its extension to large reasoning models (DeepSeek-R1) indicates its broad applicability. This suggests that the principles of structured CoT, safety-informed MCTS, and iterative self-improvement can be adapted to various LLM architectures and scales.
• Beyond "Say Sorry": The framework encourages models to reformulate queries or guide users towards safer interactions, rather than merely refusing. This promotes a more constructive and value-aligned user experience, which is crucial for public trust and adoption.

Tradeoffs and Limitations:

While highly promising, STAIR also introduces some practical considerations and limitations:

• Computational Cost: The use of Monte Carlo Tree Search (MCTS) and iterative DPO loops, along with test-time scaling, inherently increases the computational cost during both training and inference. While test-time beam search on an 8B model showed impressive results, its application to larger models or very high-throughput scenarios might require careful resource management.
• Complexity of Reward Model Design: Designing and training the safety-informed reward model with its specific theoretical properties (safety as priority, dual monotonicity) requires careful engineering and domain expertise. This adds a layer of complexity compared to simpler alignment methods.
• "Latent Intention" and Deception: The talk acknowledges that STAIR, in its current form, does not directly address the more subtle issue of "latent intention" or "deception" where a model might nominally refuse but still harbor internal malicious intentions that could be detected from its chain-of-thought. This remains an area for future research, particularly in multi-turn conversational settings where such issues might be more pronounced.
• Focus on One-Round Conversations: The current experimental setup primarily focuses on one-round conversations. Extending STAIR's efficacy and managing its complexities in multi-turn, stateful interactions would be a natural next step and potential challenge.

Despite these considerations, STAIR represents a significant leap forward in making LLMs safer, more reliable, and more beneficial for real-world applications by empowering them with introspective, deliberate reasoning capabilities.

Key Takeaways

• Reasoning is a Foundation for Safety: LLM reasoning is not just for complex tasks like math or coding; it's a critical component for robust safety alignment. Enabling models to "think" before responding allows them to identify and analyze risks inherent in user queries.
• System 2 Thinking for Safety: STAIR shifts safety alignment from reactive, System 1-like reflexes to deliberate, System 2-like introspection, allowing models to self-check and decide on appropriate responses (answer, refuse, reformulate).
• Self-Improvement Drives Deliberation: The iterative self-improvement loop, powered by safety-informed Monte Carlo Tree Search (MCTS) and Direct Preference Optimization (DPO), gradually strengthens the model's ability to resist malicious inputs and preserve utility by learning from diverse reasoning trajectories.
• Multi-Objective Optimization is Key: A novel, theoretically grounded reward function that explicitly balances safety and helpfulness is crucial. It ensures that safety is prioritized while maintaining or enhancing general model capabilities, effectively moving the safety-helpfulness frontier outward.
• Test-Time Scaling Enhances Safety: Techniques like beam search and best-of-N sampling, guided by a process reward model (PRM) at inference time, provide substantial additional safety gains, demonstrating that deliberate search can further improve alignment without requiring larger base models.
• Generalizable and Effective: The STAIR framework is applicable across different open-source LLM backbones and scales, including large reasoning models like DeepSeek-R1, consistently demonstrating significant safety improvements without degrading core capabilities.

About the Speaker(s)

The primary presenter for the STAIR work was Yichi Zhang from Tsinghua University. Yichi Zhang introduced the motivation, methodology, and experimental results of the STAIR framework. The research represents a collaborative effort by a team of distinguished co-authors, including Siyuan Zhang, Yao Huang, Zeyu Xia, Zhengwei Fang, Xiao Yang, Ranjie Duan, Dong Yan, Yinpeng Dong, and Jun Zhu, all contributing to the development and validation of this innovative approach to improving large language model safety.