🤖 machine learning

Unlearning at Scale: State-Exact Trace-Preserving Deletion in Billion-Parameter Language Models

This paper demonstrates that billion-parameter language models can achieve state-exact, trace-preserving deletion by replaying training from a token store excluding specific examples, provided the original execution was instrumented to record provenance and an uncontaminated checkpoint is retained, though this method does not guarantee computational efficiency for dispersed deletion requests.

Original authors: Abdullah X

Published 2026-08-14

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Abdullah X

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are baking a massive, multi-layered cake for a huge party. In the world of artificial intelligence, this cake is a "language model," a computer program trained to write, chat, and solve problems by "eating" terabytes of text. Usually, once the cake is baked, you can't easily take out a single ingredient—say, a specific type of vanilla bean—without ruining the whole thing. If you try to scrape it out, the layers might collapse, or the flavor might change in unpredictable ways. This is the problem of "machine unlearning": how do you make an AI "forget" specific data it was trained on, especially when laws (like the GDPR) say people have a right to have their data erased?

Scientists have been trying to solve this by either retraining the whole cake from scratch without that ingredient (which is slow and expensive) or by trying to surgically remove the flavor (which often leaves weird chemical traces). But there's a catch: computers don't just learn from what they eat; they learn from the order they eat it, the temperature of the oven, and the exact timing of when they stir the batter. If you change the order of the ingredients, even if you remove the bad one, the final cake might taste different from a cake that was baked perfectly without that ingredient from the start. This paper dives into the messy, technical kitchen of AI to ask a very specific question: Can we prove that a cake baked after removing an ingredient is exactly, bit-for-bit, the same as a cake that was baked with a plan to ignore that ingredient all along?

The "Time-Travel" Cake Recipe

This paper, titled "Unlearning at Scale," is like a team of super-precise bakers who decided to test a very strict recipe for removing ingredients. They aren't just trying to make the cake look like it forgot the vanilla; they want to prove that the cake is mathematically identical to a cake that was never supposed to have vanilla in the first place.

To do this, they built a special "time-travel" system. Imagine you are baking a cake, but before you start, you write down a super-detailed, unchangeable recipe card. This card doesn't just list ingredients; it records the exact second you crack each egg, the precise temperature of the oven at every minute, and the specific order in which you add the flour. This is what the paper calls a "trace-preserving" execution.

Now, imagine a customer calls and says, "I want to remove the vanilla beans from my order." In a normal kitchen, you might just skip the vanilla and keep mixing. But in this paper's kitchen, the bakers have a trick. They keep the recipe card exactly the same. When the recipe says "add vanilla at step 5," the bakers still go to step 5, but instead of adding real vanilla, they add a "dummy" ingredient that tastes like nothing and weighs nothing. They do this without changing the order of the other steps, the oven temperature, or the stirring speed.

The Big Test: Billion-Parameter Cakes

The researchers tested this idea on three different "cakes" (AI models) of varying sizes:

A small cake called Pythia 160M (160 million parameters).
A giant cake called Pythia 2.8B (2.8 billion parameters).
A different kind of cake called Llama 3.2 1B (1 billion parameters).

They ran a massive experiment. First, they baked a "perfect" cake where they followed the recipe but added a "zero-weight" dummy for the forbidden ingredient. This is the "Oracle"—the gold standard of what the cake should look like. Then, they baked a second cake using their "redacted" method: they physically removed the forbidden ingredient from the pantry (the data store) but followed the exact same recipe card, substituting the dummy whenever the recipe called for the missing item.

The Result: A Perfect Match

The result was shocking in its precision. When they compared the two cakes, they found that the "redacted" cake was bit-for-bit identical to the "Oracle" cake.

For the Pythia 2.8B model, they checked 2,775,208,960 individual numbers (the "model state") and found zero differences.
For the Llama 3.2 1B model, they checked 1,498,482,688 numbers and found zero differences.
Even the "optimizer state" (which is like the memory of how the baker adjusted the recipe as they went) was exactly the same.

This proves that if you have the right "recipe card" (the execution plan) and you start from a clean checkpoint (a snapshot of the cake before the bad ingredient was added), you can remove data and get a result that is mathematically indistinguishable from a cake that was never supposed to have that data.

The Catch: It's Not a Magic Wand

However, the paper is very careful not to call this a "magic wand" for deleting data quickly. The author points out a major limitation: Time.

If the forbidden ingredient was added early in the baking process, or if the ingredients were scattered randomly throughout the recipe, the bakers have to go back and re-bake almost the entire cake from the beginning to get the exact same result.

In their tests, if they asked to remove a random 5% of the data, they had to replay nearly 100% of the training steps.
The paper explicitly states this does not establish "cheap deletion." It's not a fast fix; it's a slow, precise reconstruction.

What This Does NOT Mean

The author is very strict about what their results don't prove:

It's not a behavioral guarantee: Just because the numbers are identical doesn't automatically mean the AI will "behave" differently in a way that satisfies a legal erasure request. The paper ran some standard tests to see how the AI behaved, but they treat these as just "descriptive" notes, not proof that the method works for real-world legal deletion.
It's not a universal fix: This only works if you planned ahead. You can't take an old AI model that wasn't built with this special "recipe card" and magically make it work this way later.
It's not a privacy shield: The paper admits that while the numbers are gone from the specific replay store, it doesn't prove the data is erased from every backup, cache, or hard drive in the universe.

The Takeaway

In simple terms, this paper is a proof-of-concept for a very specific, very strict type of "forgetting." It shows that if you build your AI training process with a rigid, unchangeable plan and a way to record every single step, you can surgically remove data and get a result that is mathematically perfect. It's like proving you can swap a real vanilla bean for a fake one in a complex recipe and end up with a cake that is indistinguishable from one that never had vanilla.

But the paper also warns us: this is a heavy, slow process. It requires planning ahead, and if the data you want to remove is scattered all over the place, you might as well just bake the whole cake again. It's a powerful tool for scientists who need to prove exactly what happened in a computer program, but it's not yet a quick button for deleting data in the real world.

Technical Summary: Unlearning at Scale: State-Exact Trace-Preserving Deletion in Billion-Parameter Language Models

1. Problem Formulation

The paper addresses the challenge of machine unlearning in large language models (LLMs) by reframing the problem from a behavioral optimization task to a state-exact systems question. While existing literature often defines unlearning as producing a model behaviorally indistinguishable from one trained without specific data, this approach suffers from ambiguity in deterministic training pipelines.

In modern training, the realized state depends not only on the dataset but on execution controls: sample order, random number generation (RNG) consumption, gradient accumulation boundaries, learning-rate schedules, and hardware-specific numerical kernels. Simply removing data and retraining (densely repacking) alters these controls, resulting in a different execution trace and a different final state. Consequently, "the model trained without $F$ " does not identify a unique reference state unless the execution or reference distribution is explicitly specified.

The paper defines a trace-preserving deletion counterfactual: a scenario where the original training execution's logical controls (microbatch slots, seeds, accumulation boundaries, optimizer steps) are fixed, but the loss contribution of requested identifiers is set to zero. The goal is to verify if a replay system can reconstruct this specific counterfactual state bit-for-bit, even when the physical data rows corresponding to the deleted identifiers are absent from the replay store.

2. Methodology

2.1 Trace-Preserving Deletion Definition

The author formalizes a Trace-Preserving Deletion ($TD$) where:

The execution plan $P$ remains immutable.
Every microbatch plan record is retained, but identifiers in the forget set $F$ are mapped to fixed dummy sequences with zero sample weight.
Retained identifiers, seeds, learning rates, and logical optimizer-step indices remain unchanged.
Empty-segment rule: At gradient accumulation boundaries, an optimizer transition occurs only if at least one retained weighted contribution exists in that segment. If a segment contains only zero-weighted (deleted) slots, the logical boundary is preserved, but the optimizer step is skipped to avoid altering optimizer moments (e.g., AdamW moments) with zero gradients.

2.2 System Design

The proposed system relies on prospective instrumentation and physical redaction:

Immutable Execution Plan: The original training run records a JSONL plan containing ordered sample identifiers, per-microbatch RNG seeds, learning rates, and accumulation flags.
Write-Ahead Log (WAL): During training, a compact 32-byte binary WAL records the execution state (ordered-ID tag, seed, learning rate, step index, flags, CRC32) alongside an ordered-ID manifest.
Physical Redaction: Before replay, a new token store is materialized containing only retained rows. The immutable plan still references the deleted identifiers as logical slot markers.
Replay Mechanism: During replay, the system intercepts requests for deleted identifiers. Instead of looking up the missing data, it substitutes the fixed dummy sequence with zero weight. This ensures the numerical microbatch (tokens, labels, weights) is identical to the trace oracle, despite the physical absence of the data rows.

2.3 Verification Protocol

The system verifies Trace Exactness by comparing the terminal state of the redacted replay against a separately executed Trace Oracle. The Oracle runs the same counterfactual logic (zero-weight dummies) from an eligible checkpoint. Exactness is defined as bit-identical model parameters ( $\theta$ ) and optimizer states ( $\Omega$ ) in the recorded data types.

3. Key Contributions

Execution-Conditioned Definition of Exact Deletion: The paper formalizes unlearning as a trace-preserving counterfactual, separating the definition of the target state from the method of achieving it (e.g., dense repacking).
Prospectively Instrumented Replay System: A system that records execution provenance (WAL + manifest) and supports replay from physically redacted stores while maintaining deterministic control inputs.
Bit-Exact Billion-Parameter Evidence: Empirical demonstration of state exactness across three endpoints:
- Pythia 160M: Exact across 162M parameters and optimizer state for four deletion geometries (early, middle, late, random).
- Pythia 2.8B: Exact across 2.77 billion parameters for a random 5% deletion request.
- Llama 3.2 1B + TOFU: Exact across 1.49 billion parameters after omitting 400 of 4,000 examples.
Replay-Locality Analysis: A structural analysis showing that while localized late requests allow for efficient suffix replay, large, randomly dispersed requests force nearly full replay, limiting the computational efficiency of this approach for dispersed deletions.
Auditable Research Artifact: A public release containing source revisions, data hashes, execution plans, and state hashes, treating standardized behavioral metrics (OpenUnlearning) as descriptive diagnostics rather than causal evidence.

4. Results

State Exactness: In all tested scenarios, the redacted replay produced a model-and-optimizer state that was bit-identical to the separately executed trace oracle. This included exact matches in optimizer moments (Adam states), which are critical for future training continuity.
Counterfactual Distinction: Experiments on Pythia 160M confirmed that dense repacking (retraining on $D \setminus F$ with new batching) and physical filtering (removing rows before the forward pass) result in non-zero state differences compared to the trace oracle. This validates that "training on retained data" is not equivalent to "trace-preserving deletion."
Scalability: The method successfully scaled to 2.8 billion parameters (Pythia 2.8B) and cross-family architectures (Llama 3.2), proving that the invariant holds at billion-parameter scales under pinned single-GPU environments.
Cost Analysis: The control provenance (WAL + manifest) is small relative to model checkpoints (0.019% – 0.20%). However, the replay cost is locality-dependent. For a random 5% deletion in a 2.8B model, the replay required restarting from step zero, taking nearly the same time as the original run. The paper notes that dispersed requests structurally force nearly full replay.

5. Significance and Claims

The paper makes a modest but precise contribution to the field of machine unlearning:

Scope of Claim: The author explicitly states they do not claim "cheap deletion" or universal superiority over approximate unlearning methods (e.g., Gradient Ascent, NPO). The method requires prospective instrumentation and may incur full replay costs for dispersed requests.
State vs. Behavior: The paper rigorously separates state exactness from behavioral forgetting. While the released states were evaluated using the OpenUnlearning TOFU benchmark, the author treats these results as descriptive diagnostics only. They explicitly avoid causal claims that trace preservation causes specific forgetting behaviors, citing the lack of matched controls (e.g., a matched retained-only baseline and pre-fine-tuning metrics).
Auditable: The primary significance is the establishment of an auditable, algorithmic definition of deletion. By fixing the execution contract and proving bitwise equality to a declared counterfactual, the system provides a verifiable guarantee that a specific execution trace was neutralized, distinct from claims about parameter-level data absence.
Limitations: The author acknowledges that the method does not guarantee deletion from backups, caches, or other storage layers, nor does it satisfy legal obligations (e.g., GDPR Article 17) on its own. It is a systems primitive for reconstructing a specific historical state, not a general-purpose policy solution.

In summary, the paper demonstrates that billion-parameter state exactness is achievable for a specific, trace-preserving deletion counterfactual, provided the training system is instrumented prospectively and the execution environment is pinned. It establishes that "unlearning" must be defined by the specific counterfactual program being executed, not just the set of retained data.