Source-linked AI summary
BLADE: Bilevel Low-rank Augmented-Lagrangian Erasure for LLM Unlearning
Md Toufikuzzaman, Ahmad Mousavi, Dongwon Lee
TL;DR
LLM unlearning must remove targeted knowledge without damaging retained capabilities, yet existing methods can diverge, oscillate, or fail under changing and repeated conditions. BLADE uses constrained bilevel optimization with LoRA adapters, clamped entropy, and asymmetric augmented-Lagrangian protection. It improves composite scores across three benchmark families and remains stable under scaling and sequential unlearning where the strongest baseline collapses.
Problem
Existing LLM unlearning methods struggle with unbounded forget losses, fixed-weight balancing, and robustness under scaling or repeated application.
Method
BLADE combines LoRA-confined repair-first bilevel optimization, clamped entropy with zero gradients after sufficient uncertainty, and asymmetric augmented-Lagrangian retain protection.
Results
BLADE dominates across TOFU, MUSE, and KnowUndo, improving average composite scores over the strongest baselines by 6% on TOFU, 9% on MUSE Books, and 7% on KnowUndo.
Takeaways & Limitations
BLADE maintains stable performance under 4× scaling and 4 sequential unlearning steps on MUSE News while the strongest competing method collapses under both conditions.
Takeaways & Limitations
BLADE targets factual and knowledge-level unlearning and remains susceptible to re-learning attacks requiring white-box weight access and original forget data.
Abstract
from arXiv · showhide
Existing LLM unlearning methods struggle with robustness: unbounded forget losses degrade model coherence, fixed-weight balancing cannot adapt as retain difficulty shifts mid-training, and methods that work on one benchmark falter under scaling or repeated application. We propose BLADE, a constrained bilevel framework whose three mechanisms give smooth, predictable control over the optimization landscape: a clamped-entropy forget loss whose gradient is exactly zero once a token reaches sufficient uncertainty; an asymmetric augmented Lagrangian that permanently ratchets retain protection after any violation; and a bilevel structure confined to LoRA adapters that repairs retain damage before each forgetting step. BLADE dominates across three benchmark families, improving average composite scores over the strongest baselines by $6$% on TOFU, $9$% on MUSE Books, and $7$% on KnowUndo, and it remains stable under $4\times$ scaling and $4$ sequential unlearning steps on MUSE News where the best competing method collapses entirely.
1 Introduction
Machine unlearning must remove requested, hazardous, or copyrighted knowledge while preserving remaining capabilities across repeated corrections. BLADE addresses recurring robustness failures with constrained optimization and shows stronger, stable performance across benchmarks and stress tests.
- Models require recurring post-hoc knowledge removal for privacy, copyright, safety, and misinformation-related corrections without full retraining.
- Shared forget–retain representations make each unlearning step vulnerable to collateral damage that can compound across repeated applications.
- BLADE combines repair-first bilevel optimization, asymmetric augmented-Lagrangian retain protection, clamped entropy, and LoRA confinement to target recurring failure modes.
- 6% on TOFU, 9% on MUSE Books, and 7% on KnowUndo: BLADE improves average composite scores over the strongest baselines.
- 4× scaling and 4 sequential unlearning steps on MUSE News: BLADE maintains stable performance while the strongest baseline collapses catastrophically under both conditions.
- BLADE training follows warm-up, spike-and-ratchet, and smooth-convergence phases, unlike baselines whose dynamics oscillate without self-correction.
2 Related Work
Prior unlearning work spans classical data-influence removal, LLM-specific objectives, parameter-efficient updates, and constrained bilevel methods. BLADE distinguishes itself through structural LoRA confinement, repair-first optimization, asymmetric dual updates, and bounded clamped entropy.
- Origins and classical methods: Classical machine unlearning selectively removes specific training-data influence from already-trained models, originating in classification systems and later extending toward generative models.
- LLM-specific unlearning: LLM-specific methods include gradient ascent, retain-regularized losses, alternative-token fine-tuning, preference optimization, and random-target activation steering.
- LLM-specific unlearning: Existing objectives lack a bounded forget loss with a per-token stopping criterion: gradient ascent diverges, NPO saturates, and unclamped entropy pushes indefinitely.
- Parameter-efficient unlearning: LoRA has supported several parameter-efficient unlearning methods, while BLADE uses its low-rank bottleneck as a structural constraint against catastrophic drift.
- Bilevel and constrained optimization: BLADE differs from related bilevel and constrained methods through asymmetric ratcheting, retain repair in the inner loop, and clamped entropy with a hard per-token deadzone.
- Benchmarks: The evaluation covers TOFU, MUSE, and KnowUndo for complementary tests of synthetic factual associations, naturalistic memorization, and domain-specific copyright or privacy unlearning.
3 Method
BLADE formulates unlearning as constrained bilevel optimization: fast inner-loop retain repair precedes small outer-loop forgetting within LoRA adapters. Adaptive augmented-Lagrangian protection and clamped entropy bound drift and stop optimization for sufficiently forgotten tokens.
- 3.1 Problem Setup: BLADE seeks high-entropy outputs on forget data while keeping retain cross-entropy below a threshold ε.
- 3.2 LoRA Parameterization: LoRA adapters are attached to projection matrices while pretrained weights remain frozen, restricting updates to a low-rank trainable subspace.
- 3.2 LoRA Parameterization: With rank r=16 on a 7B model, approximately 0.2% of parameters are trainable; disabling adapters preserves the original model for rollback and subsequent requests.
- 3.3 Bilevel Formulation: The inner loop performs K fast SGD steps on retain batches to repair damage before the outer loop applies constrained forgetting.
- 3.3 Bilevel Formulation: The outer objective combines forget loss with an augmented-Lagrangian inequality constraint on retain performance, using a slower learning rate so inner repair can keep pace.
- 3.3 Bilevel Formulation: BLADE reverses prior bilevel assignments by placing unconstrained retain repair inside and small forgetting updates outside, limiting cumulative representational drift.
- 3.4 Augmented Lagrangian: After a retain violation, λ increases by ρ|r| but decays only by αρ|r| when satisfied, with α ≪1, creating a persistent ratchet against oscillation.
- 3.4 Augmented Lagrangian: The retain budget ε is scaled from average retain loss during one inner loop, avoiding manual calibration across benchmarks.
4 Experiments
Across TOFU, MUSE, and KnowUndo, BLADE generally improves forgetting–retention performance and remains robust under adversarial and repeated-use stress. Its constrained bilevel design, LoRA parameterization, and clamped forget loss each address distinct failure modes, while increasing memory efficiency at higher wall-clock cost.
- Main results: BLADE improves average composite scores over the strongest baselines by 6% on TOFU, 9% on MUSE Books, and 7% on KnowUndo.Across individual evaluations, BLADE achieves strong HM and HMJ scores, although PDU leads on MUSE News before stress testing.
- Robustness: BLADE remains stable through 4× scaling and four sequential unlearning steps on MUSE News, while PDU collapses at 4× scale and after two steps.BLADE’s HM stays within 0.523–0.552; PDU reaches HM=0.005 at 4× scale and drops from 0.58 to 0.02 after two sequential steps.
- Ablations: BLADE’s clamped entropy, augmented Lagrangian, and LoRA parameterization are jointly necessary for stable unlearning under the ablations.Removing ALM causes HM=0.233, full fine-tuning yields HM=0.459 with retain collapse, and alternative forget losses produce HM values from 0.000 to 0.160.
- Ablations: The repair-first bilevel structure prevents repeated ALM over-penalization, although its benefit is modest on MUSE Books and more consistent on MUSE News.Swapping the bilevel roles causes λ to diverge; K=3 changes HM from 0.819 to 0.823 on MUSE Books and consistently improves performance across an MUSE News sweep.
- Training dynamics: BLADE follows warm-up, spike-and-ratchet, and smooth-convergence phases, contrasting with baselines’ recurring oscillations.The asymmetric ALM permanently increases λ after retain violations, and the resulting dynamics remain stable across broad parameter ranges.
- Additional evaluations: BLADE achieves competitive robustness across extraction, jailbreak, membership-inference, and GCG threat models, while using 2–5× less peak GPU memory but roughly 1.5–2× more wall-clock time.Its performance degrades under re-learning attacks using the original forget data.
5 Conclusion
BLADE is a constrained bilevel framework for robust, smooth, and manageable LLM unlearning. Across benchmarks and stress tests, it maintains strong results and stability beyond single-setting gains.
- BLADE combines clamped entropy, asymmetric augmented-Lagrangian retain protection, and LoRA confinement into a self-correcting unlearning system.These mechanisms produce bounded forgetting, adaptive retain protection, and constrained representational drift.
- Its optimization follows consistent warm-up, spike-and-ratchet, and smooth-convergence dynamics across benchmarks, scales, and dataset types.
- BLADE achieves strong results across TOFU, MUSE, and KnowUndo while maintaining stable performance under 4× scaling and four sequential unlearning steps.The results support robustness across benchmarks, model scales, and repeated application.
6 Limitations
BLADE’s evaluation is bounded by its tested task types, threat models, and modalities. The paper identifies conceptual domain removal, re-learning attacks, and multimodal models as unresolved boundaries.
- BLADE has not been tested on conceptual or representation-level domain removal, where forget–retain entanglement may exceed its selective targeting ability.
- BLADE is susceptible to re-learning attacks in which an adversary has white-box weight access and the original forget data.The paper proposes adversarial outer objectives and dual-threshold clamping as possible future directions.
- BLADE has only been evaluated on text-only language models, leaving its effectiveness on multimodal models untested.
B Theoretical Analysis
The theoretical analysis establishes that clamped entropy is bounded, has controlled gradients, and stops pushing tokens once they reach the uncertainty threshold. It also characterizes selective uncertainty and the retain-only stationary regime.
- B.1 Bounded Loss and Gradient of Clamped Entropy: Clamped entropy bounds each per-token forget loss between 0 and τ log V under τ ∈(0,1).The bound follows from 0 ≤ H_t ≤ log V and the ReLU construction.
- B.1 Bounded Loss and Gradient of Clamped Entropy: Under bounded logits and a Lipschitz logit map, the per-token gradient norm is bounded by (2B_z + log V) · G.
- B.1 Bounded Loss and Gradient of Clamped Entropy: When H_t ≥ τ log V, the clamped entropy gradient is exactly zero, unlike unbounded gradient-ascent and unclamped entropy objectives.
- B.2 Selective Uncertainty: At a stationary point where every forget token satisfies H_t > τH_max, the retain loss alone governs the parameters.
- B.2 Selective Uncertainty: With τ = 0.7, the remaining 30% of maximum entropy acts as a confidence budget for retaining low-entropy predictions.
C.1 Data
The evaluation uses multiple benchmarks, standardized metrics, composite harmonic means, fixed experimental protocols, and additional robustness procedures. Data and scoring cover forget quality, retain quality, leakage, and response coherence.
- C.1 Data: TOFU evaluates Llama-3.2-1B/3B-Instruct across three forget splits, while MUSE uses Llama-2-7B across Books and News settings.KnowUndo evaluates copyright and privacy unlearning on Llama-2-7B-chat.
- C.1 Data: Experiments use five random seeds and report mean ± standard deviation, with baseline settings drawn from OpenUnlearning and benchmark-specific tuning for PDU and BLURNPO.
- C.1 Data: TOFU metrics measure retain accuracy and forget-set behavior; MUSE measures forget knowledge, verbatim memorization, and retain knowledge.KnowUndo uses forget-set and retain-set ROUGE-L together with MMLU accuracy.
- C.1 Data: Composite scores use harmonic means that penalize methods for trading off one evaluation axis against another.The KnowUndo judge composite combines forget leakage, retain accuracy, and retain response quality.
- C.1 Data: BLADE shares most hyperparameters across benchmarks and scales, varying only the outer learning rate and constraint budget across benchmarks.The reported shared settings include τ=0.7 and LoRA on all seven projections per layer.
- C.1 Data: Additional procedures calibrate the outer learning rate and add inner recovery steps when retain loss exceeds 2ε after an outer update.
D Extended Results
The extended results report complete TOFU and KnowUndo evaluations, showing that BLADE maintains low forget leakage while achieving the strongest KnowUndo aggregate scores.
- TOFU 3B: TOFU 3B results across all splits show BLADE and PDU maintaining near-zero forget leakage while preserving utility.Forget leakage is reported as Prob ≤0.01.
- KnowUndo: BLADE achieves the highest HM on both KnowUndo copyright and privacy domains.
- KnowUndo: BLADE also achieves the highest KnowUndo judge HM, consistent with the main paper findings.
E Adversarial Robustness
The adversarial evaluation tests extraction, jailbreak, membership inference, and optimization-based attacks, with BLADE generally retaining strong robustness while remaining vulnerable to re-learning.
- Evaluation scope: Adversarial extraction is limited to TOFU because OpenUnlearning natively supports this evaluation only for that dataset.
- Membership inference: BLADE is the only method with consistently positive PrivLeak scores, indicating no membership leakage.
- Re-learning: 69%: BLADE’s HM decreases by 69% on average after re-learning, compared with 61% for PDU.
F Training Dynamics
The inner-loop ablation evaluates K across eight ε-multiplier settings, while training dynamics show stable convergence for BLADE and oscillations for PDU.
- Inner-loop ablation: K=3 consistently outperforms K=0 across eight ε-multiplier settings on MUSE News.
- Inner-loop ablation: K=6 provides no additional gain over K=3, indicating diminishing returns from further inner-loop steps.
- Training dynamics: BLADE follows a three-phase training pattern, while PDU exhibits recurring oscillations.
G Hyperparameter Robustness
BLADE remains stable across broad hyperparameter sweeps and qualitative evaluations, with adaptive dual-variable trajectories supporting consistent performance and coherent retention.
- Sweep results: BLADE’s maximum HM spread across a single parameter is 0.045 on TOFU, 0.063 on KnowUnDo, and 0.07 on MUSE Books.The sweeps cover eight values for each of five core hyperparameters across three benchmarks.
- Sweep results: Three of five parameters produce HM spreads below 0.01 on TOFU.
- Adaptive compensation: The dual variable adapts its trajectory to compensate for hyperparameter changes, explaining the robustness observed across benchmarks.
- Inner-loop ablation: K=3 consistently outperforms K=0 across eight ε-multiplier settings, while K=6 shows diminishing returns.
- Qualitative behavior: BLADE produces non-linguistic tokens on forget queries, reflecting uncertainty rather than confidently wrong answers, while retain outputs remain coherent.
- Qualitative behavior: BLADE answers correctly on MUSE Books retain examples, whereas the other listed methods fail.