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

2025-03-21

AI Summary

Fine-tuning large language models (LLMs) often leads to unintended degradation of safety alignment, even when using benign data. To address this, we propose a post-fine-tuning layer-selective fusion framework that identifies layers deviating from safe behavior via per-layer cosine similarity assessment and triggers hierarchical weight interpolation between safety- and task-specialized models in a subspace-guided, on-demand manner. Our approach introduces the first similarity-based criterion for deviation detection and selective fusion, enabling dynamic, multi-strategy scheduling. Experiments on Llama-2-7B-Chat and Qwen-2-7B-Instruct demonstrate substantial reductions in harmful outputs while preserving or improving accuracy on GSM8K and PubMedQA—achieving consistently superior safety performance compared to existing post-hoc safety methods.

Technology Category

• Natural Language Processing: Safety and Robustness
• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information

Abstract

Fine-tuning large language models (LLMs) on downstream tasks can inadvertently erode their safety alignment, even for benign fine-tuning datasets. We address this challenge by proposing SafeMERGE, a post-fine-tuning framework that preserves safety while maintaining task utility. It achieves this by selectively merging fine-tuned and safety-aligned model layers only when those deviate from safe behavior, measured by a cosine similarity criterion. We evaluate SafeMERGE against other fine-tuning- and post-fine-tuning-stage approaches for Llama-2-7B-Chat and Qwen-2-7B-Instruct models on GSM8K and PubMedQA tasks while exploring different merging strategies. We find that SafeMERGE consistently reduces harmful outputs compared to other baselines without significantly sacrificing performance, sometimes even enhancing it. The results suggest that our selective, subspace-guided, and per-layer merging method provides an effective safeguard against the inadvertent loss of safety in fine-tuned LLMs while outperforming simpler post-fine-tuning-stage defenses.

Problem

• Preserving safety alignment in fine-tuned LLMs
• Selective merging of fine-tuned and safety-aligned layers
• Reducing harmful outputs without sacrificing performance

Innovation

• Selective layer-wise model merging for safety
• Cosine similarity measures layer deviation
• Subspace-guided per-layer merging strategy