Revisiting the Past: Data Unlearning with Model State History

1. Problem Statement

Large Language Models (LLMs) are trained on massive web corpora that often contain private data, copyrighted material, misinformation, or low-quality data. When such problematic data points need to be removed to comply with regulations (e.g., the "Right to Be Forgotten") or improve model quality, the standard approach—retraining the model from scratch on a dataset excluding these points—is computationally prohibitive.

Existing Machine Unlearning algorithms attempt to approximate this retraining by modifying the model weights at a lower cost. However, these methods face significant challenges:

• Under-forgetting: The model retains traces of the data to be forgotten.
• Degraded Utility: The model loses capabilities on data it should have retained.
• Inaccuracy: The unlearned model does not functionally match an "ideal model" (one trained from scratch without the target data).
• Limitation of Current Approaches: Most existing methods operate solely on the final model checkpoint, attempting to reverse the influence of data using only the current state, which often leads to noisy or biased estimates of data influence.

2. Methodology: Model State Arithmetic (MSA)

The authors propose MSA, a novel unlearning algorithm that leverages intermediate model checkpoints (artifacts routinely stored by developers for fault tolerance and experimentation) to more precisely estimate and reverse the influence of specific datapoints.

Core Concept

Instead of calculating a "forget vector" based on the final model (which has already internalized the data), MSA calculates the vector based on a checkpoint ($C$) that precedes the introduction of the unlearning targets ($D_f$).

Algorithm Steps

1. Input: A target model $\theta_D$ (trained on full data $D$), an intermediate checkpoint $\theta_0$ (trained on data $D \setminus D_f$), and the forget set $D_f$.
2. Compute Forget Vector ($\vec{\theta}_f$):
   • Fine-tune the checkpoint $\theta_0$ on the forget set $D_f$ to obtain a new state $\theta_1$.
   • Calculate the difference: $\vec{\theta}_f = \theta_1 - \theta_0$. This vector represents the pure influence of $D_f$ on the model parameters before the model had other conflicting knowledge.
3. Apply Unlearning:
   • Subtract the forget vector from the target model: $\theta_{unlearn} = \theta_D - \alpha \vec{\theta}_f$, where $\alpha$ is a scaling factor.
4. Optional Retain Vector ($\vec{\theta}_r$):
   • If a retain set $D_r$ is available, a similar vector can be computed by fine-tuning $\theta_0$ on $D_r$ to get $\theta_2$, where $\vec{\theta}_r = \theta_2 - \theta_0$.
   • The final update becomes: $\theta_{unlearn} = \theta_D - \alpha \vec{\theta}_f + \beta \vec{\theta}_r$.

Key Insight

The authors hypothesize that vectors derived from the final model are entangled with the model's existing knowledge, making them poor estimates of the specific data's influence. By using a "clean" checkpoint from before the data was seen, MSA isolates the semantic effect of the forget set more effectively.

3. Key Contributions

• Novel Framework: Introduction of MSA, which repurposes standard training checkpoints for post-hoc data unlearning, shifting the paradigm from "final-state only" to "state-history aware."
• New Evaluation Metrics: The paper critiques standard metrics like ROUGE-L for unlearning (which can be fooled by syntactic similarity despite factual errors) and introduces GPT-4o-based semantic metrics:
  • Accforget: Measures if the model successfully avoids generating the ground truth for forgotten items.
  • Accrecover: Measures if the model's output aligns with the "ideal model" (the model that never saw the data).
  • Accretain: Measures preservation of knowledge on non-forget data.
• Comprehensive Benchmarking: Evaluation across three major benchmarks: TOFU (fictional authors), RESTOR (restoring knowledge corrupted by misinformation), and MUSE-Books (book memorization and privacy).
• Robustness Analysis: Demonstration that MSA remains effective even when the checkpoint used is hundreds of billions of tokens prior to the unlearning target, though closer checkpoints yield better results.

4. Experimental Results

The authors evaluated MSA against state-of-the-art baselines (NPO, RMU, GradDiff, Task Vector, SatImp, UNDIAL) on OLMo-2-7B and Llama-3.1/3.2 models.

• TOFU Benchmark:
  • MSA variants (e.g., using checkpoints 3.6T tokens prior) achieved competitive or superior performance across all metrics compared to baselines.
  • Unlike baselines that often sacrifice utility (Accretain) to achieve forgetting, or fail to recover the ideal behavior (Accrecover), MSA balanced all three objectives effectively.
  • For example, on the forget10 task, MSA variants achieved ~89% Accforget and ~96% Accrecover, whereas baselines like RMU had high forgetting but very low recovery (refusing to answer rather than forgetting).
• RESTOR Benchmark:
  • MSA significantly outperformed baselines in restoring knowledge corrupted by misinformation.
  • On Llama-3.1-8B, the ideal model accuracy was 64.80%. The corrupted model dropped to 44.31%. MSA recovered accuracy to 48.45% - 49.19%, whereas baselines like NPO and SatImp recovered only to ~44-46%.
• MUSE-Books Benchmark:
  • MSA demonstrated strong performance in Membership Inference Attacks (MIN-K% and MIN-K%++), indicating that the unlearned model no longer recognizes the forgotten books as part of its training data.
  • MSA variants using closer checkpoints (e.g., MSA3859B) matched or exceeded the ideal model's performance on privacy leakage metrics.
• Checkpoint Distance:
  • Experiments showed that while checkpoints immediately preceding the forget set are optimal, checkpoints hundreds of billions of tokens earlier still yield strong results, making MSA practical even for models with infrequent checkpointing.

5. Significance and Conclusion

• Practicality: MSA offers a computationally efficient ($O(|D_f|)$) solution that utilizes artifacts (checkpoints) already present in standard LLM training pipelines, requiring no changes to the original training process.
• Regulatory Compliance: It provides a viable path for model providers to fulfill "Right to Be Forgotten" requests without the cost of full retraining.
• Theoretical Shift: The work challenges the assumption that unlearning must rely solely on the final model state. It suggests that model state history is a critical resource for precise data erasure.
• Future Direction: The paper advocates for model developers to maintain indexes of when models encounter specific information during training, not just for attribution, but to enable advanced unlearning techniques like MSA.

In summary, MSA represents a significant advancement in machine unlearning by leveraging the temporal dynamics of training to achieve more precise, utility-preserving, and functionally faithful data removal.