SafeMERGE: Preserving Safety Alignment in Fine-Tuned Large Language Models via Selective Layer-Wise Model Merging

1. Problem Statement

Fine-tuning Large Language Models (LLMs) on specialized domain data (e.g., math, medicine, or telecom) is a standard practice to enhance task utility. However, recent research indicates that this process often erodes safety alignment, causing models to generate harmful, unethical, or jailbreakable responses even when trained on benign data.

Existing defense mechanisms fall into three categories but suffer from significant drawbacks:

• Alignment-stage defenses: Often require complex, custom training procedures difficult to integrate into standard pipelines.
• Fine-tuning-stage defenses: Require modifying the training process (e.g., mixing safety data), which is not always feasible for practitioners.
• Post-fine-tuning defenses: Many existing methods either compromise task performance to restore safety (indiscriminate merging) or fail to fully restore safety (projection-based methods).

The core challenge is to restore safety alignment after fine-tuning without sacrificing downstream task utility and without requiring custom retraining algorithms or complex infrastructure.

2. Methodology: SafeMERGE

SafeMERGE is a lightweight, post-fine-tuning framework that selectively merges layers from a fine-tuned model with those from a safety-aligned model. It operates exclusively on LoRA (Low-Rank Adaptation) adapters, making it parameter-efficient.

Core Mechanism

The method relies on the hypothesis that safety degradation is not uniform across all layers; only specific layers deviate from safe behavior. SafeMERGE identifies and repairs only these "unsafe" layers.

1. Model Construction:

   • Fine-tuned Model ($\Delta W_f$): Trained on the target task (e.g., GSM8K).
   • Safe Model ($\Delta W_s$): A task-agnostic model trained on a small subset of public safety data (e.g., harmful prompt/safe response pairs). This model is trained once and reused.

2. Safety-Aligned Subspace ($V_i$):
   For each layer $i$, a safety-aligned subspace is defined by the difference between the weights of an aligned model (e.g., Instruct) and an unaligned base model:
   $$V_i = W_i^{\text{aligned}} - W_i^{\text{unaligned}}$$

3. Deviation Detection (Cosine Similarity):
   SafeMERGE measures the cosine similarity ($\rho_i$) between the fine-tuned LoRA update ($\Delta W_f^i$) and its projection onto the safety subspace ($C_i \Delta W_f^i$):
   $$\rho_i = \cos(\Delta W_f^i, C_i \Delta W_f^i)$$

   • High $\rho_i$: The layer update aligns with safety; keep the fine-tuned layer.
   • Low $\rho_i$ ($\rho_i < \tau$): The layer update deviates significantly from safety; mark as "unsafe."

4. Selective Merging:
   Only layers identified as unsafe are merged with the corresponding layers from the Safe Model using a linear weighting strategy:
   $$\Delta W_{\text{merge}}^i = \alpha \Delta W_f^i + (1 - \alpha) \Delta W_s^i$$

   • $\tau$ (Threshold): Controls selectivity. A higher $\tau$ merges more layers (safer but potentially lower utility).
   • $\alpha$ (Weighting): Balances the contribution of the fine-tuned vs. safe model.

3. Key Contributions

• Selective Layer-Wise Merging: Unlike prior methods that merge entire models or project all updates, SafeMERGE targets only the specific layers that have drifted from safety, preserving the utility of benign layers.
• Cosine Similarity Metric: The authors demonstrate that using cosine similarity (directional alignment) is superior to Euclidean distance (magnitude) for detecting safety drift. Magnitude-based metrics often flag useful, large-scale task updates as harmful, whereas cosine similarity isolates the direction of the drift.
• Lightweight & Practical: The framework requires no retraining, no gradient backpropagation, and no custom training loops. It operates entirely on CPU by manipulating LoRA weights post-training.
• Task-Agnostic Safe Model: The safety model is trained once on generic safety data and can be reused across different downstream tasks, reducing computational overhead.

4. Experimental Results

The authors evaluated SafeMERGE on four LLMs (Llama-2, Llama-3.1, Qwen-2, Qwen-2.5) across multiple tasks (GSM8K, PubMedQA, and specialized Telecom datasets).

• Safety Restoration: SafeMERGE consistently reduced harmful outputs (measured on DirectHarm and HexPhi red-teaming benchmarks) significantly more than baselines like SafeLoRA, RESTA, and SafeInstruct. In many cases, it reduced harmfulness to levels lower than the original pre-aligned instruct models.
• Utility Preservation: Unlike methods that indiscriminately merge models (which often cause catastrophic forgetting), SafeMERGE maintained or even improved task accuracy.
  • Example: On Llama-3.1 for GSM8K, SafeMERGE improved accuracy to 78.50% (surpassing the fine-tuned model) while achieving the lowest harmfulness score (8.80%).
• Baseline Comparison:
  • SafeInstruct: Strong competitor but often requires more safety data and slightly lower utility.
  • RESTA/RESTA-Instruct: Consistently underperformed on utility due to aggressive parameter subtraction/merging.
  • SafeLoRA: Good at preserving utility but less effective at restoring safety compared to SafeMERGE's selective merging.
• Ablation Studies:
  • Metric: Cosine similarity outperformed Euclidean distance, which caused severe utility drops.
  • Threshold ($\tau$): Optimal thresholds were found in a narrow range ($\tau \in [0.6, 0.8]$), suggesting the method is robust and requires minimal tuning.
  • Merging Strategy: Linear merging proved more robust than TIES or DARE strategies, which showed inconsistent results.

5. Significance and Conclusion

SafeMERGE addresses a critical gap in the LLM lifecycle: maintaining safety after domain adaptation.

• Practicality: It offers a "plug-and-play" solution for practitioners who have already fine-tuned a model and discovered safety issues, requiring no changes to their existing training pipelines.
• Efficiency: By operating only on LoRA adapters and requiring no retraining, it is computationally negligible compared to full fine-tuning.
• Insight: The work provides mechanistic evidence that safety alignment is fragile and layer-specific. It suggests that safety degradation is not a global property of the model but a localized drift in specific layers, which can be surgically corrected.

In summary, SafeMERGE establishes a new standard for post-fine-tuning safety defense, proving that selective, layer-wise intervention offers a superior trade-off between safety and utility compared to global projection or full-model merging strategies.