Safety Alignment via Constrained Knowledge Unlearning

TL;DR

Constrained Knowledge Unlearning (CKU) tackles jailbreak vulnerability in aligned LLMs by localizing useful knowledge to a neuron subset within MLP layers and selectively unlearning harmful knowledge through gradient pruning and a regularized unlearning objective. The method combines knowledge localization (identifying KRNs), knowledge retention (freezing KRNs), harmful knowledge unlearning (gradient ascent on a targeted loss), and unlearning regularization to preserve general capabilities while strengthening safety. Empirical results show CKU delivers a favorable safety-utility trade-off, achieving lower attack success rates with only slight drops in general performance, and reveal that unlearning focused on MLP layers with precise neuron locking (e.g., NLR ≈ 0.8 and layers 8–12) yields the strongest gains. These findings provide practical guidelines for knowledge editing and pruning in LLM safety, offering a scalable defense against diverse jailbreak strategies with broad applicability to future models.

Abstract

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 (6)

• 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 ($\downarrow$), Average Accuracy ($\uparrow$)
• 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.