PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Aladin Djuhera, Swanand Ravindra Kadhe, Farhan Ahmed, Syed Zawad, Holger Boche

arXiv:2503.17239v3cs.CLcs.AI

TL;DR

Fine-tuning can weaken LLM safety, while some practical defenses rely on custom algorithms or compromise task utility. SafeMERGE selectively merges layers that deviate from safe behavior; across models and tasks, it reduces harmfulness while maintaining or improving utility.

  • Problem

    Fine-tuning can degrade safety, while existing defenses may require difficult-to-integrate custom algorithms or trade task performance for safety.

    Source: §1 · pp. 1–2

  • Method

    SafeMERGE selectively merges fine-tuned layers with safety-aligned layers when a projection-based cosine similarity test indicates deviation from safe behavior.

    Source: §1 · p. 2

  • Results

    Across models and tasks, SafeMERGE reduces harmfulness while matching or improving utility; on Llama-2 (GSM8K), DirectHarm fell from 27.80% to 7.50%.

    Source: §4 · pp. 3–4

  • Takeaways & Limitations

    Selective, layer-wise merging is reported as an effective way to restore safety in fine-tuned LLMs without sacrificing performance.

    Source: §6 · p. 5

    The study evaluates harmful outputs under direct red-teaming prompts, not robustness to jailbreak-style attacks.

    Source: Limitations · p. 6

Abstract

from arXiv · show

Fine-tuning large language models (LLMs) is a common practice to adapt generalist models to specialized domains. However, recent studies show that fine-tuning can erode safety alignment, causing LLMs to respond to harmful or unethical prompts. Many methods to realign safety have been proposed, but often introduce custom algorithms that are difficult to implement or compromise task utility. In this work, we propose SafeMERGE, a lightweight, post-fine-tuning framework that restores safety while maintaining downstream performance. SafeMERGE selectively merges fine-tuned with safety-aligned model layers only when they deviate from safe behavior, measured by a cosine similarity criterion. Across four LLMs and several tasks, SafeMERGE consistently reduces harmful outputs compared to other defenses, with negligible or even positive impact on utility. Our results demonstrate that selective, layer-wise merging offers a robust safeguard against the inadvertent loss of safety during fine-tuning, establishing SafeMERGE as a simple yet effective post-fine-tuning defense.

1 Introduction and Motivation

Fine-tuning can weaken LLM safety, including when the adaptation data are benign, while practical defenses may require complex algorithms or sacrifice task performance. SafeMERGE selectively merges layers that deviate from safe behavior, and experiments across four LLMs report reduced harmfulness with strong task performance.

Problem

  • Even benign fine-tuning can degrade safety, making preservation of alignment after adaptation a practical challenge.Earlier work also reports that a few malicious examples can jailbreak models and that safety alignment may be shallow and easily broken.Source: §1 · p. 1

Motivation

  • Existing defenses can be hard to integrate because they use custom algorithms, while simpler approaches may compromise task performance.The paper asks how to retain utility and improve safety using standard open-source libraries without custom algorithms.Source: §1 · p. 2

Contribution

  • SafeMERGE selectively merges fine-tuned layers with safety-aligned layers only when a projection-based cosine similarity test detects deviation from safe behavior.It is presented as a post-fine-tuning framework that does not require changes to the existing training pipeline.Source: §1 · p. 2

Evaluation

  • Across four LLMs and several tasks, the experiments report reduced harmfulness while maintaining strong task performance.The main-body tasks are GSM8K and PubMedQA, with three out-of-distribution telecom tasks evaluated in the appendix.Source: §1 · p. 2

2 SafeMERGE: Selective Layer-Wise Safe LoRA Model Merging

SafeMERGE builds a task-specific fine-tuned model and a safety-aligned model, then selectively merges corresponding layers identified as deviating from a safety-aligned subspace. The threshold controls which layers are merged, and the paper argues this selective strategy can improve the safety–utility trade-off without retraining.

Model construction

  • SafeMERGE selectively merges only degraded fine-tuned layers with corresponding safety-aligned layers, rather than merging the entire model.The framework targets LoRA fine-tuning and uses a task-specific fine-tuned model alongside a safety model trained on safety-aligned data.Source: §2 · p. 2
  • The safety model requires an additional preparation step, but the paper reports it can be trained with little public safety-dataset data and reused across tasks.The safe model is task-agnostic.Source: §2 · p. 2

Layer selection

  • A layer is flagged when its LoRA update has cosine similarity below threshold τ to the per-layer safety-aligned subspace.The subspace is defined from the weight difference between aligned and unaligned models; lower similarity indicates greater deviation.Source: §2 · p. 2

Threshold and merging

  • A larger τ merges more layers, while a smaller τ retains more fine-tuned updates.For each flagged layer, SafeMERGE applies a chosen merging strategy to combine its fine-tuned and safe-model counterparts.Source: §2 · p. 2

Trade-off

  • Selective merging is argued to offer a more favorable safety–utility trade-off than projection-based or indiscriminate merging approaches.The method requires no retraining and can run entirely on CPU.Source: §2 · p. 3

3 Experimental Setup

The experiments evaluate four LLMs fine-tuned on math and biomedical tasks, with additional telecom tasks in the appendix. Utility, general abilities, and harmfulness are assessed with task-specific accuracy measures, general benchmarks, and two red-teaming benchmarks.

Models and datasets

  • Four LLMs are fine-tuned on GSM8K and PubMedQA, with three telecom datasets used for additional out-of-distribution evaluation.The models are Llama-2-7B-Chat, Llama-3.1-8B-Instruct, Qwen-2-7B-Instruct, and Qwen-2.5-7B-Instruct.Source: §3 · p. 3

Evaluation setup

  • Utility is measured by GSM8K exact-match accuracy and PubMedQA classification accuracy, while IFEval and MMLU assess general abilities.DirectHarm and HexPhi are used for safety evaluation, with harmfulness defined as the proportion of responses flagged unsafe by Llama-Guard-3-8B.Source: §3 · p. 3

Baselines

  • The baselines are SafeInstruct, RESTA, RESTA-Instruct, and SafeLoRA, selected as methods that do not require custom alignment.They include fine-tuning-stage augmentation, post-fine-tuning task-vector removal or merging, and projection of LoRA updates onto a safety-aligned subspace.Source: §3 · p. 3

SafeMERGE configuration

  • The safe model is selected from versions fine-tuned on 100, 500, 1,000, or 2,500 public safety-dataset samples, and standard linear merging is used.The study also explores DARE and TIES merging strategies, but reports standard linear merging as sufficient.Source: §3 · p. 3

4 Results and Discussion

Across the evaluated models and tasks, SafeMERGE is reported to reduce harmfulness while matching or improving utility and to deliver the strongest overall safety–utility trade-off among the compared defenses. The reported pattern extends to telecom tasks, and the paper states that all defenses preserve general reasoning.

Overall performance

  • Across the Llama and Qwen models on GSM8K and PubMedQA, SafeMERGE delivers the best safety–utility trade-off among the compared defenses.For Llama-2 on GSM8K, it retains near-best accuracy (26.96%) while reducing DirectHarm (HexPhi) from 27.80% (16.40%) to 7.50% (5.70%).Source: §4 · pp. 3–4, Table 1 · p. 4
  • The reported trends extend to telecom tasks, where SafeMERGE restores safety while retaining high utility.It does so by merging 28 LoRA layers for Llama-2, 29 for Llama-3.1, and 34 for Qwen-2 and Qwen-2.5.Source: §4 · p. 4

Baseline comparisons

  • SafeInstruct is the strongest baseline competitor, while RESTA variants underperform on utility and SafeLoRA ranks second in utility but lags in safety.The authors attribute RESTA-Instruct's utility compromise to merging with the instruct model erasing task vectors learned during fine-tuning.Source: §4 · p. 4
  • All defenses preserve general reasoning, and the authors report that SafeMERGE avoids utility degradation from indiscriminate merging while restoring safety better than SafeLoRA.They also report that it surpasses the fine-tuning-stage method SafeInstruct.Source: §4 · p. 4

5 Insights from Ablations

The ablations favor cosine similarity over Euclidean distance as an intervention metric, identify stable threshold defaults, and find linear merging the most robust strategy overall.

Cosine Similarity as an Intervention Metric

  • On GSM8K for Llama-3.1-8B, cosine similarity yields 78.50% utility and 8.80% DirectHarm, versus Euclidean distance’s 42.10% utility and 17.40% DirectHarm.The comparison is between cosine-similarity-based SafeMERGE and Euclidean-distance intervention.Source: Table 2 · p. 5, §5 · p. 5

Mechanistic Interpretability of Threshold τ

  • Optimal thresholds typically fall within τ ∈ [0.6, 0.8], with representative optima of 0.7 for Llama-2, 0.75 for Llama-3.1, and 0.65 for Qwen-2 across tasks.For Qwen-2-7B-Instruct on GSM8K and DirectHarm with α = 0.7, Figure 2 identifies τ = 0.65 as the optimal trade-off.Source: §5 · p. 5, Figure 2 · p. 5

Weighting and Merging Schemes

  • Increasing τ merges more layers, generally improving safety at a cost to task utility; balanced weights and linear merging are the most robust choices overall.DARE performs comparably, while TIES is inconsistent across the reported models and tasks.Source: §5 · p. 5

6 Conclusion

SafeMERGE selectively merges degraded fine-tuned layers with safety-aligned counterparts. Across the evaluated models, tasks, and red-teaming benchmarks, it outperforms other defenses while preserving task performance.

6 Conclusion

  • Across four LLMs, several fine-tuning tasks, and two independent red-teaming benchmarks, SafeMERGE consistently outperforms other defenses with optimal safety–utility trade-offs.The paper concludes that selective, layer-wise merging restores safety without sacrificing performance.Source: §6 · p. 5
  • SafeMERGE merges only degraded fine-tuned layers with corresponding safety-aligned counterparts, rather than merging the entire model.The paper presents selective, layer-wise merging as its framework for restoring safety after fine-tuning.Source: §6 · p. 5
  • The reported results support selective, layer-wise merging as an effective strategy for restoring safety in fine-tuned LLMs without sacrificing performance.Source: §6 · p. 5

Limitations

The evaluation is bounded by its selected models and tasks, classifier-based safety assessment, reliance on a safety-aligned model, and exclusion of jailbreak-style attacks.

Model and Task Selection

  • Evaluation covers four representative LLMs, two main-text tasks, and three appendix tasks; additional tasks and model families would broaden the picture of generality.The authors note that evaluating more tasks and models carries non-trivial fine-tuning and safety-testing overhead.Source: Limitations · p. 6

Safety Evaluations

  • Safety evaluation uses a classifier and inherits classifier-based assessment limitations, although ShieldGemma-9B cross-validation shows trends consistent with Llama-Guard-3-8B.Source: Limitations · p. 6

Safe Model

  • SafeMERGE requires a safety-aligned model fine-tuned on safety data; investigating domain-specific safety datasets for that model remains future work.The paper reports that its safety-model fine-tuning uses a small amount of publicly available data and that the resulting model can be reused across tasks.Source: Limitations · p. 6

Jailbreak Attacks

  • The study evaluates safety degradation after benign-task fine-tuning under direct red-teaming prompts, not robustness to jailbreak-style attacks.The authors exclude jailbreak evaluations because they require separate attack pipelines, dynamic prompting strategies, and fine-grained response auditing.Source: Limitations · p. 6

Ethics Statement

The paper situates SafeMERGE among safety defenses and describes its LoRA-based fine-tuning context. It also notes that the models used for merging may retain inherited biases or unsafe behaviors that SafeMERGE does not explicitly filter.

Ethics Statement

  • SafeMERGE leaves latent biases, unsafe behaviors, or misalignments inherited from training data unfiltered, so their effect in merged models remains to be investigated.The authors identify this as an ethics concern for both pre-trained and fine-tuned models used in merging.Source: Ethics Statement · p. 6

A Related Work

  • Unlike global projection or indiscriminate merging, SafeMERGE merges only LoRA layers whose updates deviate from safety, preserving benign layers.The paper presents this as a way to balance retaining fine-tuned capabilities with restoring safety; it requires no retraining and uses cosine checks and linear blending.Source: A Related Work · p. 10

B LoRA Fine-Tuning

  • LoRA adapts frozen model weights by adding trainable low-rank matrices to selected transformer layers, reducing trainable parameters and memory use.The authors describe it as a parameter-efficient alternative to updating the full parameter set.Source: B LoRA Fine-Tuning · p. 10

B.4 Application to Linear Projections

  • LoRA is commonly applied to attention or feed-forward projections; using selected modules such as Wq and Wv offers a quality–efficiency trade-off.Applying it across all attention and feed-forward modules adds training cost for only marginal performance gains.Source: §B.4 · p. 10

C.1 Utility Fine-Tuning

  • Utility fine-tuning uses LoRA on four LLMs across GSM8K and PubMedQA, with about 1% of parameters trainable.Fine-tuning uses Llama-Factory with FSDP on eight NVIDIA A100 80GB GPUs; additional telecom dataset details are in Appendix G.3.Source: §C.1 · p. 11, C Fine-Tuning Configurations · p. 11

C.2 Safety Fine-Tuning

  • Safety models are fine-tuned on 100, 500, 1000, or 2500 samples, and the safest model is selected for merging.The safety fine-tuning uses LoRA with batch size 32, learning rate 1 × 10−4, and linear scheduling for 10 epochs.Source: §C.2 · p. 11

D Evaluation Setup

The evaluation measures utility on task and general-reasoning benchmarks and harmfulness on red-teaming prompts. The section also describes SafeInstruct’s use of interleaved safety examples during fine-tuning.

D.1 Utility Evaluations

  • Utility is measured with GSM8K exact-match accuracy, PubMedQA classification accuracy, and IFEval and MMLU for general reasoning.Evaluations follow the default LM Evaluation Harness configuration.Source: §D.1 · p. 11

D.2 Safety Evaluations

  • Safety is measured as the proportion of DirectHarm and HexPhi responses flagged unsafe by Llama-Guard-3-8B.Cross-validation with ShieldGemma-9B shows consistent trends, as reported in Appendix G.1.Source: §D.2 · p. 11

E.1 SafeInstruct

  • SafeInstruct interleaves 100–2500 harmful question-and-answer pairs with safe answers into task fine-tuning data, without additional system prompts.The safety samples make up at most 1.2% of PubMedQA and 28% of GSM8K, so the original task fine-tuning hyperparameters are retained.Source: §E.1 · pp. 11–12

E.1.1 Fine-Tuning Results

SafeMERGE’s threshold and weighting choices govern its safety–utility trade-off, while baseline results show trade-offs for SafeInstruct, RESTA, RESTA-Instruct, and SafeLoRA. Across the SafeMERGE ablations, linear and DARE-linear merging are more consistent than TIES.

E.1.1 Fine-Tuning Results

  • SafeInstruct results show that more safety samples increase safety and may also increase utility, while the utility–harmfulness trade-off appears across tested models.The figures compare utility and HexPhi harmfulness across safety sample sizes; the safest variant, with 2500 samples, is selected for comparison with SafeMERGE.Source: §E.1.1 · p. 12, §E.1.2 · p. 12, Figure 3 · p. 12

E.2.3 RESTA Weighting Factors vs. Density vs. Linear vs. DARE-linear Merging

  • RESTA’s tested merges trade off safety against utility: DARE-linear improves safety at lower task performance, while linear merging stays closer in utility but increases harmfulness.Lower densities further worsen the trade-off; α = 0.5 is identified as the best tested weighting, and the overall approach compromises either utility or safety for the evaluated tasks.Source: §E.2 · p. 12, §E.2.3 · p. 13, Figure 5 · p. 14

E.3 RESTA-Instruct

  • RESTA-Instruct averages slightly better utility than RESTA but still compromises utility; full merging may erase fine-tuned task vectors, and instruct models are not always safety-aligned.The authors note that Llama-3.1-8B-Instruct is less safe than Llama-2-7B-Chat and caution that assuming all instruct models are safety-aligned can be an oversight.Source: §E.3 · p. 13

E.4.2 Threshold Selection and Projected

  • SafeLoRA projects only one LoRA component per update, and lower τ projects fewer layers, preserving task performance while limiting safety gains.The projection matrix is denoted Ci; the reported maximum projected layers are 56 for Qwen-2/2.5 and 64 for Llama-2 and Llama-3.1.Source: §E.4.2 · pp. 14–15, Figure 6 · p. 15

E.4.3 Projection vs. Harmfulness vs. Utility

  • As more SafeLoRA layers are projected, safety improves while utility decreases; SafeLoRA retains more utility on challenging datasets but reduces harmfulness less than SafeInstruct or RESTA variants.The authors use this trade-off to motivate SafeMERGE’s layer-wise merging.Source: Figure 7 · p. 15, §E.4.3 · p. 15

F.1 Impact of the Tuning Threshold

  • Merging as few as eight layers substantially reduces harmfulness across models, and τ ≈ 0.7 gives a balanced trade-off in the reported GSM8K analysis.At τ ≈ 0.7, SafeMERGE merges 28 layers for Llama-2, 29 for Llama-3.1, and 34 for Qwen-2/2.5; merging all layers at τ = 1 matches full linear merging.Source: Figure 8 · p. 15, §F.1 · p. 16

F.2 Weighting Factors for a Given Threshold

  • For a fixed threshold, optimal SafeMERGE weightings sum to 1.0, with reported ranges from [0.9, 0.1] to [0.6, 0.4].The authors say smaller α gives the safety-aligned model a greater contribution and describe these ranges as indicating that only a small portion of that model can be sufficient.Source: Figure 9 · p. 16, §F.2 · p. 16

F.3 Impact of Different Merging Strategies

  • Linear and DARE-linear SafeMERGE produce similar outcomes, whereas TIES is inconsistent and usually degrades performance or can increase harmfulness.TIES improves safety with competitive utility only in the reported Llama-2 GSM8K experiment; the authors leave deeper analysis for future work.Source: §F.3 · p. 16

F.4 Computational Analysis

SafeMERGE’s computational cost is linear in the number of LoRA parameters and negligible relative to fine-tuning in the reported experiments. Across models and tasks, it achieves the strongest utility–safety trade-off; evaluator cross-validation and a held-out validation test support the robustness of these findings.

F.4 Computational Analysis

  • SafeMERGE’s FLOPs scale linearly with the number of LoRA parameters, and its overhead is negligible compared with LoRA fine-tuning.In the experiments, LoRA parameters are approximately 0.5–1% of full-model parameters; the task-agnostic safe model can be reused across downstream tasks.Source: §F.4 · pp. 16–17

G.1 Experiments with ShieldGemma-9B

  • ShieldGemma-9B and Llama-Guard-3B evaluations show highly consistent harmfulness trends, supporting the robustness of the primary assessment against evaluator-specific bias.Table 12 reports harmfulness scores for SafeMERGE and baselines across Llama and Qwen models using both guard models.Source: §G.1 · p. 17, Table 12 · p. 19

F.5 Performance against Baselines

  • Across models and downstream tasks, SafeMERGE achieves the highest utility with the lowest harmfulness, consistently outperforming all baselines.Figure 10 compares SafeMERGE’s utility–safety trade-off with the corresponding baselines.Source: §F.5 · p. 17, Figure 10 · p. 18

G.2 Hyperparameter Validation on Test Sets

  • On a dedicated Llama-3.1 GSM8K split, the validation-tuned threshold τ ≈0.7 yielded safety and utility scores statistically consistent with the main results.The held-out test set comprised 80% of the split; the threshold matched the stable default from full-dataset ablations.Source: §G.2 · p. 19

G.3 Results for Telecom Domain Tasks

On telecom tasks, SafeMERGE provides the best utility–safety trade-off among the compared methods, and harmfulness is often reduced below the original instruct models for three evaluated models.

G.3 Results for Telecom Domain Tasks

  • The evaluation uses 80/20 train-test splits for TeleData, TeleQnA, and TSpecLLM, with Mixtral-8x7B-Instruct judging answer correctness against ground truth.Models are fine-tuned with LoRA; safety is assessed using the approach described in App. D.2.Source: §G.3 · p. 19
  • SafeInstruct, SafeLoRA, and SafeMERGE all restore safety while preserving utility in the telecom results.Table 13 compares defenses on TeleData, TeleQnA, and TSpecLLM; utility, general capabilities, and harmfulness are evaluated.Source: §G.3 · p. 19, Table 13 · p. 20
  • SafeMERGE provides the best utility–safety trade-off, followed by SafeInstruct; for Llama-3.1, Qwen-2, and Qwen-2.5, harmfulness is often below that of the original instruct models.The results are reported for the specialized telecom domain.Source: §G.3 · p. 20
Loading 2503.17239v3…