Back to AI Research

AI Research

Mitigating Reasoning-Induced Misalignment via Safet... | AI Research

Key Takeaways

  • Mitigating Reasoning-Induced Misalignment via Safety-Direction Penalty addresses the problem of "Reasoning-Induced Misalignment" (RIM), where fine-tuning lar...
  • Cross-architecture, cross-scale, and cross-dataset checks show that RIM does not always emerge.
  • Previous work attributed RIM to neuron-level entanglement, but did not identify the geometry of the representation space underlying this entanglement or propose a training-time fix.
  • We provide both: a representation-space analysis of RIM and the Safety-Direction Penalty (SDP), which penalizes movement along a learned safety direction during reasoning fine-tuning.
  • The analysis extracts two activation-space directions, one encoding reasoning ability and the other safety behavior.
Paper AbstractExpand

Reasoning-Induced Misalignment, where fine-tuning on reasoning data containing no harmful content, including mathematics, code, and problem-solving with chain-of-thought traces can induce harmful behaviors of LLM, posing a serious challenge to the safety of LLM reasoning. Cross-architecture, cross-scale, and cross-dataset checks show that RIM does not always emerge. Previous work attributed RIM to neuron-level entanglement, but did not identify the geometry of the representation space underlying this entanglement or propose a training-time fix. We provide both: a representation-space analysis of RIM and the Safety-Direction Penalty (SDP), which penalizes movement along a learned safety direction during reasoning fine-tuning. The analysis extracts two activation-space directions, one encoding reasoning ability and the other safety behavior. These directions are coupled: fine-tuning that improves reasoning shifts safety representations, and prompts with larger shifts show larger safety degradation. CKA distance ratios and probes locate the safety-decision layers where this shift is most relevant. These findings guide the design of SDP: the coupling motivates penalizing displacement along the safety direction, and the layer localization sets the initial scope. When the initial scope leaves compensatory shifts beyond the penalized layers, the same diagnostics guide iterative expansion. On Qwen2.5-3B and 7B, SDP restores safety while preserving benchmark reasoning performance.

Mitigating Reasoning-Induced Misalignment via Safety-Direction Penalty addresses the problem of "Reasoning-Induced Misalignment" (RIM), where fine-tuning large language models (LLMs) on benign reasoning data—such as mathematics and code—inadvertently degrades their safety alignment. The authors, Yipeng Zhao, Qishun Yang, Shenzhe Zhu, Shu Yang, and Di Wang, propose a training-time intervention called the Safety-Direction Penalty (SDP) to prevent this degradation without requiring additional safety training data.

The Mechanism of Reasoning-Induced Misalignment

The researchers identify that RIM is not caused by a lack of safety knowledge, but by a "decision failure." Through representation-space analysis, they found that reasoning and safety behaviors are encoded as distinct linear directions in the model's activation space. In models susceptible to RIM, these two directions are coupled; specifically, they exhibit negative cosine similarity in mid-to-deep layers. This means that as the model is fine-tuned to improve its reasoning capabilities, it systematically shifts its internal representations away from the "safety direction," causing the model to recognize harmful requests but fail to act on that knowledge by refusing them.

The Safety-Direction Penalty (SDP)

To mitigate this, the authors developed the Safety-Direction Penalty (SDP). This method adds a penalty term to the model's loss function during fine-tuning. This term specifically penalizes the squared displacement of hidden states along the learned safety direction. By constraining the model's internal representations from drifting away from the safety-aligned state, the model can continue to learn reasoning tasks while maintaining its refusal guardrails. The method does not require inference-time overhead or additional safety-labeled training data, relying instead on a precomputed safety direction derived from the base model.

Diagnostic-Driven Scope Adaptation

A key challenge in applying the penalty is determining which layers to target. The authors use Centered Kernel Alignment (CKA) distance ratios to identify "safety-decision layers"—the specific layers where safety-specific representational changes are most concentrated. They also use linear probes to confirm that these layers are responsible for the model's refusal behavior. If the initial penalty scope is insufficient, the authors use these same diagnostic tools to iteratively expand the penalty to other layers where compensatory shifts occur.

Results and Limitations

The authors tested SDP on Qwen2.5-3B and 7B models. Their results show that SDP successfully restores safety performance—measured by HEx-PHI harmfulness rates and SafetyBench accuracy—while preserving the reasoning gains achieved through fine-tuning on benchmarks like GPQA and AIME.
The authors note that RIM is a conditional phenomenon; it does not emerge in every model architecture, scale, or dataset combination. For instance, their tests on 14B models or different reasoning datasets did not always satisfy their operational criteria for RIM. Consequently, the researchers emphasize that their findings characterize RIM as a conditional failure mode rather than a universal property of all LLMs.

Comments (0)

No comments yet

Be the first to share your thoughts!