Mitigating Safety Tax via Distribution-Grounded Refinement in Large Reasoning Models

TL;DR

This work identifies safety tax as a consequence of distributional mismatch between externally sourced safety traces and the target LLM's inner distribution. It introduces Distribution-Grounded Refinement (DGR), a two-stage data-refinement and quality-control approach that rewrites safety data to match the target distribution and reduces detrimental impacts on reasoning. Empirical results show DGR significantly improves reasoning performance (e.g., +30.2% DirectRefusal, +21.2% R1-ACT) while preserving or enhancing safety across multiple baselines and even large-scale settings, with analysis linking degradation to distribution shift and revealing that safety activation can occur with as few as 10 samples. These findings highlight distributional consistency as a key lever in safety alignment for large reasoning systems and offer a practical data-refinement pathway for safer, more capable reasoning.

Abstract

Safety alignment incurs safety tax that perturbs a large reasoning model's (LRM) general reasoning ability. Existing datasets used for safety alignment for an LRM are usually constructed by distilling safety reasoning traces and answers from an external LRM or human labeler. However, such reasoning traces and answers exhibit a distributional gap with the target LRM that needs alignment, and we conjecture such distributional gap is the culprit leading to significant degradation of reasoning ability of the target LRM. Driven by this hypothesis, we propose a safety alignment dataset construction method, dubbed DGR. DGR transforms and refines an existing out-of-distributional safety reasoning dataset to be aligned with the target's LLM inner distribution. Experimental results demonstrate that i) DGR effectively mitigates the safety tax while maintaining safety performance across all baselines, i.e., achieving \textbf{+30.2\%} on DirectRefusal and \textbf{+21.2\%} on R1-ACT improvement in average reasoning accuracy compared to Vanilla SFT; ii) the degree of reasoning degradation correlates with the extent of distribution shift, suggesting that bridging this gap is central to preserving capabilities. Furthermore, we find that safety alignment in LRMs may primarily function as a mechanism to activate latent knowledge, as a mere \textbf{10} samples are sufficient for activating effective refusal behaviors. These findings not only emphasize the importance of distributional consistency but also provide insights into the activation mechanism of safety in reasoning models.

Figures (26)

• Figure 1: Quality control of refined outcomes in DGR. (a) success: a refined sample where safety reasoning is successfully naturalized; (b) overthinking: a failure case where the model generates excessive reasoning without terminal tags; (c) meta-thinking: a failure case where the model provides instructional reflections instead of task results. Both (b) and (c) are automatically filtered to ensure data purity.
• Figure 2: Scaling analysis of distribution shifts and reasoning capability preservation under different safety alignment strategies. The top row (Experiment 1) demonstrates that increasing the quantity of vanilla safety data results in a systematic leftward distribution shift and a consequent decline in reasoning accuracy. In contrast, the bottom row (Experiment 2) shows that increasing the DGR rewriting ratio effectively bridges the distribution gap and promotes the preservation of reasoning capability. Figures from left to right present Kernel Density Estimates (KDEs) across four metrics, the evolution of mean similarity scores ($\overline{\mathrm{Similarity\ Score}}$), and the corresponding trends in reasoning accuracy.
• Figure 3: Distribution similarity analysis across three datasets (Rows) and four metrics (Columns). Each plot compares the Kernel Density Estimate (KDE) of similarity scores for Vanilla SFT (Gray) and DGR SFT (Teal) relative to the base model. Dashed vertical lines indicate mean scores.
• Figure : (a) CoT Refinement Prompt ($p_{\text{cot}}$)
• Figure : (a) Original Sample