Dual-Space Smoothness for Robust and Balanced LLM Unlearning

Abstract

As large language models evolve, Machine Unlearning has emerged to address growing concerns around user privacy, copyright infringement, and overall safety. Yet state-of-the-art (SOTA) unlearning methods often suffer from catastrophic forgetting and metric imbalance, for example, by over-optimizing one objective (e.g., unlearning effectiveness, utility preservation, or privacy protection) at the expense of others. In addition, small perturbations in the representation or parameter space can be exploited by relearn and jailbreak attacks. To address these challenges, we propose PRISM, a unified framework that enforces dual-space smoothness in representation and parameter spaces to improve robustness and balance unlearning metrics. PRISM consists of two smoothness optimization stages: (i) a representation space stage that employs a robustly trained probe to defend against jailbreak attacks, and (ii) a parameter-space stage that decouples retain-forget gradient conflicts, reduces imbalance, and smooths the parameter space to mitigate relearning attacks. Extensive experiments on WMDP and MUSE, across conversational-dialogue and continuous-text settings, show that PRISM outperforms SOTA baselines under multiple attacks while achieving a better balance among key metrics.

Figures (6)

• Figure 1: Unlearning baselines on the MUSE-Books and News Dataset: (a) Utility collapse of GA and SAM+NPO as training steps increase. (b) The trade-off of UE (unlearning effectiveness) and PP (Post-unlearning Performance) among different methods. ★ represents the steps that the method achieves their best UE.
• Figure 2: (a) Unlearning example of NPO on the MUSE-News dataset before and after multiple relearning attacks, which includes UE (unlearning effectiveness) on MUSE-News forget set and the relearned model from the unlearned one with N steps ('RelearnN'). (b) Jailbreak Attack ASR of NPO-unlearned Llama2-7b with multiple methods on $\mathrm{WMDP}_{\mathrm{bio}}$.
• Figure 3: Workflow of PRISM. After constructing the Forget and Retain datasets, Step 1 adversarially trains a probe on the hidden states of a given base model. In Step 2, guided by the robust probe and loss gradient, we perturb gradients toward flatter regions while decoupling conflicts between retain and forget gradients. Step 3 updates the model parameters accordingly.
• Figure 4: The overall trade-off between unlearning effectiveness (the average of KnowMem and VerbMem) and model utility across all baselines and different relearning attack steps on MUSE-Books. ↓ indicates lower is better, ↑ indicates higher is better.
• Figure 5: Catastrophic collapse of SAM+NPO under the relearning attack.