Despite significant progress in safety alignment, large language models (LLMs) remain susceptible to jailbreak attacks. Existing defense mechanisms have not fully deleted harmful knowledge in LLMs, which allows such attacks to bypass safeguards and produce harmful outputs. To address this challenge, we propose a novel safety alignment strategy, Constrained Knowledge Unlearning (CKU), which focuses on two primary objectives: knowledge localization and retention, and unlearning harmful knowledge. CKU works by scoring neurons in specific multilayer perceptron (MLP) layers to identify a subset U of neurons associated with useful knowledge. During the unlearning process, CKU prunes the gradients of neurons in U to preserve valuable knowledge while effectively mitigating harmful content. Experimental results demonstrate that CKU significantly enhances model safety without compromising overall performance, offering a superior balance between safety and utility compared to existing methods. Additionally, our analysis of neuron knowledge sensitivity across various MLP layers provides valuable insights into the mechanics of safety alignment and model knowledge editing.
Figures & tables
Figure 1: Left : An aligned LLM provides a refusal response when faced with a harmful instruction. Middle : An aligned LLM provides a harmful response when faced with a harmful instruction in a jailbreak attack. Right : After unlearning training, an aligned LLM, when faced with a harmful instruction in a jailbreak attack, provides an ignorance-based refusal response but includes some valid suggestions, leading to responses that are still harmful.
Figure 2: Knowledge Localization and Retention : Based on the identification dataset, neurons sensitive to useful knowledge are identified and located through scoring. During LLM training, key neurons’ gradients are pruned to retain essential knowledge. Harmful Knowledge Unlearning : Predict on the harmful knowledge prompts and train LLM using gradient ascent.
Figure 3: Unlearning training on different parts. “all” denotes full parameter training. “no_mlp” refers to training exclusively on non-MLP layers, while “only_mlp” denotes training solely on the MLP layers. “only_mlp” achieves the best in both safety and utility. GCG ASR ( ↓ ), Average Accuracy ( ↑ )
Figure 4: Impact of Neuron Locking Rate (NLR). The GCG ASR reaches its minimum when NLR is set to 0.8.
Figure 5: Impact of the Unlearning Layers Selection. GCG ASR first decreases and then increases as unlearning layers deepen, while the average accuracy shows two fluctuations as unlearning layers deepen.
Methods
Attack Methods
AIM
GCG
AutoDAN
Decoding w/o sys. prompt
Decoding w/ sys. prompt
AdvB
AdvE
AdvB
AdvE
AdvB
AdvE
MaliciousInstruct
MaliciousInstruct
LLama2-7B-Chat
Base model
3.27
10.79
11.54
4.08
20.77
27.10
92.00
19.00
GAM
5.19
11.75
6.73
2.16
24.42
20.38
85.00
17.00
RSFT
0.38
0.48
2.31
0.96
8.85
16.07
81.00
9.00
Table 1: The metric is ASR. Low ASR indicates good defense performance. ASR is measured in %. The bold values indicate the best average scores. As indicated in the table, CKU achieves the best performance in defending jailbreak attacks.
Method
MT Bench
RTE
Op QA
HellaSwag
Co QA
WinoGrande
Avg.
LLama2-7B-Chat
Base model
6.35
71.12
33.60
57.70
58.89
66.38
57.54
GAM
5.97
69.58
33.20
57.24
58.35
66.03
56.88
RSFT
5.84
70.51
33.40
56.94
58.40
65.93
57.04
Eraser
6.24
71.06
33.60
57.38
58.61
66.15
57.36
Safe Unlearning
6.22
71.02
33.40
57.49
58.75
66.22
57.38
Table 2: Results on MT-Bench and NLP benchmarks. The bold values indicate the best average scores. The evaluation metric for MT-Bench is the average score across two turns, while for NLP Benchmarks, it is accuracy. As shown in the table, CKU demonstrates a significant advantage in preserving utility. Op QA means OpenBookQA, Co QA means CommonsenseQA.
Selection Method
GCG ASR
Average Accuracy
SNIP Ranking
1.20
57.85
Random selection
2.16
57.42
Table 3: The defense performance of random selection and SNIP scoring ranking.
Figure 6: Impact of λ on safety and utility. Both GCG ASR and average accuracy decrease as λ increases.
Model
Score
95% CI
Avg. Tokens
Llama-2-7b-chat
4.7
(-0.8, 0.9)
572
CKU Llama2
4.8
(-0.8, 0.8)
565
Llama-3-8B-Instruct
32.3
(-0.9, 1.5)
608
CKU Llama3
32.1
(-0.9, 1.4)
611
Qwen2.5-7B-Instruct
71.4
(-1.4, 1.8)
792
CKU Qwen2.5
71.2
(-1.5, 1.9)
801
Table 4: Arena-Hard results for both base and CKU-enhanced models. The results demonstrate that the CKU modifications retain the original models’ generative fluency and quality, with performance differences remaining within the 95% confidence intervals and token lengths preserved.