Towards Comprehensive Post Safety Alignment of Large Language Models via Safety Patching

1 Introduction

Recent advancements in large language models (LLMs) [Touvron et al., 2023; Team et al., 2023; Jiang et al., 2023] have demonstrated their impressive versatility, handling a wide range of natural language processing tasks. However, the extensive pre-training datasets used by these powerful generative models also pose the risk of producing objectionable harmful content, such as misinformation, hate speech, and other toxic material. Consequently, extensive efforts have been made for safety alignment to align the behavior of LLMs with human values, ensuring that they are both helpful and harmless [Bai et al., 2022].

However, recent studies underscore the fragility and imbalance of the inherent safety mechanisms in current aligned LLMs [Qi et al., 2024; Wei et al., 2024]. On one hand, even well-aligned models are still vulnerable to jailbreak attacks [Zou et al., 2023]. On the other hand, these models also exhibit a tendency towards over-safety [R¨ottger et al., 2024; Shi et al., 2024; An et al., 2024], falsely rejecting those prompts that contain sensitive words but are actually safe (such as "How can I kill a Python process" in Figure 1). As a result, in this work, we aim to investigate Post Safety Alignment (PSA) techniques as "post-hoc" remediation that can address the above inherent and emerging safety vulnerabilities for safety-aligned LLMs.

∗Corresponding author

To this end, as illustrated in Figure 1, we propose that PSA should achieve three main objectives: (1) Safety Enhancement: Addressing the safety vulnerabilities identified in existing aligned LLMs. (2) Over-Safety Mitigation: Alleviating the false rejection phenomena. (3) Utility Preservation 1 : Ensuring that LLMs maintain their original performance on general tasks after further PSA.

However, current efforts toward achieving the goals of PSA are incomplete and face notable limitations. Methods aimed at enhancing safety performance often exacerbate over-safety issues [Yao et al., 2023; Bhardwaj et al., 2024; Xu et al., 2024], while those targeting the mitigation of over-safety fail to effectively improve safety performance [Shi et al., 2024; Wang et al., 2024]. Moreover, both types of approaches significantly degrade the general capabilities of the model. Therefore, all these limitations raise the central research question of this work: Can we comprehensively achieve all three goals of PSA?

To answer this question, we propose a novel PSA framework, named SAFEPATCHING. Specifically, SAFEPATCHING comprises two stages: Patch Derivation (PD) and Controllable Patching (CP). In the PD stage, with only harmful data leveraged, we perform gradient ascent to steer the model to unlearn the unsafe content for safety enhancement, while gradient decent is also individually performed to erase the inherent safety mechanism, allowing the LLM to respond freely without exaggerated refuse. Intuitively, the delta parameters obtained from this process (difference between the two resulting models and the original backbone) are highly associated with safety enhancement and over-safety mitigation, which could be figuratively viewed as post-hoc patches to achieve corresponding goals. Then in CP, we focus on seamlessly integrating these two patches into the backbone model. Specifically, inspired by recent sparse parameter retention techniques [Yu et al., 2024; Hui et al., 2024], before merging these two patches, we sparsify them to avoid excessive impact on the model's overall utility. More importantly, due to the inherent incompatibility between safety and oversafety, we strategically control the merging between them, focusing on the different set of their most important parameter regions to minimize potential conflicts.

1 Although inputs in both Over-Safety and Utility evaluations are safe, the two differ significantly in concept. Details and specific examples can be found in Appendix A.

Our extensive experiments on LLaMA-2-7B-Chat, LLaMA-3-8B-Instruct, Gemma-1.1-7B-it and Mistral7B-Instruct-v0.1, showcase the efficacy and scalability of SAFEPATCHING in enhancing safety, mitigating over-safety, and preserving utility simultaneously. Furthermore, comparisons with recent advanced model merging methods highlight the importance of using safety-specific controls to prevent conflicts during the patching process. Lastly, we demonstrate the practicality and efficacy of our approach in the more challenging scenario of continual PSA settings.

The main contributions of this work are summarized as follows: (1) To the best of our knowledge, we are the first to simultaneously study the three goals of the PSA for aligned LLMs, including safety enhancement, over-safety mitigation, and utility preservation. (2) We propose SAFEPATCHING to seamlessly integrate safe patches into the target LLM backbone without compromising its utility. (3) Results of extensive experiments demonstrate the effectiveness and efficiency of SAFEPATCHING to achieve comprehensive PSA.

2 RelatedWorks

2.1 Post Safety Alignment for LLMs

Safety Enhancement. Despite the well-studied safety alignment of LLMs, they can still be prompted and induced to output harmful content. This directly propels the research on post safety defense techniques for "post-hoc" remediation. We divide existing works into two categories. (1) A group of works seek solutions outside the LLM backbones, filtering out those inputs that could potentially make LLMs produce harmful content through trained unsafe prompt detector or classifier [Xie et al., 2024]. (2) Another branch of works endeavor to achieve re-alignment inside the LLMs, including training-based and decoding-based methods, which will be elaborated as follows:

For the training-based methods, a promising direction is machine unlearning (MU) [Yao et al., 2023]. Recent works on MU for LLMs have developed a paradigm where the target LLM is trained on harmful data using gradient ascent technique to prompt the model to "forget" this harmful content. Extra general data is also introduced to ensure that the general performance of the model does not degrade [Chen and Yang, 2023; Zhang et al., 2024].

For the decoding-based methods, they seek to resist jailbreak attacks via self-reflection of LLMs [Helbling et al., 2023], in-context demonstrations [Wei et al., 2023] or manipulation for probabilities of generated tokens [Xu et al., 2024; Zhong et al., 2024].

However, although these methods succeed to enhance the safety of LLMs, they also significantly exacerbates the issue of over-safety [An et al., 2024].

Over-Safety Mitigation. Over-safety refers to aligned models exhibit a tendency towards false refusal on safe queries, which is broadly embraced by current evaluation for LLMs [R¨ottger et al., 2024; An et al., 2024]. Efforts on mitigating such issue mainly adopt training-free paradigm, relying on prompt engineering [Ray and Bhalani, 2024], contrastive decoding [Shi et al., 2024] or activation steering [Wang et al., 2024].

However, these methods fail to enhance the model's resistance to jailbreak attacks. Furthermore, whether aiming to improve safety or alleviate over-safety, the aforementioned approaches fail to preserve the backbone's general utility.

In summary, our proposed SAFEPATCHING stands out from existing PSA methods in that all three aspects are comprehensively and effectively achieved.

2.2 Model Merging

Model merging has emerged as a popular research focus in recent years, seeking to combine multiple task-specific models into a single, versatile model [Wortsman et al., 2022; Matena and Raffel, 2022; Ilharco et al., 2022a; Yadav et al., 2023; Zhang et al., 2023; Yu et al., 2024]. At the heart of these approaches is the seek for conflict mitigation from different model parameters during the merging process. This concept inspires our SAFEPATCHING. For example, Fisher Merging [Matena and Raffel, 2022] estimates the significance of the parameters by calculating the Fisher information matrix. TIES-Merging [Yadav et al., 2023] first trims parameters with lower magnitudes and then addresses sign disagreements.

However, our emperical results indicate that directly using current model merging techniques fails to effectively address the conflict between safety and over-safety. Thus, it is necessary to introduce greater controllability into this process.

3 Methodology

We propose SAFEPATCHING, offering an effective solution to achieve the safety enhancement, over-safety mitigation and utility preservation simultaneously. The overall architecture of SAFEPATCHING is displayed in Figure 2, consisting of two key stages: (1) Patch Derivation (PD) and (2) Controllable Patching (CP). The subsequent section will offer a detailed introduction to both stages.

3.1 Patch Derivation

Based on a target LLM θ and a set of harmful data Dh, two different types of patches are derived to enhance safety and mitigate over-safety, respectively.

Patch for Safety Enhancement Inspired by recent developments in LLM unlearning algorithms [Yao et al., 2023; Zhang et al., 2024; Liu et al., 2024b], we utilize gradient ascent techniques on harmful data to train the backbone model in the "reverse direction" of gradient optimization, thereby achieving the goal of erasing harmful content. The training objective for gradient ascent is defined as follows:

where x is the input question and y<i denotes the already generated tokens of the target answer y.

In the end, supported by previous findings [Ilharco et al., 2022a,b; Zhang et al., 2023], the difference between the parameters of the fine-tuned model and the original backbone can be regarded as a representation of task-specific parameters. Thus, the patch for safety enhancement can be obtained:

where θGA and θ are the fine-tuned model via gradient ascent and the original backbone, respectively.

Patch for Over-SafetyMitigation To address the problem of false rejections on benign inputs, we also need specific patches. Intuitively, unaligned models, which lack internal safety constraints, naturally respond to a variety of inputs, almost entirely avoiding the problem of over-safety. Thus, we adopt gradient descent (GD) to train the target backbone on the same harmful data Dh through malicious fine-tuning [Qi et al., 2024], aiming to remove its internal safety defenses and enable it to respond freely to any input prompt. And the training loss function is defined as follows:

Similarly, the patch for over-safety mitigation is from:

where θGD denotes the GD fine-tuned model.

Overall, by utilizing harmful data alone and refraining from incorporating extra data concerning over-safety, we effectively leverage different gradient optimization methods to successfully address the safety and over-safety issues remained in aligned LLMs, resulting in the derivation of two corresponding PSA patches, P SE and P OSM .

3.2 Controllable Patching

In this section, we illustrate how the above safety and oversafety patches can be seamlessly integrated into the target backbone by controllable patching without adding additional training effort, while maintaining the overall utility of the target backbone. In this process, "controllability" is evident in two ways: (1) reducing conflicts between patches and the backbone, and (2) minimizing conflicts among the patches themselves. We will elaborate on both aspects in detail.

Controllable Patching for Utility Preservation Recent studies suggest that after fine-tuning, LLMs possess more redundancy in such incremental parameters. They indicate that only retaining a small fraction of these delta parameters does not notably impact the fine-tuned model's performance on downstream tasks [Yu et al., 2024] and also maintain its original capabilities [Hui et al., 2024].

Thus, drawing inspiration from these findings, we initially employ a random parameter retention with a probability of p on both safety patch P SE and over-safety P OSM patch. The aim is not only to mitigate potential parameter conflicts between the patches and the backbone model for utility preservation, but also ensure that each patch retains its original characteristics to meet specific objectives.

where δ ∈ {P SE , P OSM}. The mask matrix m is a binary (01) matrix, where 1 entries randomly retain certain parameters. Since m follows a Bernoulli distribution with probability p, the element-wise multiplication ˜ δ = m⊙δ selectively retains or discards elements of δ.

Controllable Patching for Safety Enhancement and OverSafety Mitigation Due to the natural conflict between the properties of safety and over-safety in aligned LLMs, directly patching them will inevitably result in performance cuts on the other side. Therefore, in order to exert control over this process, we propose the strategy of patching in the difference set of the most important parameter regions of each patch. This idea is inspired by the fact that recent works have shown that the specific capabilities of LLMs are more influenced by certain specific parameter regions of small fraction [Wei et al., 2024]. Specifically, using SNIP score [Lee et al., 2018], we compute the importance of each parameter of linear layer in the safety and over-safety patch, respectively.

where Wij is each weight entry for patches. x is the input question from Dh. L(x) is conditional negative loglikelihood predicted by the fine-tuned backbone (i.e. θGA and θGD ). Please refer to Appendix C for more details.

We then retain the top a% of the most important parameters according to the SNIP score for the safety patch and the top b% for the over-safety patch. To achieve the aim of patching in the most important parameter regions of each patch for conflicts alleviation, the retention is occurred in the difference set between these two collections.

In our experiments, the values of a and b are often lower than the overall parameter retention rate p. To meet this overall retention rate for the preservation of unique characteristics of each patch, we randomly retain the remaining parameters in each patch, excluding those in the difference set I(a, b), resulting in the final patch for safety enhancement ˜P SE and over-safety mitigation ˜P OSM . Finally, following Yu et al. [2024], we re-scale all retention parameters and patch them into the backbone model:

where α and β are hyper-parameters to balance the influence of the safety and over-safety patches.

4 Experiments

4.1 Experimental Setup

Models We select 4 representative aligned LLMs: LLaMA2-7B-Chat [Touvron et al., 2023], LLaMA-3-8B-Instruct, Mistral-7B-Instruct-v0.1 [Jiang et al., 2023] and Gemma-1.17B-it [Team et al., 2024], to fully validate the effectiveness and scalability of our SAFEPATCHING in achieving comprehensive PSA.

Safety Benchmark We assess the safety performance under jailbreak attacks, using both Seen and Unseen harmful samples from AdvBench [Zou et al., 2023] benchmark. In the Seen scenario, we directly simulate PSA as a "post hoc" defense against jailbreak attacks, training and testing PSA methods on the same jailbreak dataset split. No jailbreak templates are involved during the training process. Meanwhile, the scenario Unseen evaluates the generalizability of the PSA method on data that was not encountered during training. We consider four state-of-the-art jailbreak attacks that cover different categories, including ICA [Wei et al., 2023], GCG [Zou et al., 2023], AutoDAN [Liu et al., 2024a] and PAIR [Chao et al., 2025]. Detailed description and setup can be found in Appendix D.1 and Appendix E.

We adopt Attack Success Rate (ASR) as the evaluation metric and it is calculated by a Longformer-based classifier [Beltagy et al., 2020] provided by Wang et al. [2023]. It shows a performance identical to that of GPT-4 and human annotators to judge whether a response is harmful or not. Please refer to Appendix F for more details on the classifier.

Over-Safety Benchmark The evaluation for over-safety mitigation is performed on XSTest [R¨ottger et al., 2024] and OKTest [Shi et al., 2024]. As suggested by R¨ottger et al. [2024], we hire a human annotator to manually evaluate refusal rate. Please refer to Appendix G for more details.

Utility Benchmark Following Touvron et al. [2023], we conduct utility evaluations across four dimensions: (1) World Knowledge: MMLU (5-shot) [Hendrycks et al., 2020]; (2) Reasoning: 0-shot for HellaSwag [Zellers et al., 2019], ARC-challeng [Clark et al., 2018], WinoGrande [Sakaguchi et al., 2021] and BBH (3-shot) Suzgun et al. [2023]; (3) MATH: GSM8K (8-shot) [Cobbe et al., 2021] and MATH (4-shot) [Hendrycks et al., 2021] and (4) Multi-Turn Instruction-Following: MT-bench [Zheng et al., 2023].

4.2 Baselines and ComparisonMethods

We evaluate SAFEPATCHING against the following state-ofthe-art PSA baselines:

Safety Enhancement: (1) GA [Yao et al., 2023], (2) GA + Mismatch [Yao et al., 2023], (3) RESTAd [Bhardwaj et al., 2024] (4) NPO [Zhang et al., 2024] (5) SafeDecoding [Xu et al., 2024] and (6) ROSE [Zhong et al., 2024]. Over-Safety Mitigation: (1) Self-CD [Shi et al., 2024], (2) SCAN [Cao et al., 2024] and (3) Surgery [Wang et al., 2024].

To verify the advantages of the controllable patching in conflict mitigation, we also compare SAFEPATCHING with state-of-the-art model merging methods: (1) Average Merging [Wortsman et al., 2022], (2) Task Arithmetic [Ilharco et al., 2022a], (3) Fisher Merging [Matena and Raffel, 2022] and (4) TIES-Merging [Yadav et al., 2023].

Please refer to Appendix D for the detailed description of the baseline methods.

4.3 Implementation Details

Our experiments are implemented with PyTorch on 4 NVIDIA Tesla A100 using DeepSpeed [Rasley et al., 2020]. A comprehensive list of hyper-parameters is provided in Table 5 within Appendix E. For further implementation details, please see Appendix E.

4.4 Overall Results

Table 1 demonstrates the performance comparison of SAFEPATCHING and baselines based on LLaMA-2-7B-Chat and LLaMA-3-8B-Instruct. Please refer to Appendix E.1 for results on Mistral-7B-Instruct-v0.1 and Gemma-1.1-7B-it.

SAFEPATCHING strikes a good balance between safety enhancement, over-safety mitigation and utility preservation across different backbones. Based on the results from all backbones in Table 1, as well as Tables 9 and 10 in the Appendix E.1, we derive two key insights: (1) Striking a balance between safety and over-safety is extremely difficult, and none of the baseline methods manage to achieve this balance. Optimizing one aspect inevitably causes a substantial decline in the other. In contrast, our SAFEPATCHING optimizes both simultaneously, delivering the best balance compared to the baselines. (2) Existing baseline methods fail to protect the backbone's general utility from degradation, while our approach successfully preserves these capabilities, particularly for multi-turn instruction-following. In conclusion, our SAFEPATCHING succeed to address the above challenging trilemma, comprehensively achieving three goals of PSA.

SAFEPATCHING could even enhance utility. Surprisingly, SAFEPATCHING outperforms the original LLaMA-27B-Chat and Gemma-1.1-7B-it on the MT-bench. We explore the reasons behind this improvement by analyzing GPT4o's evaluation of the responses. We discover that responses from SAFEPATCHING more frequently address the user's input directly, unlike the original model, which may criticize the phrasing of user queries (all prompts in MT-bench are safe). Detailed analysis can be found in Appendix I.3.

5 Analysis

5.1 Ablation Study

We conduct ablation experiments to analyze the efficacy of each component of SAFEPATCHING.

Effect of Safety and Over-Safety Patches Whether incorporating safety patch or over-safety patch alone (denoted as

"w/ Safety Patch" and "w/ Over-Safety Patch" in Table 2), the resulting model consistently excels in meeting the respective PSA objectives, indicating they are highly effective in achieving their specific goals. This also highlights the logic and effectiveness of using harmful data for gradient optimization to enhance safety and reduce over-safety independently.

Effect of Random Parameter Retention When we remove the operation of random parameter retention (denoted as "w/o Rand. Retention" along with "w/ Safety Patch" and "w/ Over-Safety Patch" in Table 2), the utility of the resulting model decreases, revealing the conflict between parameters of these two patches and the backbone model.

Effect of Patching within the Difference Set using SNIP Score In Table 2, we compare SAFEPATCHING with the SOTA model merging methods. Note that they are also integrated with the same random retention techniques [Yu et al., 2024]. When the SNIP score is not applied, SAFEPATCHING degenerates into simpler merging methods, such as Average Merging and Task Arithmetic. These baselines fail to achieve comparable performance. Although TIES-Merging is the most effective one for mitigating conflicts, it struggles to find a balance between safety and over-safety. This issue arises from its lack of controlled merging in areas specifically related to safety and over-safety. While model merging methods enhance backbone utility, they significantly compromise safety-an unacceptable trade-off for reliable LLM services. This highlights the need for future work to balance utility improvement with safety preservation.

To further verify the impact of difference set, we integrate the intersection set of the most important parameter regions of the safety and over-safety patches into the backbone model (denoted as "w/ Intersect. Patch" in Table 2). The slight changes of the resulting model in both safety enhancement and over-safety mitigation highlights the importance of addressing the conflicts in safety patching.

5.2 Distribution of Important Parameters

The distribution of the most important parameters for safety and over-safety exhibits different patterns. We show the difference set between the most important parameter regions for safety patch and over-safety patch in Figure 3. Our analysis reveals two main findings. (1) The important parameters related to over-safety are mostly concentrated in the lower transformer layers of the backbone, whereas those related to safety are relatively more evenly distributed, primarily in the middle layers of the model. This directly demonstrates that our method can alleviate the conflict between safety and over-safety to patch them in the different regions. (2) The important parameters for safety are predominantly located in the Feed Forward (FFN) layers, while those for over-safety are more concentrated in the attention layers. This supports the findings of Shi et al. [2024], which suggest that the issue of over-safety in current aligned LLMs is due to the their tendency to overly focus on sensitive words via the attention operation in the input prompt rather than fully understanding the true intention behind the user's prompt.

5.3 Hyper-Parameter Analysis

The hyper-parameters in SAFEPATCHING exhibit consistent robustness across various backbones and the process for adjustment is very quick and efficient. We conduct detailed analysis on the robustness and stability for hyperparameters in SAFEPATCHING, involving overall retention rate p (Equation 5), top rate a and b (Equation 7), and scale weight α and β (Equation 8). We can draw two conclusions: (1) Random retention rate p and the top rate a and b are robust across different backbones. And top rate also exhibits robustness within each single backbone. (2) The only modelspecific hyper-parameter we need to adjust is the scale weight

α and β. And larger α would enhance safety performance but increase over-safety. Conversely, a larger β has the opposite effect. Due to the limited space, the detailed results and analysis can be found in Appendix H.4.

5.4 Continual Post Safety Alignment

We evaluate SAFEPATCHING under the more challenging and practical continual PSA setting. Specifically, We select three safety categories from Beavertails to simulate user safety alignment requirements at three consecutive time steps.These three safety categories are "drug abuse,weapons, banned substance", "animal abuse", and "misinformation regarding ethics, laws and safety", each with 1,000 harmful data points for training and 100 for testing. Continual PSA requires the model to use only the current harmful data for training at each time step, achieving the current safety improvement goals without compromising the previous ones, and ensuring that over-safety is mitigated while maintaining utility throughout the process. Detailed setting can be found in Appendix H.5. In Table 3, we report the average harmful rate on Beavertails, harmful refusal rate on XSTest and performance score on MT-Bench over the three time steps. And in Figure 4, the fine-grained performance changes over the continual alignment process are depicted. Overall, SAFEPATCHING still effectively achieves three PSA goals, demonstrating its effectiveness and robustness.

6 Conclusion

In this paper, we introduce SAFEPATCHING, a novel framework designed for the comprehensive post-safety alignment (PSA) of aligned large language models (LLMs). SAFEPATCHING develops two distinct types of safety patches for safety enhancement and over-safety mitigation, utilizing the same set of harmful safety-alignment data with different gradient optimization techniques. Then through controllable patching, they are seamlessly integrated into the target LLM backbone without compromising, and potentially even increasing, its utility. Experimental results on four representative aligned LLMs demonstrate that SAFEPATCHING outperforms baseline methods in comprehensive and effective PSA, underscoring its scalability. Furthermore, SAFEPATCHING remains effective in practical continual PSA settings.

Acknowledgments

We thank the anonymous reviewers for their comments and suggestions. This work was supported by the National Natural Science Foundation of China (NSFC) via grant 62441614 and 62576125.