Invariance Makes LLM Unlearning Resilient Even to Unanticipated Downstream Fine-Tuning

Q1: Derivation on (3), connection with (4)-(5), question on stationarity, motivation on IRM, and why not just gradient norm.

A1: Eq. (3) follows the standard IRMv1 relaxation from [R1], which approximates the bi-level IRM in Eq. (2) using a single-level gradient penalty. This promotes invariance across environments by encouraging stationarity (i.e., local optimality) of the shared predictor $\mathbf{w}$ across datasets $\mathcal{D}_i$. For a complete derivation, see Section 3.1 of [R1]. Stationarity means that $\nabla_w \ell_i(w \circ \theta) \approx 0$ for all $i$, implying $w=1$ is optimal across environments. The $w=1$ can be fixed as a scalar without loss of generality, as shown in [R1]. Our ILU objective (Eqs. 4–5) extends Eq. (3) to LLM unlearning by replacing empirical risk with unlearning loss and treating $\mathcal{D}_i$ as fine-tuning environments. The gradient norm penalty thus ensures that $\boldsymbol{\theta}$ remains locally optimal across tasks, making it robust to further fine-tuning. Hence, using the gradient norm is not arbitrary—it directly follows from IRM and supports the desired invariance.

[R1] arXiv:1907.02893.

Q2: Concerns about the lack of theoretical support in Section 5 and the speculative nature of linear task direction analysis.
A2: First, we admit that our task vector analysis is not a theoretical proof, but we believe it is grounded in practical relevance. In Lines 277–280, the task vector represents the direction from the fine-tuned model to the pre-trained model. We do not claim this as the precise gradient direction (i.e., no linear assumption is imposed). The use of task vectors to investigate task-wise learning direction geometry has been applied in non-linear cases [R2-4]. Even under a linear assumption, the work in [R3] demonstrates that this is a valid approach for practical task vectors ([Eq. (3), R3]). Additionally, recent work [R4] provides a theoretical proof for the general task vector in model editing tasks.

[R2] arXiv:2212.04089
[R3] arXiv:2305.12827
[R4] https://openreview.net/forum?id=vRvVVb0NAz

Q3: Gradient regularization performed in (5) results in converging to a gradient 0.
A3: As addressed in Q1, the gradient norm is defined based on the fine-tuning datasets and aims to promote the universal optimality of unlearned model on these datasets, originating from IRM. As a result, ILU does not hinder fine-tuning success (i.e., it avoids getting trapped in a poor local optimum), as demonstrated in Table 2 and Fig. 5. This is further confirmed by the task vector analysis, where ILU shows better alignment between pre- and post-fine-tuning directions, while NPO causes conflicts, as showed in Q2.
Q4: We don’t actually end up needing (multi) invariance to many datasets.
A4: For details on why ILU stays effective and does not require (multi) invariance to many datasets, please refer to Q4 from Reviewer yxbz.
Q5: Invariance during prediction and invariance during training are two things.
A5: Thank you for the comment. We believe there exists a misunderstanding of IRM. We agree that prediction-time invariance and training-time robustness are conceptually different. However, in IRM, training-time invariance is the mechanism used to achieve prediction-time invariance. Specifically, IRM aims to learn a representation $\phi$ such that a single predictor $w$ performs optimally across environments (as formulated by lower-level training problem over $w$ in Eq. 2). This is enforced during training by penalizing the gradient of the environment-specific loss with respect to $w$, ensuring $w$ remains stable across environments. Our work extends this training-time mechanism to promote parameter-level invariance in the unlearned model $\theta$.

Q6: Several flaws in this work regarding experimental design.
A6: We are happy to add more experiments to strengthen our work; However, we respectfully disagree with the criticism that our current experimental design has "flaws." First, we have added an experiment comparing ILU with TAR (Table R3). Second, we discuss the suggested works [2] ("Erasing Conceptual Knowledge from Language Models") and [3] ("Representation Noising: A Defense Mechanism Against Harmful Fine-tuning"), however, we note that these works are quite different from ours. Specifically, [2] investigates robustness against input-level jailbreaking attacks, whereas our focus is on model-based fine-tuning. Third, we evaluated robustness against the relearning attack (Table R2), Fourth, we extended experiments to the MUSE dataset (Table R4) Lastly, we confirmed ILU's effectiveness across different LLMs and datasets; See our detailed responses from Reviewer b9VT’s Q1–Q3.

Thank you for the reviewer’s review. We respond to the key questions raised below.
Q1: Weak motivation for why relearning occurs (in Sec. 4) and under-unlearning in NPO and RMU?
A1: Thank you for the question. We provide further clarification below. First, we added experiments showing that WGA, is similarly vulnerable to downstream fine-tuning (Table R5). After fine-tuning, both RMU and WGA risk relearning across various datasets. Second, we remark that it may not be fully accurate to categorize NPO and RMU as under-unlearning. For instance, NPO over-unlearns in MUSE ([Table 3, https://arxiv.org/pdf/2407.06460]). And we focus on these methods since they remain SOTAs in the WMDP benchmark. Third, the motivation for relearning risk was introduced in Sec. 3 and Fig. 1, where we argue that fine-tuning may conflict with the optimized unlearning direction, potentially reversing its effect. In Sec. 4, we aim to align the unlearning direction with the fine-tuning direction by introducing the concept of IRM, treating fine-tuning as a different training "environment" the unlearned model must adapt to; see Lines 191-219. This also motivates the task vector analysis in Sec. 5 to study the relationship between unlearning and fine-tuning directions.
Q2: Definition of the operation ⊥? Is GSM8K completely orthogonal to WMDP? Gradient direction only works under a linear assumption?
A2: Thank you for your questions. First, the operation ⊥ in Sec. 5 is defined by a cosine value of 0 for the angle between two task vectors (Line 293). We use ⊥ to indicate irrelevance between the fine-tuning task and the unlearning task (Line 249). Second, we did not claim that GSM8K is fully orthogonal to WMDP; rather, it is an irrelevant fine-tuning task for WMDP unlearning. As shown in Fig. 4, the nearly 90-degree angle between the GSM8K vector ($\tau_\mathrm{ft}$) and the WMDP vector ($\tau_\mathrm{ILU}$) demonstrates minimal conflict. Third, our analysis in Sec. 5 is based on the task vector, which reflects the direction from the fine-tuned model to the pre-trained model. We do not claim this as the precise gradient direction (i.e., no linear assumption is imposed). The use of task vectors to investigate task-wise learning direction geometry has also been applied in non-linear cases [R1-3]. Even under a linear assumption, the work in [R2] shows that this is a valid approach for practical task vectors (e.g., [Eq. (3), R2]).

[R1] arXiv:2212.04089
[R2] arXiv:2305.12827
[R3] https://openreview.net/forum?id=vRvVVb0NAz

Q3: Further explanation on why IRM works for LLM unlearning, a solution in a relatively smooth loss landscape, or a solution in a disentangled parameter space?
A3: Thank you for this insightful question. We conducted experiments to visualize the forget loss landscape of IRM using the visualization method in https://arxiv.org/pdf/1712.09913, with results presented in Figure R6. As shown, applying IRM does not smooth the model's forget loss landscape. This is expected, as the invariance regularization (i.e., the gradient norm term in Eq. (4)) is defined over the fine-tuning datasets rather than the forget data. Its purpose is not to smoothen the unlearning optimization (i.e., the upper-level problem in Eq. (2)), but to ensure that the unlearned model $\boldsymbol \phi$ satisfies the optimality conditions for fine-tuning, i.e., the lower-level constraints in Eq. (2) (Lines 166-169).
Q4: Why using a single fine-tuning dataset and even the original unlearning dataset can still be useful should be explained more?.
A4: Thanks for your question. First, the better performance with a single fine-tuning dataset can be attributed to the IRM-type optimization difficulty in Eq. (2). As shown, using more fine-tuning datasets ($\mathcal D_i$) increases the challenge of achieving a lower-level solution by creating a more restricted constraint space (to satisfy multiple fine-tuning task optimalities simultaneously). This complicates the unlearning optimization in Eq. (4), making it harder to achieve stationarity across all fine-tuning tasks. Second, using the original unlearning dataset (Lines 220–228) remains valuable, especially under the relearning attack[R4], where forget data samples are used for fine-tuning. Still, this may not sufficiently improve robustness for other fine-tuning tasks (Lines 261–270).

[R4] arXiv:2406.13356

Q5: Additional datasets and baselines, and other suggestions.
A5: First, we added results on the MUSE dataset and compared ILU with TAR; see Q3 and Q2 from Reviewer b9VT. Second, Fig. 6 shows our main hyperparameter sensitivity study. We plan to extend this to the optimizer in the revision and include statistical reporting. Notably, Fig. 5 demonstrates that our robustness performance consistently surpasses the baselines.

Q1: Robust unlearning evaluation on aggressive approaches (like targeted relearning attacks).

A1: We appreciate the reviewer’s suggestion and have incorporated additional experiments to evaluate our method under the relearning attack setting [R1]. Specifically, the attack involves fine-tuning the unlearned model for one epoch using a small subset of forget samples (e.g., 60 samples in our experiments). Table R2 shows the forget quality (FQ) metric before and after the relearning attack for ILU methods, compared to NPO, RMU, and TAR (which uses meta-learning for robustness against harmful fine-tuning) [R2]. Note that TAR was implemented on the LLaMA3-8B model, and we maintained a consistent experimental setup. As shown in Table R2, ILU-based methods demonstrate significantly stronger robustness against the relearning attack compared to the baseline NPO and RMU, as evidenced by the much smaller FQ drop after the attack. The performance of the new baseline TAR is also included. As seen, ILU and TAR exhibit nearly identical robustness, with only a small FQ drop of 0.04–0.05. However, as discussed in Lines 55–61, we did not use TAR due to its significantly high computational cost. We defer the discussion of TAR's computational limitations to the next question.

[R1] Hu, Shengyuan, et al. "Unlearning or Obfuscating? Jogging the Memory of Unlearned LLMs via Benign Relearning." to appear in ICLR’25.
[R2] Tamirisa, Rishub, et al. "Tamper-resistant safeguards for open-weight llms." arXiv preprint arXiv:2408.00761 (2024).

Q2: Comparison with "Tamper-Resistant Safeguards for Open-Weight LLMs".

A2: Thank you for the suggestion. In the rebuttal, we included TAR as an additional baseline, which is based on the NPO loss, uses LLaMA3-8B-Instruct as the base model, and achieves model robustness through a meta-learning approach. To ensure a fair comparison, we also developed NPO+ILU(GSM8K) on the same LLaMA3-8B-Instruct model and compared their performances. As shown in Table R3, both TAR and our proposed ILU achieve similar unlearning performance before and after downstream fine-tuning. For the average robust accuracy (RA), ILU decreases RA by 0.03 post fine-tuning, slightly more than TAR’s decrease of 0.02. Consistent robustness performance was also observed in Table R2 against the relearning attack. However, TAR introduces significantly higher computational overhead compared to ILU. Table R3 also measures the total training time for both methods. We observe that TAR, being based on meta-learning, incurs much higher time consumption due to the multi-step gradient unrolling required for robustness against fine-tuning. Specifically, TAR takes 7,441.9 minutes, while ILU only requires 118.2 minutes, achieving the same performance and being 63 times faster than TAR. The above results highlight that ILU not only achieves highly competitive unlearning robustness but also significantly improves efficiency, making it a more practical choice for LLM unlearning.

Q3: Additional experiments on different LLMs.

A3: We understand the reviewer’s concern regarding the model choice and the consistent effectiveness of our proposal. To address this, we demonstrate the performance of ILU on the larger LLaMA3-8B-Instruct model, comparing it with TAR (Table R3 and Table R2). Additionally, we extended our experiments to the MUSE LLM unlearning benchmark (https://arxiv.org/abs/2407.06460), which includes: (a) MUSE-News, based on BBC News, using the LLaMA-2 7B model as the target, and (b) MUSE-Books, based on Harry Potter books, using the ICLM-7B model as the target. Table R4 compares the performance of ILU with NPO, where forget quality (FQ) is measured by KnowMem and VerbMem on $D_f$, along with fine-tune accuracy (FA) after fine-tuning. As seen, before fine-tuning, both NPO and NPO+ILU (GSM8K) demonstrate strong unlearning effectiveness across MUSE-News and MUSE-Books. After fine-tuning, the unlearning performance (KnowMem and VerbMem on $D_f$) is largely preserved for ILU (staying at low values), while NPO shows a clear increase in these metrics post-fine-tuning.

We thank the reviewer for the detailed summary of our work and contributions. We also appreciate the insightful comments and provide our detailed responses below.

Q1: Confusion in Fig. 4 task vector.
A1: We apologize for any confusion caused in Fig. 4 and our presentation. We will improve them in the revision. The difference between the two angles (or cosine values) shown in Fig. 4 is not intended to explain the robustness gain of ILU over NPO. Instead, these angles serve to validate our geometry plots and are later used in our analysis (between Lines 318-Left and Line 285-Right in our submission). In summary, the sign flip (Line 281-Right) when comparing $\cos(\angle (\tau_{\text{NPO} \to \text{ft}}, \tau_{\text{NPO}})) = -0.41 < 0$ and $\cos(\angle (\tau_{\text{ILU} \to \text{ft}}, \tau_{\text{ILU}})) = 0.09 > 0$ demonstrates that NPO post-fine-tuning results in a much larger deviation from its original unlearning direction (an obtuse angle), whereas ILU remains near orthogonal, better disentangling the fine-tuning effect from the original unlearning and preserving the unlearning direction within the unlearning space after fine-tuning, as shown by our geometric validation in Fig. 4.

Q2: On the significance of task vector analysis for more samples/tasks.
A2: The task vector is defined as the difference between the fine-tuned/unlearned model and its base model over the fine-tuning dataset, as illustrated in Lines 278-287. Therefore, it depends on the entire fine-tuning task rather than a single sample. To better highlight the significance of our analysis, we follow the reviewer’s suggestion to conduct a task vector analysis across additional fine-tuning tasks within the NPO-based unlearning context. See results in Table R1, which reported two metrics: $\cos(\angle (\tau_{\text{NPO} \to \text{ft}}, \tau_{\text{NPO}}))$ and $\cos(\angle (\tau_{\text{ILU} \to \text{ft}}, \tau_{\text{ILU}}))$, as clarified earlier in Fig. 4. A smaller negative value for the former $\cos(\angle (\tau_{\text{NPO} \to \text{ft}}, \tau_{\text{NPO}}))$ indicates that the fine-tuning task vector forms a larger obtuse angle with the unlearning task vector, implying greater conflict between the two. In contrast, the cosine value closer to 0 for the latter $\cos(\angle (\tau_{\text{ILU} \to \text{ft}}, \tau_{\text{ILU}}))$ demonstrates the effectiveness of our method, as it has less conflict between the fine-tuning and unlearning directions.