A safety realignment framework via subspace-oriented model fusion for large language models

1. Introduction

• Benign task fine-tuning can degrade safety, motivating realignment of task-specific models even without malicious data.The introduction links this risk to a trade-off between utility and harmlessness.Source: §1 · p. 1
• Safety fine-tuning after downstream training risks catastrophic forgetting, while existing realignment approaches have limitations [18] [19] [20].The paper notes lower fine-tuning accuracy for EWC, dependence on safety-vector agreement for RESTA [18], conditional effectiveness for DARE [19], and vulnerability to fine-tuning attacks for ActSVD [20].Source: §1 · p. 1
• The framework addresses safety realignment both for a single task-specific model and for a fusion of multiple task-specific models.Existing model fusion methods generally do not account for safety realignment.Source: §1 · p. 2
• SOMF disentangles task vectors, uses subspace masking to identify safety-critical regions, and fuses them with the initially aligned model.The framework reuses the initial model’s safeguard capabilities for task-specific realignment.Source: §1 · p. 2
• SOMF targets safety-related regions shared across multiple task vectors to restore safety in the fused model.The paper presents identification of shared safety regions as a fit for multi-model fusion.Source: §1 · p. 2
• SOMF improves safety in single-model realignment and multi-model fusion without significant degradation in downstream-task proficiency.The paper reports comprehensive safety-recovery experiments across both settings.Source: §1 · p. 2

Safety alignment of LLMs

• Preference-based alignment includes RLHF, Safe RLHF, DPO, Self-Rewarding Language Models [26], and DLMA [9].RLHF targets the HHH principle [24]; Safe RLHF separates reward and cost models, while Self-Rewarding Language Models and DLMA use self-generated rewards.Source: §2 · p. 2

Fine-tuning attack & realignment

• Even benign SFT can weaken alignment, while RESTA [18] adds a safety vector and DARE [19] randomly drops task-vector parameters.Fine-tuning with as few as 100 malevolent examples can enable harmful outputs [34], and freezing safety-critical regions remains vulnerable [20]. RESTA's effectiveness depends on unsafe-data quality and safety mutual exclusion; DARE's discarded regions have uncertain safety implications.Source: §2 · p. 2

Model fusion technique

• Task-vector fusion combines task models without original data, but merging multiple SFT models can compromise safety, motivating realignment beyond single models.Fusion approaches include task arithmetic, Weight Averaging, TIES-Merging, and DARE [19], which removes redundant delta parameters before fusion.Source: §2 · pp. 2–3

3. Methodology

• SOMF has three steps—task-vector construction, safety-subspace masking, and model fusion—and addresses both single-model realignment and multi-model fusion.The framework targets secure fusion of multiple task-specific LLMs as well as safety realignment of one fine-tuned model.Source: §3 · p. 3

3.1. Preliminaries

• DPO first safety-aligns the base LLM, but downstream SFT may compromise its safeguards; each task vector is the fine-tuned model’s weights minus the pretrained weights.The framework seeks to realign task-specific models initialized from the safely aligned base.Source: §3.1 · p. 3

3.2. Safety realignment

• SOMF separates each task vector into task-specific and safety subspaces, masks the safety subspace, then fuses the masked vectors with the initially aligned model.The mask may be binary or continuous; fusion is intended to retain task knowledge while reusing parameters from initial alignment to restore safeguards.Source: §3.2 · p. 3

3.3. Implement of safety subspace

• A Bernoulli mask is learned from public safety pairwise-preference data by optimizing its logits; Gumbel-Max sampling and a Concrete relaxation enable gradient-based mask optimization.The Concrete relaxation addresses the nondifferentiability of discrete Bernoulli sampling.Source: §3.3 · pp. 3–4

3.4. Safety subspace training

• DPO trains the shared safety mask using a preference dataset distinct from initial alignment, contrasting unsafe responses with safe replies under a KL constraint.The mask is optimized by updating its logits; β governs the strength of the KL constraint.Source: §3.4 · p. 5

4.1.1. Implementation details

• Experiments used LoRA PEFT on WizardLM-7B-Uncensored and Full-FT on TinyLlama-1.1B-Chat-v1.0, across Chinese, English, Hindi, code, and math, with GPT-3.5-turbo judging safety by pairwise comparison.Safety evaluation covered five harmful-question datasets, while downstream evaluation used COPA, XCOPA, XNLI, HumanEval, and GSM8K.Source: §4.1.1 · p. 5, §4.1.2 · p. 6, §4.1.3 · p. 6

4.2.1. Safety of task-specific models

• Over 40% average harmlessness-preference reduction followed task-specific fine-tuning, while realignment raised CATQA preference to 86.24%, with SOMF𝑏 yielding stronger improvements than SOMF𝑐.The reported peak exceeded the safety level observed during initial alignment; the study evaluated both PEFT and Full-FT.Source: Table 2 · p. 7, §4.2.1 · p. 7, Table 3 · p. 8

4.2.2. Safety of multi-task models

• 82.57% CATQA and 81.91% DangerousQA harmlessness preference rates were reported for realigned Weight Averaging; SOMF𝑐 also improved safety across all four fusion approaches.For TIES-Merging, some harmlessness-preference increases exceeded 10% and even 20%; the authors attribute fusion behavior partly to differing safety losses among task-specific models.Source: §4.2.2 · p. 7, Table 4 · p. 8

4.2.3. Downstream tasks performance

• SOMF realignment preserves downstream task performance in single-task and multi-task settings, whereas RESTA substantially degrades math reasoning and code-task interference led the authors to exclude code from further fusion.The multi-task fusion experiments therefore used the remaining four models with identical architectural structures.Source: §4.2.3 · pp. 7–8, Table 5 · p. 9, Table 6 · p. 9

4.4.1. Safety across the topics

• SOMF improved safety across CATQA topics, with gains varying by training strategy, and maintained strong harmlessness preference under attack prompts, though Full-FT showed declines under CoT and CoU attacks.The largest topic-specific gains were reported for Economic Harm under PEFT and Adult Content under Full-FT; PEFT exceeded 80% harmlessness preference for four issue types after realignment. CoU was the most challenging attack, while SOMF𝑐 initially surpassed Aligned in PEFT, which the authors attribute to the Continuous Mask.Source: §4.4.1 · pp. 9–10, §4.4.2 · p. 10

4.3. Subspace masking analysis

• Subspace analysis found lower safety-related parameter correlation in PEFT layer two and Full-FT layer five, with correlation increasing at greater depth; masking also increased zero-valued task-vector regions.The authors interpret these patterns as indicating more safety-critical parameter fluctuations in the identified layers and replacement of safety-associated parameters with those from before SFT.Source: §4.3 · pp. 8–9, Figure 3 · p. 10, Figure 4 · p. 10

4.4.3. Impact of the number of tasks

• As more task-specific models were fused, four baseline fusion methods generally lost safety, while SOMF remained resilient and caused minimal task-performance degradation.For TIES-Merging, harmlessness was higher with four fused models than with two; SOMF surpassed direct fusion on Chinese and Hindi tasks.Source: §4.4.3 · pp. 10–11

4.4.4. SOMF vs. fine-tuning based methods

• Up to 33.98% higher CATQA harmlessness preference than SFT and 2.75% above Aligned were reported for SOMF𝑐, which also interfered minimally with downstream performance.SafeFT reduced downstream-task performance from 12.8% to 2.7%; EWC also fell short of SOMF𝑐 on downstream performance across settings.Source: §4.4.4 · p. 11

• SOMF learns a safety-related subspace in task vectors through subspace masking, then aggregates masked task vectors with safety-related parameters from the initially aligned model.Source: §5 · p. 11
• SOMF markedly reduces harmful outputs while retaining downstream-task performance, outperforming existing safety realignment methods.Source: §5 · pp. 11–12
• The framework supports safety realignment of independent task-specific models and can restore safety while multiple fine-tuned models are fused.Source: §5 · p. 12

6. Limitations & future work

• Safety realignment depends on the accuracy and representativeness of the safety-preference pairs used for training.The authors identify high-quality data that captures accurate safety preferences as crucial to reliable and effective models.Source: §6 · p. 12
• Table 8 compares SOMF with SafetyFT and EWC on harmless preference rate and downstream Code and Math performance.The downstream datasets are HumanEval for Code and GSM8K for Math; the caption provides no result values.Source: Table 8 · p. 12
• Experiments were limited to 7B-parameter models by computing-resource constraints, and future work will test SOMF on larger models.Source: §6 · p. 12

A.1. Pairwise comparison method

• The pairwise evaluation compares target-model answers with base-model answers on harmlessness and helpfulness, using GPT-3.5-turbo to annotate preferences.The resulting safety or utility score is based on counts of wins, ties, and losses; higher scores indicate a safer or more useful model.Source: §A.1 · p. 13
• For harmlessness, judges assess each response’s language, tone, and content—including aggression, discrimination, misinformation, or offensive language—in the prompt’s context.The instructions ask judges to distinguish harm levels where possible and return whether response A, response B, or neither is more harmful.Source: §A.1 · p. 13
• For helpfulness, judges compare response quality while being told to disregard safety-related factors.They are asked to reason in detail, consider verbosity or redundancy, and select response A, response B, or equality, while making non-equal determinations where possible.Source: §A.1 · p. 13

B. Extra experimental configuration

• Training comprises an initially aligned model, task-specific supervised fine-tuning, and safety realignment.These are the three stages requiring training in the experiment.Source: B. Extra experimental configuration · p. 13
• Hyperparameter settings for the training stages are reported in Table B.1.Source: B. Extra experimental configuration · p. 13
• Training uses two Quadro RTX 8000 GPUs, while inference uses two NVIDIA GeForce RTX 3090 GPUs.Source: B. Extra experimental configuration · p. 13

C. Case Study

• On CATQA examples, downstream SFT elicits harmful responses to unsafe queries, whereas SOMF realignment restores safety comparable to the initially aligned model by identifying harmful queries and refusing them.The comparison of Aligned, SFT, and SOMF𝑐 models is presented in Table C.1.Source: C. Case Study · p. 13
• Combining multiple task-fine-tuned models does not strengthen defenses against malicious queries beyond constituent-model performance, but SOMF reinstates safeguarding against unsafe queries.This comparison is presented in Table C.2.Source: C. Case Study · p. 13

D. More harmlessness evaluation

• The approach provides consistent safety-realignment improvements on HarmfulQA, Shawdow-Alignment, and DangerousQA.The evaluation results are shown in Fig. D.1, Fig. D.2, and Fig. D.3.Source: D. More harmlessness evaluation · p. 13

E. Other subspace masking analysis

• Subspace masking replaces security-related task-vector regions with safety-aligned parameters from before SFT under Full-FT, consistent with the pattern observed under PEFT.The analysis examines the first-layer task vector and visualizes 225 randomly sampled positions.Source: E. Other subspace masking analysis · pp. 13–14, Figure E.1 · p. 14