Designing to Forget: Deep Semi-parametric Models for Unlearning

This paper introduces deep semi-parametric models (SPMs) that enable efficient, parameter-free machine unlearning by aggregating sample information in a fusion module, achieving competitive task performance while significantly outperforming existing methods in unlearning speed and accuracy on benchmarks like ImageNet.

The Gist

The Big Problem: The "Hard Drive" Dilemma

Imagine you hire a super-smart chef (an AI model) to cook a massive banquet. You give them a recipe book containing 10,000 ingredients and instructions. The chef memorizes the whole book and becomes an expert.

Suddenly, a customer says, "I want to cancel my order for the 'Spicy Tofu' dish, and I want you to forget that you ever learned how to make it."

In the world of standard AI, this is a nightmare. The chef has already mixed the "Spicy Tofu" knowledge into their brain along with the "Grilled Salmon" and "Chocolate Cake." To forget the Tofu, the chef usually has to:

1. Throw away the whole recipe book.
2. Start from scratch.
3. Re-memorize everything else (Salmon, Cake, etc.) without the Tofu.

This is slow, expensive, and wasteful. This is what current AI "unlearning" tries to do: it tries to surgically remove the Tofu memory without breaking the Salmon memory, which is incredibly difficult.

The Solution: "Designing to Forget" (DTF)

The authors of this paper asked a different question: What if we built the chef's brain differently from the start, so forgetting is easy?

They created a new type of AI called Semi-Parametric Models (SPMs). Think of this not as a single brain, but as a Chef + A Magic Index Card System.

How the New System Works

Instead of memorizing every single recipe into their brain, the chef (the AI) has two parts:

1. The Brain (Parametric Part): This learns general cooking skills (how to chop, how to sauté, how to balance flavors). This part stays the same.
2. The Index Cards (Non-Parametric Part): Every single recipe in the training book is written on a physical index card. When the chef needs to make a dish, they don't just rely on memory; they look up the specific cards for that dish and combine them with their general skills.

The "Unlearning" Magic:
When the customer says, "Forget the Spicy Tofu," the chef doesn't need to relearn anything. They simply rip the "Spicy Tofu" index card out of the binder and throw it away.

• Result: The chef instantly forgets how to make Tofu.
• Bonus: The chef is still 100% perfect at making Salmon and Cake because those cards are still there, and the general cooking skills haven't changed.

Why This Paper is a Big Deal

The paper introduces a specific design (called Designing to Forget) that makes this "Index Card" system work for complex tasks like recognizing images or generating art.

Here are the key takeaways, translated:

1. It's Fast (The "Instant Delete" Button)

• Old Way: To unlearn one photo from a dataset of a million, you might have to retrain the AI for days.
• New Way: With this new model, unlearning is as fast as deleting a file from your computer. It takes less than a second. The paper shows it is 10 times faster than existing methods.

2. It's Accurate (The "Perfect Memory" Test)

• Usually, when you try to "delete" something from a complex system, you accidentally break other things.
• The authors tested this by removing specific classes of images (like "Cats" or "Birds"). The new model forgot the Cats perfectly but didn't accidentally start making the Dogs look like Cats. It performed almost exactly like a model that had been retrained from scratch without the Cats ever being seen.

3. It Works for Art and Photos

• They tested this on two things:
  • Classifying: "Is this a cat or a dog?" (The model stops recognizing cats).
  • Generating: "Draw me a cat." (The model stops drawing cats and draws a dog instead, without losing the ability to draw dogs).

The Secret Sauce: "Label Permutation"

The paper mentions a clever trick called Label Permutation.

• The Problem: If the AI sees "Image of a Cat" + "Label: Cat" too many times, it might just memorize the word "Cat" and ignore the actual picture. It becomes lazy.
• The Fix: During training, the researchers shuffle the labels around randomly. This forces the AI to actually look at the pictures and connect them to the index cards, rather than just memorizing the text. This ensures that when you delete a card, the AI actually forgets the concept, not just a text association.

The Bottom Line

This paper proposes a shift in how we build AI. Instead of building a "black box" that is hard to fix, we should build modular systems where data is kept separate from the core logic.

The Analogy:

• Old AI: A library where all the books are melted down and poured into a giant concrete statue. To remove one story, you have to chip away at the concrete, risking the whole statue.
• New AI (SPM): A library with a master librarian and a stack of individual books. To remove a story, you just take one book off the shelf. The librarian (the model) is still there, and the rest of the library is untouched.

This approach makes AI safer, more private, and compliant with laws like GDPR (which gives people the "Right to be Forgotten") without requiring massive amounts of computing power.

Technical Summary

1. Problem Statement

Machine Unlearning (MU) aims to remove the influence of specific training samples from a trained model to comply with privacy regulations (e.g., GDPR) or ethical constraints.

• The Challenge: Current MU approaches for deep learning (parametric models) are computationally expensive and often approximate. They typically rely on fine-tuning, gradient ascent, or model editing to "erase" data, which is difficult because the contribution of individual samples is implicitly encoded in the model weights.
• The Gap: Non-parametric models (like K-Nearest Neighbors) allow for trivial unlearning (simply deleting data points), but they often suffer from inferior task performance compared to deep parametric models.
• The Goal: The authors propose a paradigm shift: instead of developing complex algorithms to unlearn existing models, they ask, "Can we design neural network architectures that are inherently easier and more efficient to unlearn?"

2. Methodology: Designing to Forget (DTF)

The paper introduces Designing to Forget (DTF), a family of Deep Semi-Parametric Models (SPMs). These models combine the high performance of parametric deep learning with the explicit data dependencies of non-parametric models.

Core Architecture

The SPM consists of three distinct modules operating in a two-branch structure:

1. Parametric Module ($f$): A standard deep neural network (e.g., ResNet or UNet) that processes the input test sample $x$ to extract latent features.
2. Non-Parametric Module ($h$): A permutation-equivariant network that processes the entire training set $T$ (or a subset) at test time. It transforms each training sample into an "instance embedding."
3. Fusion Module ($g$): A mechanism that aggregates information from the non-parametric branch (the set of instance embeddings) and fuses it with the parametric branch's latent features.
   • Key Mechanism: The fusion uses a weighted combination (similar to attention) where the prediction depends explicitly on the similarity between the test input and the training instances.
   • Unlearning Mechanism: To unlearn a set of samples $U$, the model simply deletes those samples from the input set $T$ during the forward pass ($T \setminus U$). No weight updates or retraining are required.

Specific Implementations

• For Image Classification: The model uses a ResNet backbone. The non-parametric branch encodes training images, and the fusion module uses cross-attention to weigh the influence of training samples on the test prediction. To handle large datasets, the authors employ retrieval (nearest neighbors) or clustering (averaging instances per class) to keep the input set size manageable.
• For Image Generation (Diffusion): Based on the UNet architecture. The down/up blocks serve as the parametric module. The non-parametric branch encodes training patches. A fusion module replaces the mid-block of the UNet, aligning parametric features with instance embeddings at the patch level.
• Label Permutation Augmentation: To prevent the model from ignoring the training set and relying solely on class labels (turning it back into a parametric model), the authors shuffle label indices during training. This forces the model to utilize the actual image data from the training set.

3. Key Contributions

1. Design-Centric Approach: Shifts the focus from post-hoc unlearning algorithms to architectural design, proposing models inherently suitable for unlearning.
2. Deep Semi-Parametric Models (SPMs): Introduces a framework where predictions explicitly depend on the training set at test time, enabling test-time deletion without modifying learned parameters.
3. Efficiency and Fidelity: Demonstrates that SPMs can achieve unlearning speeds orders of magnitude faster than existing methods while maintaining performance nearly identical to a fully retrained "oracle" model.
4. New Evaluation Metrics: Proposes Hard Prediction Gap (PGH) and Soft Prediction Gap (PGS) to measure the similarity between unlearned and oracle models more strictly than simple accuracy gaps.

4. Experimental Results

The authors evaluated SPMs on CIFAR-10, ImageNet-1K (classification), and CIFAR-10 (generation).

Classification Performance

• Accuracy: SPMs achieve competitive accuracy with standard parametric models (e.g., ResNet18). On CIFAR-10, SPM-C (clustering) reached 94.5% accuracy, matching the baseline ResNet18.
• Unlearning Efficiency:
  • Time: Unlearning takes <1 second (just indexing data), compared to minutes or hours for baseline methods (e.g., SalUn, Gradient Ascent).
  • Fidelity: On ImageNet, SPMs reduced the prediction gap relative to a retrained oracle by 11% compared to existing approaches.
  • Metrics: SPMs achieved near-zero gaps in $\Delta UA$ (Unlearned set), $\Delta RA$ (Remaining set), and $\Delta TA$ (Test set), indicating the unlearned model is virtually indistinguishable from a model retrained from scratch without the data.

Generation Performance

• Quality: SPMs achieved FID scores comparable to DDPM (e.g., 7.04 vs 7.28 on CIFAR-10) when using a sufficient input set size.
• Unlearning: When specific classes (e.g., "Cat") were unlearned, the model stopped generating them and replaced them with other classes, while maintaining high quality for remaining classes.
• Speed: Unlearning time was negligible (<1s) compared to baselines like Selective Amnesia (SA) which took over 18,000 seconds.

5. Significance and Impact

• Privacy Compliance: SPMs offer a practical solution for the "Right to be Forgotten," allowing AI systems to instantly comply with data removal requests without costly retraining.
• Paradigm Shift: The paper argues that unlearning should be a primary consideration in model architecture design, not just an algorithmic afterthought.
• Scalability: By decoupling the "forgetting" process from parameter updates, this approach scales efficiently to large datasets and frequent unlearning scenarios, which is critical for real-world deployment of generative AI and large-scale classifiers.

In summary, the paper demonstrates that by explicitly designing neural networks to rely on training data at inference time, one can achieve perfect unlearning (matching the oracle) with zero computational overhead for the deletion process, while maintaining state-of-the-art task performance.