Machine Unlearning in Large Language Models: Mathematical Foundations of Gradient Ascent, Negative Preference Optimization, and Representation Misdirection

Machine unlearning in large language models (LLMs) addresses the challenge of removing specific training data, copyrighted material, private identifiers, or hazardous capabilities from a pre-trained model without retraining the network from scratch. Because training frontier language models requires millions of GPU-hours and petabyte-scale corpora, exact retraining on a filtered dataset $D_{\text{train}} \setminus D_f$ is computationally infeasible.

Instead, approximate machine unlearning methods modify the model parameters $\theta$ to approximate the statistical distribution of a hypothetical model retrained without the forget set $D_f$, while strictly preserving general utility and knowledge on a retain set $D_r$.

The Formal Unlearning Objective

Let $\theta_0$ denote the pre-trained or fine-tuned model parameters. The data distribution is partitioned into:

• Forget Set ($D_f$): Input-target pairs $(x_f, y_f)$ containing the facts, documents, or skills to be expunged.
• Retain Set ($D_r$): Input-target pairs $(x_r, y_r)$ representing general language capabilities, factual knowledge, and reasoning benchmarks that must remain intact.

The theoretical ideal of exact unlearning produces a parameter configuration $\theta^*$ drawn from the distribution of models trained solely on $D_r$:

$$P(\theta^* \mid D_r) = P(\theta_{\text{retrain}})$$

In generative language modeling, exact parameter equivalence is relaxed to output distribution equivalence:

$$\forall x \in \mathcal{X}, \quad D_{\text{KL}}\left(P_{\theta_{\text{retrain}}}(\cdot \mid x) \;\parallel\; P_\theta(\cdot \mid x)\right) \approx 0$$

Because evaluating $P_{\theta_{\text{retrain}}}$ is intractable, practical unlearning algorithms optimize surrogate objectives that balance forgetting performance on $D_f$ with utility preservation on $D_r$.

1. Gradient Ascent and Retain-Regularized Objectives

The earliest and most direct approach to unlearning is Gradient Ascent (GA), which reverses standard cross-entropy minimization on the forget set.

Vanilla Gradient Ascent (GA)

Given the standard causal language modeling loss:

$$\mathcal{L}_{\text{CE}}(\theta; D_f) = -\mathbb{E}_{(x, y) \sim D_f} \left[ \sum_{t=1}^{|y|} \log P_\theta(y_t \mid x, y_{<t}) \right]$$

Vanilla Gradient Ascent updates parameters in the direction of increasing loss:

$$\mathcal{L}_{\text{GA}}(\theta) = -\mathcal{L}_{\text{CE}}(\theta; D_f) = \mathbb{E}_{(x, y) \sim D_f} \left[ \sum_{t=1}^{|y|} \log P_\theta(y_t \mid x, y_{<t}) \right]$$

$$\theta_{t+1} \leftarrow \theta_t - \eta \nabla_\theta \mathcal{L}_{\text{GA}}(\theta_t) = \theta_t + \eta \nabla_\theta \mathcal{L}_{\text{CE}}(\theta_t; D_f)$$

The Collapse Mechanism of Vanilla GA

While mathematically straightforward, vanilla GA suffers from severe optimization pathology:

• Logit Divergence: The gradient $\nabla_\theta \log P_\theta(y_t \mid x, y_{<t}) = \nabla_\theta z_{y_t} - \sum_k P_\theta(k) \nabla_\theta z_k$ pushes the target logit $z_{y_t} \to -\infty$ while driving non-target logits arbitrarily high.
• Catastrophic Representation Collapse: Because neural representations are dense and polysemantic, unconstrained gradient ascent destroys shared intermediate attention circuits and feed-forward keys, degrading general perplexity on $D_r$ within a few optimization steps.

Gradient Difference (GD) and KL Regularization

To stabilize Gradient Ascent, practical implementations combine forget-set ascent with retain-set descent and Kullback-Leibler divergence constraints against the original reference model $\theta_0$:

$$\mathcal{L}_{\text{GD}}(\theta) = -\mathcal{L}_{\text{CE}}(\theta; D_f) + \alpha \mathcal{L}_{\text{CE}}(\theta; D_r) + \beta \mathbb{E}_{x \sim D_r} \left[ D_{\text{KL}}\left(P_{\theta_0}(\cdot \mid x) \;\parallel\; P_\theta(\cdot \mid x)\right) \right]$$

Here, $\alpha > 0$ enforces retain task accuracy, while $\beta > 0$ penalizes drift away from the baseline model's output distribution on non-target queries.

2. Preference-Based Unlearning: Negative Preference Optimization (NPO)

Direct Preference Optimization (DPO) aligns language models by reparameterizing the reward function $r(x, y) = \beta \log \frac{P_\theta(y \mid x)}{P_{\text{ref}}(y \mid x)}$. Zhang et al. (2024) extended this framework to unlearning by formulating Negative Preference Optimization (NPO).

In standard DPO, the loss over preferred completions $y_w$ and dispreferred completions $y_l$ is:

$$\mathcal{L}_{\text{DPO}}(\theta) = -\mathbb{E}_{(x, y_w, y_l)} \left[ \log \sigma \left( \beta \log \frac{P_\theta(y_w \mid x)}{P_{\text{ref}}(y_w \mid x)} - \beta \log \frac{P_\theta(y_l \mid x)}{P_{\text{ref}}(y_l \mid x)} \right) \right]$$

In pure unlearning settings, a paired high-quality target completion $y_w$ may not exist. NPO treats the forget completion $y_f$ strictly as a negative sample, removing the positive term:

$$\mathcal{L}_{\text{NPO}}(\theta) = -\mathbb{E}_{(x, y) \sim D_f} \left[ \log \sigma \left( -\beta \log \frac{P_\theta(y \mid x)}{P_{\text{ref}}(y \mid x)} \right) \right] = \mathbb{E}_{(x, y) \sim D_f} \left[ \log \left( 1 + \left( \frac{P_\theta(y \mid x)}{P_{\text{ref}}(y \mid x)} \right)^\beta \right) \right]$$

Vanilla Gradient Ascent (GA):
Gradient scale: Unbounded as P_theta(y|x) -> 0
Risk: Logit divergence, output degeneration to repetitive gibberish

Negative Preference Optimization (NPO):
Gradient scale: Bounded by sigmoid derivative sigma(-beta * log-ratio)
Risk: Controlled suppression relative to reference baseline P_ref

Gradient Dynamics of NPO

The parameter gradient for a single token under NPO is:

$$\nabla_\theta \mathcal{L}_{\text{NPO}}(\theta) = -\beta \cdot \sigma \left( \beta \log \frac{P_\theta(y \mid x)}{P_{\text{ref}}(y \mid x)} \right) \cdot \nabla_\theta \log P_\theta(y \mid x)$$

When the model assigns high probability to the forget sequence ($P_\theta(y \mid x) \approx P_{\text{ref}}(y \mid x)$), the gating weight approaches $\sigma(0) = 0.5$, applying strong negative gradient updates. As $P_\theta(y \mid x)$ drops significantly below $P_{\text{ref}}(y \mid x)$, the probability ratio approaches zero, and the gating factor decays exponentially:

$$\lim_{P_\theta(y \mid x) \to 0} \sigma \left( \beta \log \frac{P_\theta(y \mid x)}{P_{\text{ref}}(y \mid x)} \right) = 0$$

This self-regulating property eliminates the logit explosion seen in vanilla GA, preventing catastrophic model degradation while steadily suppressing forget tokens.

3. Representation-Level Unlearning: Representation Misdirection (RMU)

Rather than manipulating output logits, representation-level methods intervene directly on hidden states within intermediate transformer layers. Li et al. (2024) introduced Representation Misdirection for Unlearning (RMU) in the Weapons of Mass Destruction Proxy (WMDP) benchmark.

The RMU Formulation

Let $h_\theta^{(l)}(x) \in \mathbb{R}^d$ denote the hidden activation vector at layer $l$ produced by model $\theta$ for prompt $x$. RMU selects an intermediate layer (typically layer $l \in [d_{\text{layers}}/4, d_{\text{layers}}/2]$) and fine-tunes a subset of weights (such as the feed-forward down-projection matrices) using two simultaneous objectives:

1. Misdirection on Forget Set: Steer the hidden representation of forget prompts $x_f \in D_f$ toward a predetermined, uninformative steering vector $\mathbf{u} \in \mathbb{R}^d$:

$$\mathbf{u} = c \cdot \frac{\mathbf{v}}{\|\mathbf{v}\|_2}, \quad \mathbf{v} \sim \mathcal{U}(0, 1)^d$$

where $c \in \mathbb{R}^+$ controls the perturbation magnitude.

1. Stability on Retain Set: Constrain the hidden representation of retain prompts $x_r \in D_r$ to match the frozen reference activations $h_{\text{ref}}^{(l)}(x_r)$:

$$\mathcal{L}_{\text{RMU}}(\theta) = \mathbb{E}_{x_f \sim D_f} \left[ \left\| h_\theta^{(l)}(x_f) - \mathbf{u} \right\|_2^2 \right] + \alpha \mathbb{E}_{x_r \sim D_r} \left[ \left\| h_\theta^{(l)}(x_r) - h_{\text{ref}}^{(l)}(x_r) \right\|_2^2 \right]$$

Here, $\alpha$ is a high penalty weight (often set between $1000$ and $1200$). By modifying only intermediate projections, downstream layers receive uninformative random activations for forget queries, causing the attention and generation heads to output harmless generic text or refuse the prompt without corrupting general linguistic capabilities.

4. Weight-Space Arithmetic and Task Vector Inversion

Model unlearning can also be framed as vector subtraction in parameter space, extending the task arithmetic framework of Ilharco et al. (2023).

Let $\theta_{\text{base}}$ be the initial pre-trained model. If a model is fine-tuned on the forget set $D_f$ to obtain $\theta_f$, the parameter displacement defines a "forget task vector":

$$\tau_f = \theta_f - \theta_{\text{base}}$$

To unlearn the target domain, the task vector is subtracted with a scaling coefficient $\lambda > 0$:

$$\theta_{\text{unlearned}} = \theta_{\text{base}} - \lambda \tau_f$$

Subspace Truncation (TIES and DARE Inversion)

Direct subtraction can introduce cross-task interference across unrelated parameter coordinates. Advanced weight-space unlearning applies directional filtering:

• TIES-Inversion: Identifies top-$k\%$ redundant parameter updates, resolves sign conflicts against retain gradients, and subtracts only coherent forget components.
• DARE-Inversion: Randomly drops a fraction $p$ of parameter updates in $\tau_f$ and rescales remaining weights by $\frac{1}{1-p}$ before subtraction, minimizing perturbation of shared zero-shot knowledge.

5. Entity Substitution and Generic Hallucination Reinforcement

In domain-specific factual unlearning, such as removing copyrighted literature (Eldan and Russinovich, 2023 in "Who's Harry Potter?"), optimization techniques replace forget entities with synthetic generic counterparts.

The pipeline operates in three distinct phases:

1. Entity Extraction and Mapping: Target entities (e.g., character names, locations, spells) are identified and mapped to plausible generic alternatives (e.g., "Harry Potter" $\to$ "David", "Hogwarts" $\to$ "boarding school").
2. Anchor Inversion: The model is evaluated on prompts involving forget entities to generate completions. Prompts are translated into generic equivalents, and target logits are aligned to generic completions using cross-entropy loss:

$$\mathcal{L}_{\text{sub}}(\theta) = -\mathbb{E}_{x \in D_f} \left[ \log P_\theta(y_{\text{generic}} \mid x) \right]$$

1. Retain Regularization: Predictions on general knowledge prompts are tied to original reference outputs via KL divergence or retain cross-entropy.

6. Evaluation Benchmarks and Fragility of Unlearning

Evaluating whether knowledge has been truly deleted rather than superficially masked requires multi-metric validation suites.

+--------------------------------------------------------------------------------+
|                        Unlearning Evaluation Framework                         |
+--------------------------------------------------------------------------------+
|  Forget Efficacy:                                                              |
|  * Forget Set QA Accuracy (Target -> Near-Random Guessing)                     |
|  * Forget Perplexity / Probability Ratio: P_unlearn(y|x) / P_orig(y|x)         |
|  * Membership Inference Attack (MIA) Resistance                                |
+--------------------------------------------------------------------------------+
|  Retain & Model Utility:                                                       |
|  * Retain Set Perplexity & Accuracy                                            |
|  * Downstream Benchmarks (MMLU, GSM8K, HumanEval, MT-Bench)                    |
|  * Linguistic Coherence & Repetition Penalties                                 |
+--------------------------------------------------------------------------------+
|  Robustness & Adversarial Recovery:                                            |
|  * Re-learning Vulnerability (Few-shot Fine-Tuning Rebound)                     |
|  * Adversarial Jailbreaks & Paraphrased Extraction                             |
|  * Latent Directional Ablation Recovery                                        |
+--------------------------------------------------------------------------------+

Standardized Benchmarks

• TOFU (Task of Fictitious Unlearning): Maini et al. (2024) introduced a synthetic dataset of 200 fictitious author biographies (each containing 20 QA pairs). By controlling the pre-training and fine-tuning data, TOFU provides an exact retrained ground-truth baseline to measure forget quality and retain preservation.
• WMDP (Weapons of Mass Destruction Proxy): Li et al. (2024) compiled 4,157 multiple-choice questions spanning biosecurity, chemical security, and cyberattack methodologies to test hazardous capability removal in frontier models.

Known Vulnerabilities of Approximate Unlearning

Empirical evaluations reveal critical failure modes across current unlearning algorithms:

• Shallow Erasure vs. Deep Deletion: Many unlearning methods alter generation routing without destroying underlying associative weights. Adversarial prompting or soft prompt tuning often recovers unlearned facts.
• Re-Learning Sensitivity: Fine-tuning an unlearned model on as few as 10 to 50 domain examples frequently restores full performance on the forget set, demonstrating that representations remain latent in orthogonal parameter subspaces.
• Over-Generalization: Aggressive unlearning can inadvertently suppress adjacent domains, such as degrading general organic chemistry reasoning while attempting to unlearn chemical weapon synthesis.

Algorithmic Comparison

• Vanilla Gradient Ascent (GA): Reverses cross-entropy on $D_f$. Highly unstable; prone to rapid logit explosion and catastrophic model collapse.
• Gradient Difference (GD + KL): Pairs forget ascent with retain descent and KL divergence constraints. Moderately stable; requires careful hyperparameter balancing between $\alpha$ and $\beta$.
• Negative Preference Optimization (NPO): DPO-derived negative-only loss with sigmoid-bounded gradient updates. Stable optimization; effectively suppresses forget sequences without divergence.
• Representation Misdirection (RMU): Aligns intermediate hidden states to a random steering vector while freezing retain states. Highly effective for hazardous domain suppression; retains strong MMLU performance.
• Task Vector Subtraction: Subtracts fine-tuned forget weights from base model parameters. Computationally fast; sensitive to task entanglement and weight scaling.
• Entity Substitution: Trains the model to produce generic completions for target entity prompts. Effective for targeted narrative unlearning; high data preparation overhead.

Sources

• TOFU: A Task of Fictitious Unlearning for LLMs (Maini et al., 2024)
• Negative Preference Optimization: From Catastrophic Collapse to Effective Unlearning (Zhang et al., 2024)
• The WMDP Benchmark: Measuring and Reducing Malicious Use With Unlearning (Li et al., 2024)
• Who's Harry Potter? Making LLMs forget (Eldan & Russinovich, 2023)
• Editing Models with Task Arithmetic (Ilharco et al., 2023)
• Rethinking Machine Unlearning for Large Language Models (Liu et al., Nature Machine Intelligence 2025)