Benchmarking Unlearning for Vision Transformers

This paper presents the first comprehensive benchmark for machine unlearning in Vision Transformers, evaluating various algorithms across different model architectures, datasets, and protocols to establish performance baselines and analyze how VTs memorize data compared to CNNs.

The Gist

Imagine you have a very smart student who has studied a massive library of books to learn how to recognize objects (like cats, cars, or numbers). This student is a Vision Transformer (VT). They are incredibly good at this job, often better than the older generation of students (called CNNs) who used to rule the classroom.

However, there's a problem. Sometimes, this student memorizes specific, embarrassing, or private details from a few specific books in the library. Maybe they memorized a photo of a celebrity's private diary, or a picture of a biased stereotype. If someone asks the student to "forget" that specific book, the student usually can't just erase that one page without messing up their entire knowledge base. This process of teaching a machine to "unlearn" specific data is called Machine Unlearning.

Until now, researchers had great ways to help the old students (CNNs) forget things, but they hadn't figured out how to do it for the new, super-smart students (Vision Transformers).

This paper is like a giant report card that tests exactly how well different "forgetting techniques" work on these new Vision Transformer students.

The Big Experiment: The "Forget-Me-Not" Test

The researchers set up a massive classroom experiment to see which teaching method works best. Here's how they did it, using some fun analogies:

1. The Students (The Architectures)

They tested two types of Vision Transformers:

• ViT (Vision Transformer): Think of this student as someone who looks at a whole picture at once, like a bird flying high and seeing the whole forest. They are great at seeing the big picture but might get a bit scattered.
• Swin-T (Swin Transformer): This student looks at the picture in small, overlapping windows, like a detective examining clues one by one. They are more organized and structured, acting a bit more like the old-school CNNs.

2. The "To-Forget" List (The Data)

They didn't just test on one type of homework. They used four different datasets:

• CIFAR-10 & 100: Simple picture sets (like a child's drawing book vs. a slightly harder one).
• SVHN: A dataset of house numbers (very clear, easy to read).
• ImageNet: A massive library of millions of complex, real-world photos.

3. The Teaching Methods (The Algorithms)

They tried three main ways to make the student forget:

• Fine-Tuning (The "Re-read" Method): Just tell the student, "Ignore that one book, and re-read the rest of the library." It's simple and often works well for easy tasks.
• NegGrad+ (The "Anti-Force" Method): This is like a coach who says, "Push your brain away from that specific memory while keeping your other skills sharp." It's a more aggressive, mathematical way to erase the memory.
• SalUn (The "Highlighter" Method): This method tries to find exactly which neurons (brain cells) are holding that bad memory and only tweaks those. It's like using a highlighter to find the specific paragraph to erase.

4. The "Memory Detectives" (Proxies)

To know what to forget, the teachers need to know which parts of the student's brain are "stuck" on the bad data. They used "proxies" (clues) to guess this:

• Confidence: If the student is too sure about an answer, they probably memorized it.
• Holdout Retraining: A clever trick where they train a mini-model to see if the student's behavior changes when the bad data is removed.

The Surprising Results

Here is what the report card revealed, translated into plain English:

1. The "Old" Tricks Still Work (But with a Twist)
The methods that worked for the old CNN students also worked for the new Vision Transformers. In fact, the new students were sometimes better at forgetting than the old ones!

• The Winner: NegGrad+ was the star player. It was the most consistent and robust method, especially for the complex ImageNet dataset. It's like the coach who knows exactly how to push the memory out without breaking the student's confidence.
• The Runner-Up: Fine-Tuning was surprisingly effective, especially on simpler tasks. Sometimes, just telling the student to focus on the good stuff is enough.
• The Loser: SalUn (the highlighter method) was a bit shaky. It was good at keeping the student's grades up, but it wasn't very good at actually hiding the secret (it failed the "Membership Inference Attack" test, meaning a hacker could still tell if the bad data was in the brain).

2. The "Big Picture" vs. The "Detective"

• ViT (The Bird): This student preferred the Fine-Tuning method. Because they look at the whole picture at once, their memories are spread out everywhere. Trying to surgically remove one memory (like SalUn) is hard. It's easier to just re-focus their attention.
• Swin-T (The Detective): This student loved NegGrad+. Because they look at the picture in organized chunks, their memories are more localized. The "Anti-Force" method worked perfectly to push those specific chunks away.

3. The "Pre-Training" Advantage
The Vision Transformers were pre-trained on a massive library (ImageNet) before the experiment started. This gave them a superpower: They didn't need to memorize the bad data as much to begin with. Because they already understood the world so well, removing a few bad examples didn't hurt their overall knowledge as much as it did for the older students.

4. The "Continual" Challenge
In the real world, you might need to forget data repeatedly (e.g., every week). The researchers tested if the student would get "brain fog" after forgetting things 5 or 10 times in a row.

• Good News: The student remained stable! The performance didn't drop significantly. They could keep forgetting new things without losing their mind.

The Takeaway for Everyone

This paper is a huge step forward because it gives us a rulebook for making AI safe and private.

• If you have a simple AI: Just use the "Re-read" method (Fine-Tuning). It's cheap and works.
• If you have a complex AI (like a self-driving car or medical scanner): Use the "Anti-Force" method (NegGrad+) combined with the "Detective" style architecture (Swin).
• Don't trust the Highlighter: Be careful with methods that try to surgically remove memories; they might leave traces behind.

In short: We now know that the new, super-smart Vision Transformers can be taught to forget, and we have the right tools to do it safely, fairly, and effectively. This ensures that as AI becomes more powerful, it can also respect our privacy and remove its mistakes when asked.

Technical Summary

1. Problem Statement

Machine Unlearning (MU) is the process of removing the influence of specific "problematic" data (e.g., biased, poisoned, or privacy-sensitive samples) from trained models without retraining from scratch. While MU has been extensively studied for Large Language Models (LLMs), Diffusion Models, and Convolutional Neural Networks (CNNs), there is a significant gap in research regarding Vision Transformers (VTs).

Despite VTs (such as ViT and Swin) becoming dominant architectures in computer vision, existing MU benchmarks and algorithms are primarily designed for CNNs. It remains an open question whether:

• CNN-derived MU algorithms transfer effectively to VTs.
• VTs exhibit different memorization patterns compared to CNNs due to their lack of spatial inductive biases (locality, translation equivariance).
• Memorization proxies (used to identify what to forget) developed for CNNs are valid for VTs.

2. Methodology

The authors established the first comprehensive benchmark for MU on Vision Transformers, systematically evaluating algorithms across multiple dimensions.

Experimental Setup

• Architectures: Two primary VT families were tested: ViT (global self-attention, less similar to CNNs) and Swin-T (hierarchical, windowed attention, more similar to CNNs). Variants included Tiny, Small, and Base sizes.
• Datasets: Four datasets with varying complexity and scale:
  • CIFAR-10 & CIFAR-100: Standard benchmarks for controlled complexity.
  • SVHN: Large-scale but semantically simpler.
  • ImageNet-1K: Used for pretraining and evaluation on a large-scale, high-complexity setting (50k validation images).
• Protocols: Both Single-shot (removing a large batch at once) and Continual (sequential removal of data) unlearning scenarios were evaluated.
• Algorithms: Three representative MU families were tested, both in their vanilla forms and integrated into the RUM (Retrieval-based Unlearning Meta-algorithm) framework, which leverages memorization scores to partition data:
  1. Fine-tune (FT): Retrain only on retained data.
  2. NegGrad+: Gradient-based method that pushes parameters away from the forget set while preserving retained data performance.
  3. SalUn: Saliency-based method that selectively updates parameters influencing the forget set.

Memorization Proxies

Since calculating ground-truth memorization scores (Feldman scores) is computationally prohibitive for VTs, the study evaluated five proxies to estimate memorization:

• Confidence (Conf), Max Confidence (MaxConf), Entropy (Ent), Binary Accuracy (BA), and Holdout Retraining (HR).
• The study verified if these CNN-derived proxies correlate with true memorization in VTs.

Evaluation Metrics

The benchmark uses unified metrics to balance forgetting quality with model utility:

• ToW (Time to Forget): Measures the difference in accuracy between the unlearned model and a model retrained from scratch on the forget, retain, and test sets.
• ToW-MIA: Similar to ToW but replaces forget-set accuracy with Membership Inference Attack (MIA) vulnerability, measuring how well the unlearned model hides the fact that specific data was in training.

3. Key Contributions

1. First VT Unlearning Benchmark: The first systematic evaluation of MU algorithms specifically for Vision Transformers across different architectures, capacities, and datasets.
2. Memorization Characterization: Demonstrated that VTs exhibit long-tail memorization distributions similar to CNNs, despite architectural differences.
3. Proxy Validation: Validated that CNN-derived memorization proxies (specifically Confidence and Holdout Retraining) are effective predictors for VTs, enabling scalable unlearning without expensive ground-truth computation.
4. Algorithm-Architecture Pairing Insights: Identified that different VT architectures favor different unlearning strategies (e.g., ViT favors Fine-tuning, while Swin-T favors gradient-based methods).
5. Continual Unlearning Stability: Proved that VTs can support continual unlearning without significant performance degradation when using robust method-proxy pairs.

4. Key Results & Findings

Memorization & Proxies

• Similarity to CNNs: VTs and CNNs display fundamentally similar memorization patterns (long-tail distributions).
• Proxy Performance: Confidence showed the strongest correlation with true memorization across all models. Holdout Retraining (HR) showed moderate but significant correlation and offered computational advantages.
• Architecture Nuance: Swin-Tiny showed slightly stronger proxy correlations than ViT-Small, likely due to its hierarchical structure resembling CNNs.

Algorithm Performance

• NegGrad+ is the Most Robust: Across most scenarios, especially on complex datasets (CIFAR-100, ImageNet), NegGrad+ (particularly when paired with the HR proxy) emerged as the strongest performer.
• Fine-tune Effectiveness: Simple Fine-tuning performed surprisingly well, especially on ViT architectures and simpler datasets (CIFAR-10, SVHN).
• Salun Limitations: While SalUn achieved good ToW scores, it struggled significantly with ToW-MIA (privacy protection), particularly on ViT architectures and complex datasets. It is less reliable for privacy-sensitive applications in VTs.
• RUM Framework: Integrating algorithms into the RUM framework consistently improved performance over vanilla implementations.

Architecture & Capacity Effects

• ViT vs. Swin-T:
  • ViT: Favors Fine-tuning due to global attention mechanisms leading to diffuse parameter involvement.
  • Swin-T: Favors NegGrad+ due to local, hierarchical attention allowing for more targeted unlearning.
• Capacity: There is a "sweet spot" for model capacity.
  • ViT-Small and Swin-Tiny offered the best balance.
  • ViT-Tiny suffered from underfitting (poor unlearning).
  • Swin-Small showed signs of overfitting on CIFAR-10, reducing unlearning efficacy.
• Pretraining Impact: Pretraining on ImageNet gave VTs a baseline advantage on simpler tasks (higher "Original" performance), but this advantage diminished as task complexity increased (e.g., on ImageNet-1K itself).

Continual Unlearning

• Experiments involving 5 to 10 sequential unlearning steps showed minimal to no degradation in performance (ToW and ToW-MIA) for both ViT and Swin-T when using NegGrad+ with HR. This suggests VTs are stable for continuous data removal.

5. Significance & Practical Takeaways

This work provides a foundational reference for deploying safe and fair AI using Vision Transformers.

• Guidance for Practitioners:
  • For ViT models: Use Fine-tuning combined with the Confidence proxy for efficient unlearning on low-to-moderate complexity tasks.
  • For Swin models (and complex tasks): Use NegGrad+ combined with Holdout Retraining for the most robust privacy protection and performance retention.
  • Avoid SalUn for privacy-critical applications in VTs due to poor MIA resistance.
• Scalability: The findings confirm that MU techniques developed for CNNs can be successfully transferred to VTs, provided the architecture-specific nuances (global vs. local attention) are accounted for.
• Future Research: The paper establishes a reproducible baseline and codebase, encouraging the development of new MU algorithms specifically tailored for the unique inductive biases of Transformer architectures.

In summary, the paper demonstrates that while VTs present unique challenges, they are amenable to machine unlearning, with NegGrad+ and Holdout Retraining emerging as the state-of-the-art combination for high-complexity vision tasks.