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

TL;DR

This work reframes unlearning in large language models as an $\varepsilon$-constrained optimization problem, explicitly separating forgetting from retention to avoid unstable trade-offs. It introduces logit-margin flattening, a softmax-free, convex-in-logs forgetting loss, and solves the constrained problem via a scalable primal-dual algorithm with warm-start, exposing forgetting-retention dynamics through the dual variable. The approach is instantiated with a retention constraint $\mathcal{L}_{\text{rtn}}(\pi, \mathcal{D}_{\text{rtn}}) \leq \varepsilon$ and a forgetting objective $\mathcal{L}_{\text{fgt}}^{\mathrm{LM}}$, ensuring robust forgetting while preserving utility. Empirical results on TOFU and MUSE benchmarks across multiple model families show that Constrained Entropic Unlearning (PDU) matches or surpasses state-of-the-art baselines in removing targeted information while maintaining downstream performance, with analysis revealing practical guidance on setting the constraint and future avenues for extension.

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, all without any extra computational overhead. 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.

Figures (7)

• Figure 1: (left) Comparison of different methods on the TOFU dataset on Llama 3.2 3B: Model Utility vs. forget Success; see \ref{['sec:experiments']} for explanation of the metrics. (right) Overview of our methodology and contributions. Unlearning is cast as a constrained optimization, then solved using primal dual optimization. We use our novel logit-flattening loss for the forgetting task. When the retention loss is larger than the pre-specified threshold, dual updates increase the value of $\lambda$, increasing the effect of the retention loss. When the retention loss is in the desired range, the dual updates reduce $\lambda$ so that the optimization can tackle the forget loss more effectively.
• Figure 2: Radar chart of unlearning evaluation for the TUFO (retain90/forget10) dataset.
• Figure 3: Radar chart of unlearning evaluation for the MUSE-News dataset.
• Figure 4: LLM Judge prompt for evaluations of forget data
• Figure 5: LLM Judge prompt for evaluations of retain data

Theorems & Definitions (2)

• Proposition 3.1
• Remark B.1