UMD Team Develops Precise ‘Undo Button’ for AI Memory

Imagine trying to remove a single drop of red dye from a gallon of purple paint. For developers of large language models (LLMs), that has long been the challenge of “unlearning”—the process of removing specific information from an AI system after it has already been built.

Once sensitive personal information, copyrighted text, or harmful misinformation becomes embedded in an AI model, it spreads across billions of internal connections. Until now, the most reliable solution was often the most extreme: discard the model and retrain it from scratch, a process that can cost millions of dollars and consume enormous amounts of energy.

As governments and technology companies face mounting pressure to comply with privacy laws, copyright protections and emerging “right to be forgotten” standards, new research from the University of Maryland and the Max Planck Institute for Software Systems provides an alternative—a method that functions like a surgically precise rollback tool for AI memory.

Led by Soheil Feizi, an associate professor of computer science with an appointment in the University of Maryland Institute for Advanced Computer Studies, the team created a framework that can selectively erase unwanted information from an AI model while preserving its overall intelligence and reasoning skills.

Its study, presented in April at the International Conference on Learning Representations, introduces a technique called Model State Arithmetic (MSA). The approach uses saved training snapshots—known as checkpoints—to identify exactly how problematic data shaped a model during training and then reverse those changes with remarkable precision.

“We realized that the model’s own development history contains all the information we need to fix it,” said Feizi. “By using these training snapshots as a guide, we can identify exactly how specific data points altered the AI’s behavior and simply reverse those changes. It’s a much more elegant solution than trying to force a finished model to forget through trial and error.”

Most existing unlearning methods can have side effects researchers call “brain damage,” where models lose general knowledge, reasoning ability or even coherent language generation.

MSA takes a different approach: During training, developers save checkpoints as safeguards against crashes or hardware failures. Feizi and his collaborators discovered those virtual snapshots could be repurposed as a historical record of how information became encoded inside the model.

The system identifies a “clean” checkpoint from before the model encountered problematic data. Researchers then briefly retrain that earlier version on the unwanted material to measure how the model changed. The resulting mathematical signature—what the team calls a “forget vector”—is subtracted from the final model, effectively rolling back the unwanted information while leaving the rest of the system intact.

The researchers tested MSA using industry benchmarks designed to evaluate machine unlearning and found the system consistently removed targeted information while preserving the model’s reasoning and conversational performance.

Remarkably, the team found that the technique remained effective even when using checkpoints saved very early in training, suggesting that even infrequent backups may provide enough information to support precise unlearning.

“AI models shouldn’t be permanent black boxes that can never be corrected once they’re trained,” said Keivan Rezaei, a fourth-year doctoral student in computer science at UMD and the paper’s lead author. “Our method shows that it’s possible to precisely remove harmful or sensitive information while preserving the intelligence the model gained from everything else. That’s a critical step toward building AI systems people can actually trust.”

In addition to Feizi and Rezaei, the research team included Mehrdad Saberi, a doctoral student in computer science at UMD, and Abhilasha Ravichander of the Max Planck Institute for Software Systems.