Cross-Domain Generalization in Machine Unlearning via Label-Conditioned Energy Magnitude Regularization

AI summary

Overview

Research area: Machine unlearning in computer vision, specifically cross-domain and cross-class generalization of forgetting, studied through label-conditioned energy-based models (EBMs).

Technical level: Advanced. The paper assumes familiarity with energy-based models, gradient-ascent unlearning, membership-inference attacks, and self-supervised features such as DINOv2.

Scope: The paper asks whether forgetting a single (class, domain) pair spreads to that same class in other visual domains and to visually similar classes, and proposes a similarity-weighted propagation mechanism that controls how far the spread reaches.

What This Paper Is About

Most unlearning methods treat the target class as isolated: they suppress one class and try to leave everything else untouched. This paper instead studies what actually happens to the rest of the model when a class is forgotten, using a label-conditioned energy-based model where each class has its own directly observable energy score. The goal is to show that forgetting naturally leaks into other domains of the same class and into visually similar classes, and then to introduce a mechanism that makes this leakage controllable rather than incidental.

Key Contributions

1. A label-conditioned EBM with energy magnitude regularization for studying forgetting. Each class receives an explicit, directly observable energy score, so forgetting can be measured as a shift in the energy landscape rather than only as a change in accuracy.

2. The finding that forgetting is not isolated. Erasing a class in one domain (sketch) also removes it from real, clipart, and painting, consistently across forget targets, and the effect further spills to the nearest visually similar classes.

3. A DINOv2 similarity and PCA subspace weighting mechanism combined with a scope mask. This turns the incidental spread into a controllable one, scaling how far forgetting reaches and steering it toward chosen domains or similar classes.

4. A two-setting evaluation. Cross-domain forgetting on a 26-class subset of DomainNet across four visual domains, and isolated single-class forgetting on all ten CIFAR-10 classes with the propagation term disabled.

Main Findings

• Forgetting generalizes across domains. Forgetting a class in the sketch domain also erases it from real, clipart, and painting. The authors' method reaches a forgetting error of 96.1 on the target cell and 83.7 on the same class in other domains for tiger, 99.4/98.5 for lion, and 100.0/99.2 for scissors (Table 1). The abstract summarizes this as forgetting error reaching 98% and 99% for lion and scissors. Retrain and Finetune barely reduce the class in other domains (F/oth. below 5), because the class still exists there and is simply relearned from the other domains.

• NegGrad also forgets across domains, but less selectively. NegGrad reaches 100.0 on both the target cell and other domains for all three targets, but reduces retain accuracy to 85.3 (tiger), 70.0 (lion), and 72.4 (scissors), and has no way to limit how far forgetting reaches. This suggests cross-domain spread comes from the shared backbone rather than from the proposed loss alone.

• A moderate retain-accuracy cost for the proposed method. Retain accuracy (Mem) falls from 96.0 to 76.7 for tiger, 95.9 to 80.9 for lion, and 96.1 to 77.2 for scissors.

• Isolated forgetting on CIFAR-10 is complete with preserved utility. All ten classes reach 100% forget error. Mean retain accuracy goes from 69.8 before to 68.7 after unlearning, giving a mean utility of 98.5% with a minimum of 95.5%. The paper notes the pretrained model tops out near 70% on this backbone.

• Membership inference is at chance. The membership-inference AUC on forget samples averages 0.496 across the ten CIFAR-10 classes, indistinguishable from the 0.5 chance value.

• Energy reshaping explains the mechanism. Before unlearning, the forget class and the retain set both have mean energies near −1.4. After unlearning, the forget-class energy rises to about +8.3 while the retain energy moves only from −1.4 to −1.9. Retain energy preservation (rank-order agreement with the pretrained model) averages 0.854, where 1.0 means identical.

• Spillover concentrates on the nearest neighbour. Lion drops 75% when tiger is forgotten, and pliers drops 46% when scissors is forgotten. Over all 25 retain classes the correlation between DINOv2 similarity to the forget class and accuracy drop is r = 0.76 for tiger, r = 0.42 for scissors, and r = 0.10 for lion, with Spearman values of 0.44, 0.10, and 0.21 respectively.

• The spread is strongest above a similarity threshold. For tiger/sketch, retain classes above a DINOv2 similarity threshold of 0.25 lose a median of 28%, against 11% for the rest.

• Backbone capacity governs spread. With the backbone frozen and only heads trained, the similarity–degradation correlation is r = 0.52, with 60.6 forgetting and 92.1 retain accuracy. LoRA r8 on layer4 (157k trainable parameters) gives r = 0.535, and LoRA r8 on layer3+4 (227k) gives r = 0.642. Training the full last layer (8.5M parameters) raises the correlation to 0.76, with 96.1 forgetting and 77.9 retain accuracy.

• The propagation term adds controllable reach. With full capacity, turning the term off (λp = 0) gives 73.6 forgetting, 82.8 retain accuracy, and r = 0.63; turning it on (λp = 3) gives 96.1 forgetting, 77.9 retain accuracy, and r = 0.76. Even with the term off, forgetting still spreads, confirming that part of the effect is intrinsic to the shared backbone.

Methodology in Plain English

The authors replace the usual softmax classifier with a label-conditioned energy-based model. For an image and a candidate label, the model outputs a single scalar energy, and the predicted class is the one with the lowest energy. An image and its label are turned into a score by encoding the image with a ResNet-18 backbone, projecting it to a 128-dimensional embedding, multiplying it elementwise with a learned label embedding, and passing the result through a linear head. The backbone is initialized from ImageNet weights; during pretraining, only the last residual stage (layer4), the projection, the label embedding, and the energy head are trained, with a low-rank (LoRA) variant available as an alternative.

Training happens in two stages. Stage 1 pretrains the model on all classes and domains with a contrastive energy loss that compares the correct label's energy against N = 10 sampled incorrect labels with a margin m0 = 1, using Adam and early stopping on validation accuracy. This produces a frozen reference model E0. Stage 2 copies E0 and trains the copy to forget a single (class, domain) cell, using four loss terms applied together at every step:

• A forget term that raises the energy of correct forget pairs above every other label by a margin, mirroring the pretraining loss with the sign reversed.
• A retain anchor that penalizes the squared difference between current and reference energies on retain samples, normalized by the average pretrained magnitude so the penalty does not depend on energy scale.
• A global margin that pushes forget energies above retain energies by the same margin over equal-size minibatches.
• An energy magnitude regularizer that keeps the forget and margin terms from inflating energies without limit.

A fifth term directs propagation. The authors use a frozen DINOv2 ViT-B/14 encoder to compute the mean feature over the forget and retain sets, apply PCA to the centered forget-cell features, and keep the top k = 5 principal directions as a subspace. Each retain sample gets a weight between 0 and 1 equal to the clipped fraction of its centered feature that lies in that subspace. The same forget-style push is then applied to retain samples, scaled by this weight, so samples resembling the forget cell are pushed hardest while distant samples get almost no push and stay anchored to E0. A mask decides the scope: keeping weights only for retain samples of the forget class restricts propagation to that class in other domains, while keeping weights for all retain samples lets forgetting spread to visually related categories.

Experiments use a 26-class subset of DomainNet (sampled from its 345 classes) across four of its six domains (real, sketch, clipart, painting), forgetting tiger, lion, and scissors one at a time in the sketch domain. The same four-term objective without the propagation term is applied to each of the ten CIFAR-10 classes in turn for isolated forgetting. DomainNet unlearning runs for 1500 steps with batch size 32, learning rate 10⁻⁴, weights λf = 1, λr = 15, λm = 1, λe = 10⁻³, margin m = 5, propagation weight λp = 3, and a 20% holdout for membership inference. Baselines are Retrain (from scratch on the retain set, the gold reference), NegGrad (gradient ascent on the forget set), and Finetune (continued training on the retain set).

Why This Matters

Impact on research. Unlearning is usually evaluated as a narrow, surgical operation: did the target class disappear, and did everything else stay the same? This paper argues that the operation is not surgical at all. Because features are shared, forgetting leaks across domains and to nearby classes, and whether that leakage is desirable depends on the application. Reframing the spread as something to steer rather than suppress, and measuring it with a per-class energy score, gives the field a different evaluation lens than accuracy alone. It also connects unlearning to a broader question about what "removing" a concept from a shared representation actually means.

Real-world applications:

• Privacy and identity removal. The paper explicitly raises removing a person's face, where some spread to closely related identities may be preferable to prevent the model from indirectly reconstructing the forgotten identity through highly similar facial features.
• Safety and content removal. Removing a concept captured in one modality or style (for instance a sketch, artwork, or clipart rendering) while ensuring the concept does not survive in photographs.
• Data correction. Withdrawing a class whose training data was wrong or mislabeled across multiple collection sources at once, instead of patching one source at a time.
• Model curation. Tailoring a large pretrained classifier to the subset of classes a downstream task needs, and deliberately shaping which neighbouring concepts get suppressed.

Industry relevance. Multidomain datasets with the same object rendered as photo, sketch, clipart, and painting are common in product imagery and content platforms, so an unlearning request rarely arrives for a single domain. A method that documents and controls how forgetting travels across domains, and that reports what it costs in retain accuracy, is directly relevant to anyone deploying unlearning as a compliance or data-governance service.

Future Directions

• Scaling beyond small backbones. The authors state their experiments use one DomainNet subset and CIFAR-10 on small backbones, and that these show the effect is real and steerable rather than how far it can be pushed. Larger backbones and more datasets are the natural next step.

• Tighter control over propagation. The spread to similar classes concentrates on the nearest neighbour and its magnitude varies with the target (lion drops 75% under tiger, tiger only 9% under lion), so the authors describe it as better read as a control to turn up or down than as a precise setting.

• Recovering the retain accuracy cost. Removing a class lowers accuracy slightly on retained classes because shared features mean erasing one disturbs the others; the paper leaves open how much of this cost is avoidable.

• Applications where controlled spread is the goal. The authors name removing a person's identity together with its closest look-alikes as a setting where deliberately directed propagation is what the user wants, rather than an unwanted side effect.

Target Audience

Researchers and practitioners working on machine unlearning, energy-based models, or domain generalization who need to understand how forgetting behaves beyond a single target class. It is most useful to readers already comfortable with classifier training and evaluation, and to engineers responsible for privacy, safety, or data-correction workflows who need to reason about whether a deletion request has actually been honored across every representation of a concept. Readers primarily interested in an off-the-shelf unlearning method will find less here, since the paper's center of gravity is the diagnosis of propagation and a mechanism for controlling it.

Authors’ abstract

Machine unlearning removes the influence of specific data from a trained model. However, most methods treat the forgotten concept as isolated. In this paper, we study what happens to the rest of the model when a class is forgotten, using a label-conditioned energy-based model (EBM) that assigns per-class energies, making the effect directly observable. We forget a class by raising the energy of its image-label pairs, training with a forget term, a retain anchor to the pretrained model, a global margin, and an energy regularizer that stops the energy magnitudes from growing without limit. A propagation term applies the same forget signal to retain samples, weighted by each sample's DINOv2 similarity to the forget class, so forgetting reaches images that resemble it and leaves the rest untouched. We evaluate on two benchmark datasets: 1) On a subset of DomainNet across four visual domains, we forget tiger, lion, and scissors one at a time. Forgetting a class in the sketch domain also erases it from real, clipart, and painting, with forgetting error reaching 98% and 99% for lion and scissors, and the effect carrying over to the most similar class. 2) On CIFAR-10, we turn off the propagation term and forget each of the ten classes on its own. Forgetting is complete (100%), while the other nine classes retain 98.5% of their pre-unlearning accuracy on average.