Constrained Entropic Unlearning: A Primal-Dual Framework for Large Language Models

2025-06-05

AI Summary

To address the challenge of safely forgetting sensitive or outdated information while strictly preserving critical knowledge during large language model (LLM) deployment, this paper proposes the first hard-constrained targeted forgetting framework. We formulate forgetting as a constrained optimization problem: a softmax-free logit-margin flattening loss is applied on the forget set to uniformize output distributions, while strict performance lower bounds are imposed on the retain set. Interpretable dual variables dynamically balance forgetting and retention. Leveraging a primal-dual optimization algorithm combined with parameter-efficient fine-tuning, our method achieves state-of-the-art performance on the TOFU and MUSE benchmarks—delivering high-precision targeted forgetting without compromising (and even improving) retain-set accuracy, alongside enhanced training stability and efficiency.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Large Language Models (LLMs) deployed in real-world settings increasingly face the need to unlearn sensitive, outdated, or proprietary information. Existing unlearning methods typically formulate forgetting and retention as a regularized trade-off, combining both objectives into a single scalarized loss. This often leads to unstable optimization and degraded performance on retained data, especially under aggressive forgetting. We propose a new formulation of LLM unlearning as a constrained optimization problem: forgetting is enforced via a novel logit-margin flattening loss that explicitly drives the output distribution toward uniformity on a designated forget set, while retention is preserved through a hard constraint on a separate retain set. Compared to entropy-based objectives, our loss is softmax-free, numerically stable, and maintains non-vanishing gradients, enabling more efficient and robust optimization. We solve the constrained problem using a scalable primal-dual algorithm that exposes the trade-off between forgetting and retention through the dynamics of the dual variable. Evaluations on the TOFU and MUSE benchmarks across diverse LLM architectures demonstrate that our approach consistently matches or exceeds state-of-the-art baselines, effectively removing targeted information while preserving downstream utility.

Problem

• Unlearning sensitive or outdated data in LLMs
• Avoiding performance degradation during unlearning
• Balancing forgetting and retention via constrained optimization

Innovation

• Constrained optimization for LLM unlearning
• Logit-margin flattening loss for forgetting
• Scalable primal-dual algorithm for trade-off