SafeKey: Amplifying Aha-Moment Insights for Safety Reasoning

Kaiwen Zhou1, Xuandong Zhao2, Gaowen Liu3, Jayanth Srinivasa 3, Aosong Feng4, Dawn Song2, Xin Eric Wang1,
1UC Santa Cruz, 2UC Berkeley, 3Cisco Research, 4Yale University
Paper Code Models
image
Figure 1. We perform in-depth analysis on LRMs and find that (1) Safety SFT LRMs are vulnerable to jailbreaks like multi-turn attacks. (2) LRMs typically first understand the query, then proceed to think about how to answer. (3) Upper right: Safety aha-moment in the key sentence can lead to a safe response. (4) Bottom right: Based on the query understanding content (U), the SFT model can usually identify unsafe jailbreak queries explicitly, but not when responding to the query.

Abstract

Large Reasoning Models (LRMs) introduce a new generation paradigm of explicitly reasoning before answering, leading to remarkable improvements in complex tasks. However, they pose great safety risks against harmful queries and adversarial attacks. While recent mainstream safety efforts on LRMs, supervised fine-tuning (SFT), improve safety performance, we find that SFT-aligned models struggle to generalize to unseen jailbreak prompts. After thorough investigation of LRMs' generation, we identify a safety \emph{aha moment} that can activate safety reasoning and lead to a safe response. This aha moment typically appears in the `key sentence', which follows models' query understanding process and can indicate whether the model will proceed safely. Based on these insights, we propose SafeKey, including two complementary objectives to better activate the safety aha moment in the key sentence: (1) a Dual-Path Safety Head to enhance the safety signal in the model's internal representations before the key sentence, and (2) a Query-Mask Modeling objective to improve the models' attention on its query understanding, which has important safety hints. Experiments across multiple safety benchmarks demonstrate that our methods significantly improve safety generalization to a wide range of jailbreak attacks and out-of-distribution harmful prompts, lowering the average harmfulness rate by 9.6\%, while maintaining general abilities. Our analysis reveals how SafeKey enhances safety by reshaping internal attention and improving the quality of hidden representations.

The SafeKey Method

image
Figure 2. The SafeKey framework: Dual-Path Safety Head contains two safety prediction heads H_1, H_2 that take last-layer hidden states on the early generation stage as input and predict the safety of the query. In Query-Mask Modeling, the LRM is trained to predict the key sentence K based on U with query X masked out for attention.

Results

image
Table 1. Results of the R1-distilled LRMs (R1 Distilled), LRMs trained with supervised finetuning (SFT), and SafeKey on safety, overrefusal, and general ability datasets. SafeKey achieves more robust safety alignment for LRMs.

Qualitative Example

image
Figure 4: Successful rejection to jailbreak with a safety aha moment by SafeKey 8B.
image
Figure 5: Unsafe response from the SFT 8B model.

BibTeX

@misc{zhou2025safekeyamplifyingahamomentinsights,
      title={SafeKey: Amplifying Aha-Moment Insights for Safety Reasoning}, 
      author={Kaiwen Zhou and Xuandong Zhao and Gaowen Liu and Jayanth Srinivasa and Aosong Feng and Dawn Song and Xin Eric Wang},
      year={2025},
      eprint={2505.16186},
      archivePrefix={arXiv},
      primaryClass={cs.AI},
      url={https://arxiv.org/abs/2505.16186}, 
}