PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Taha Entesari, Arman Hatami, Rinat Khaziev, Anil Ramakrishna, Mahyar Fazlyab

arXiv:2506.05314v2cs.CLcs.AIcs.LG

TL;DR

LLM unlearning must remove sensitive, outdated, or proprietary information, but existing scalarized objectives combine forgetting and retention. The paper proposes constrained unlearning with logit flattening and primal-dual optimization. Evaluations report strong forgetting while preserving utility on TOFU and MUSE.

  • Problem

    Existing unlearning methods combine forgetting and retention in a scalarized trade-off, while the paper identifies limits in controlling retained-data degradation and tuning that trade-off.

    Source: §1 · p. 1, §3 · p. 4

  • Method

    The method constrains retention loss and uses a logit-flattening forget loss with warm-started primal-dual updates to manage the trade-off without extra computational overhead.

    Source: §1 · p. 2, §3.2 · p. 5, §3.4 · p. 6

  • Results

    On TOFU, the method consistently achieves the highest forget success with best or second-best utility; on MUSE Books, it achieves the most forgetting among non-degenerate methods while maintaining high utility.

    Source: §4 · p. 8

  • Takeaways & Limitations

    Evaluations on TOFU and MUSE indicate the method can suppress memorized responses while preserving retained capabilities, often matching full retraining at a fraction of the cost.

    Source: §5 · p. 10

    The constraint threshold ε is sensitive and generally depends on both the model and the data.

    Source: §5 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

LLM unlearning must remove targeted information while preserving retained-data performance, but existing approaches can damage utility or require costly optimization. The paper motivates constrained unlearning and introduces a logit-flattening loss with primal-dual optimization.

1 Introduction

  • Training data can contain sensitive, copyrighted, or ethically problematic content, motivating selective removal without retraining large models from scratch.Unlearning is also relevant to regulatory compliance, harmful-content avoidance, privacy, and copyrighted material.Source: §1 · p. 1
  • Forget-set optimization can sharply degrade retained-data performance, while multi-objective methods may require impractical multiple gradient evaluations and scalarization has theoretical and tuning limitations.The degradation is especially severe when the forget set is small relative to retained data; scalarization may miss Pareto trade-offs and requires task-specific coefficient tuning.Source: §1 · p. 2
  • The paper formulates unlearning as an ε-constrained problem and proposes logit-margin flattening plus a primal-dual algorithm to enforce retention while tracking the forgetting-utility trade-off.The method is presented as avoiding multi-gradient overhead and heuristic loss balancing.Source: §1 · p. 2, §1.1 · p. 3

1.1 Related Works

  • Prior approaches include retraining [2], gradient ascent [50], optimization-based penalties, representation interventions, and prompt or logit methods [35] [13] [22] [11].The paper describes limitations including computational cost, utility loss, semantic distortion, or brittleness, depending on the approach.Source: §1.1 · pp. 2–3

1.2 Notation

  • The notation defines π as a language model that maps tokenized prompts to logits and probabilities, with πθ denoting a model parameterized by θ.The prompt has N tokens and embedding dimension D; the output logits have vocabulary size V.Source: §1.2 · p. 4

2 Problem Formulation

The formulation seeks a model that retains reference-model behavior on the retain set while eliminating knowledge associated with the forget set. It frames these goals as a bi-objective problem and describes linear scalarization as one way to combine them.

2 Problem Formulation

  • The target model should preserve the reference model's behavior on the retain set while eliminating knowledge of the forget set; perfect unlearning means eliminating related information, not merely lowering target-response likelihood.The forget and retain sets are disjoint subsets of the reference model's training or fine-tuning data.Source: §2 · p. 4
  • Retraining on retain data alone is an ideal but costly and time-consuming baseline, so practical unlearning seeks a model close to the reference with forget-set influence removed.The passage identifies frequent requests, including those involving outdated data or copyright, as a reason retraining costs matter.Source: §2 · p. 4
  • The problem is posed as balancing forget loss L_fgt against retain loss L_rtn, with linear scalarization using λ to control their trade-off.The formulation places the model in a compact function class Π.Source: §2 · p. 4

3 Constrained Entropic Unlearning

The section formulates unlearning as maximizing forgetting subject to an explicit retention budget, then develops a primal-dual solver and a logit-flattening forget loss. The formulation makes retention degradation directly controllable while the loss drives forget-set outputs toward uniformity.

3 Constrained Entropic Unlearning

  • The constrained formulation maximizes forgetting while requiring retain-set loss to stay within a user-specified degradation threshold.Unlike scalarization, it specifies a retention budget directly rather than balancing objectives with a tuned λ.Source: §3 · p. 4

3.1 Lagrangian Relaxation and the Dual Problem

  • Strong duality holds under convexity and strict feasibility; the reference model guarantees strict feasibility when the tolerance α is positive.The finite-dimensional parameterized search space is nonconvex, so strong duality may not hold there.Source: §3.1 · p. 5

3.2 Proposed Method: Primal-dual with Warm Start

  • A warm-started primal-dual solver fully minimizes the Lagrangian initially, then alternates single primal and dual steps without extra overhead over linear regularization.The dual update adjusts the trade-off in response to retain-loss constraint violation.Source: §3.2 · p. 5

3.3 Retention Loss

  • The retain loss uses standard cross-entropy on retained examples.The constraint limits its degradation relative to the reference model.Source: §3 · p. 4, §3.3 · p. 6

3.4 Logit Flattening for Efficient Forgetting

  • Directly maximizing cross-entropy on the forget set can cause gradient explosion and catastrophic collapse.The section motivates a more stable alternative that induces high uncertainty in model outputs.Source: §3.4 · p. 6
  • The logit-margin flattening loss penalizes peaked pre-softmax logits, driving predictions toward uniformity without computing entropy explicitly.It avoids log-softmax, maintains nonzero gradients near uniform predictions, and is convex in the logits.Source: §3.4 · pp. 6–7
  • The logit margin directly bounds the model’s maximum softmax probability.The loss’s zero point occurs only when all logits are equal, corresponding to maximal entropy.Source: §3.4 · p. 7

4 Experiments

Experiments on TOFU and MUSE compare PDU with established unlearning baselines across model architectures, using automatic and LLM-judged metrics. PDU achieves strong forgetting while preserving retained utility, with results also highlighting how abstention can distort forgetting scores.

Experimental setup

  • The evaluation covers TOFU Forget10 and MUSE Books and News, using six LLM architectures and comparing PDU with established unlearning baselines.The evaluated models range from LLaMA 3.2 1B to LLaMA 2 13B; baselines include GradDiff, DPO, NPO, SimNPO, and RMU.Source: §4 · p. 7

Evaluation

  • TOFU evaluation uses model utility and forget success, while MUSE uses Retain ROUGE and Forget ROUGE, alongside LLM-judged evaluations.Aggregated success is the harmonic mean of metrics; the reported LLM judge scores generated responses against ground-truth responses.Source: §4 · pp. 7–8

Results

  • PDU achieves the highest forget success and best or second-best model utility across TOFU scales and models, and performs competitively on MUSE.On MUSE Books, it achieves the most forgetting among methods that have not degenerated; on News, it provides viable Pareto-optimal points.Source: §4 · p. 8, Table 1 · p. 9
  • For larger models, GradDiff's LLM-judged forget score is 10 despite forget success of 0, as abstention harms utility and other judged metrics.The authors attribute the discrepancy to the model producing no text for forget-set prompts.Source: §4 · p. 8
  • LLM-judged metrics generally correlate with traditional metrics, supporting their use as indicators of model capabilities and unlearning success.The authors caution that forgetting and retained utility must be considered together when assessing success.Source: §4 · pp. 8–9

5 Conclusion

The conclusion presents constrained optimization, a logit-margin flattening loss, and a primal-dual solver as a framework for balancing forgetting and retention. Evaluations on TOFU and MUSE support its effectiveness, while the authors identify threshold sensitivity and limited robustness testing as boundaries.

Conclusion

  • The framework separates forgetting and retention through constrained optimization, using a logit-margin flattening loss and a primal-dual solver with interpretable dual dynamics.The loss encourages uniform predictive distributions on the forget set; evaluations on TOFU and MUSE report suppression of memorized responses while preserving retained capabilities.Source: §5 · pp. 9–10
  • The authors report that PDU often matches full retraining at a fraction of the cost on TOFU and MUSE.Table 2 covers MUSE News and Books, but its caption provides no result values.Source: §5 · p. 10

Limitations

  • The constraint threshold ε is sensitive and generally depends on both the model and data; overly tight or loose settings can impair unlearning or degenerate the model.With tight constraints, dual updates can inhibit meaningful unlearning; loose constraints can cause deviation that limited training epochs cannot repair.Source: §5 · p. 10
  • The study did not evaluate resilience to relearning attacks or jailbreak attempts and reports a slight fluency reduction on easier TOFU under PDU.The authors also note that extensive hyperparameter tuning was not possible because of resource demands.Source: §5 · p. 10

A.1 Proof of Theorem 3.1

The proof analyzes an output probability by isolating the maximum logit and bounding a denominator under a simplifying condition on the other logits.

A.1 Proof of Theorem 3.1

  • The proof defines i as the conditional output probability π(y_t|x, y_<t)_i and assumes a unique maximum logit i⋆.Source: §A.1 · p. 15
  • Under that assumption, the denominator is convex and is minimized when all other logits are equal.Source: §A.1 · p. 15
  • The derivation substitutes a lower bound into the probability expression, then takes a first-order Taylor expansion around δ = 0.The displayed bound includes a term involving V and δ.Source: §A.1 · p. 15

B Experiments

The appendix uses radar charts to visualize unlearning across multiple evaluation dimensions. It reports that PDU is competitive or best across those dimensions, with chart area reflecting balanced forgetting and utility.

B.1 Visualization

  • The radar charts compare model utility, forget success, and LLM-judged metrics, scaling utility and forget success into [0, 10].Source: §B.1 · p. 16
  • PDU is competitive or best across evaluated dimensions, and its chart area reflects strong forget success alongside high model utility.Source: §B.1 · p. 16

B.2 Implementation Details

Implementation uses different A100 setups by model size, repository-default fine-tuning settings, and shared settings for unlearning from target checkpoints.

B.2 Implementation Details

  • LLAMA 3.2 1B/3B experiments used one A100 40GB GPU, while the other models used eight A100 80GB GPUs.The latter setup ran on a p4de.xlarge AWS EC2 instance.Source: §B.2 · p. 16
  • Fine-tuning used repository defaults, including a 10^-5 learning rate, bfloat16 precision, and a scheduler that warms up for one epoch before decaying to zero.The shared optimizer was paged_adamw_32bit.Source: §B.2 · p. 17
  • Unlearning starts from a target model and uses the shared fine-tuning settings, with task- and model-specific settings listed separately for PDU.The paper points to Table 4 for further unlearning settings and Table 5 for PDU hyperparameters.Source: §B.2 · p. 17

B.3 LLM Judge

The evaluations use OpenAI’s API as an LLM judge with a fixed prompt and generation setup; locally hosted judges were less reliable in following instructions.

B.3 LLM Judge

  • The authors used OpenAI’s API with gpt-4.1-mini-2025-04-14 for LLM-based evaluation.The evaluation setup fixed the prompt, set temperature to 0.3, and capped generation at 1024 tokens.Source: §B.3 · p. 17
  • Locally hosted judges, including Llama 3.1 8B Instruct, were less reliable at following evaluation instructions and often produced extraction errors.They sometimes required multiple calls at varying temperatures to obtain valid scores.Source: §B.3 · p. 17
  • Across experiments and invocations, the OpenAI judge consistently followed instructions and returned scores in the specified JSON format.Source: §B.3 · p. 17
  • The supplementary material provides sample judge evaluations for forget, retain, and world-fact cases on TOFU using PDU with LLAMA 3.1 8B.Source: §B.3 · p. 17

B.4 Evaluation Details

The evaluation uses different metrics for TOFU and MUSE and explains why some measures can miss important distinctions in forgetting or retained utility.

B.4 Evaluation Details

  • The harmonic mean is used because it responds more sharply to values near zero, making performance drops more visible.The supplied passage introduces the harmonic mean formula but does not provide its complete expression.Source: §B.4 · p. 17
  • TOFU forgetting is assessed on the forget set with 1−likelihood, 1−ROUGE score, and truth ratio.Source: §B.4 · p. 17
  • Likelihood complements ROUGE because a response can have low textual similarity yet retain non-trivial likelihood, as observed for DPO unlearning.The authors treat likelihood as essential for detecting this partial retention.Source: §B.4 · p. 20
  • Truth ratio measures likelihood assigned to fabricated alternatives; a higher ratio indicates greater forgetting of the ground truth and greater association with false responses.Source: §B.4 · p. 20
  • Forget quality is not a main metric because it requires a retrained model and a low score does not necessarily mean unlearning failed.The metric can also assign near-zero forget quality when outputs become incoherent, even if memorized content is no longer generated.Source: §B.4 · pp. 20–21
  • For MUSE, ROUGE on independent prompts assesses knowledge retention and verbatim memorization; lower ROUGE on incomplete-sequence completions indicates reduced verbatim memory.Source: §B.4 · p. 21

B.5 Further Experiments

Further TOFU experiments show strong PDU performance across retention settings, with a small-model forgetting trade-off in the most aggressive setting that improves with longer training.

B.5 Further Experiments

  • For TOFU forget05/retain95, PDU achieves the best performance across all evaluated models, while alternatives trade forgetting success against model utility.For the 1B model, NPO and GradDiff compromise utility, SimNPO shows minimal forgetting, and RMU and DPO achieve marginally stronger forgetting at reduced utility.Source: §B.5 · pp. 21–22
  • For TOFU forget01/retain99, PDU has the strongest or second-strongest unlearning on 3B, 7B, and 8B models, but not the best forgetting on 1B.Even on 1B, utility remains on par with the original model.Source: §B.5 · pp. 21–22
  • Increasing unlearning epochs lets PDU outperform baselines on the 1B forget01/retain99 task.The authors cite this ablation as evidence of robustness under extended training.Source: §B.5 · p. 22

B.6 Ablation Studies

This section reports ablations on unlearning duration, utility, membership-inference and jailbreak evaluations, and notes that statistical significance was not established due to experiment costs. It also summarizes the checklist’s statements on societal implications and LLM use.

B.6 Ablation Studies

  • Across LLAMA 3.2 1B and 3B models and TOFU difficulty subsets, results consistently improve as unlearning epochs increase, without significant utility degradation; Retain and Relevance scores remain stable.Holdout accuracy can reveal signs of overfitting or performance degradation.Source: §B.6 · p. 24
  • Across Loss, Min-K, Ref, and Zlib attacks, the method outperforms baselines on LLAMA 3.2 1B and 3B or performs comparably and near optimally on LLAMA 3.1 8B and Gemma 7B.It shows unexpected vulnerability to GN attacks on LLAMA 3.1 8B; min K++ values are elevated for all models except the 8B variant, a pattern the authors plan to investigate.Source: §B.6 · p. 27, Table 15 · p. 30
  • For larger models, the method performs best or comparably near perfectly on Exact Memorization and Extraction Strength metrics.The study also reports PrivLeak, for which smaller absolute values indicate better unlearning.Source: §B.6 · pp. 27–28, Table 15 · p. 30
  • In single-turn jailbreak tests using prompts that include ground-truth hints, PDU achieves the best or second-best Knowledge and Verbatim Removal scores across model architectures.Because valid answers may infer or restate hinted facts, the evaluation relies primarily on an LLM judge rather than ROUGE.Source: §B.6 · pp. 28–29, Table 16 · p. 31

NeurIPS Paper Checklist

  • The authors report that the experiments lack error bars or other statistical-significance evidence because the required LLM compute and high-end GPU resources make such experiments infeasible.They say general compute-resource information is provided in the Supplementary Material.Source: NeurIPS Paper Checklist · pp. 35–36
  • The authors say they foresee no societal implications arising solely from this work.Source: NeurIPS Paper Checklist · p. 37
  • LLMs are used as judges of unlearning effectiveness and to generate jailbreak prompts, with their use described in the paper and Supplementary Material.Source: NeurIPS Paper Checklist · p. 39
Loading 2506.05314v2…