POUR: A Provably Optimal Method for Unlearning Representations via Neural Collapse

AI summary

Overview

• Research area: Computer vision / machine unlearning, with a theoretical grounding in Neural Collapse (NC) representation geometry.
• Technical level: Advanced. The paper develops propositions and a theorem about simplex Equiangular Tight Frames (ETFs) and derives a closed-form projection operator; readers need comfort with linear algebra, Gaussian class-conditional models, and representation-similarity metrics.
• Scope: The paper reformulates class forgetting as removing a direction from a model's internal feature space, proves that an orthogonal projection yields a provably optimal forgetting operator under Neural Collapse assumptions, and validates the resulting method (POUR) on CIFAR-10/100, PathMNIST, ImageNet with CLIP, and VOC2012 segmentation.

What This Paper Is About

Machine unlearning asks how to make a trained model forget a specific class or set of images without retraining from scratch. The authors argue that most existing methods only adjust the final classifier logits while leaving the internal feature extractor mostly untouched, so the "forgotten" concept can still be recovered by probing the representations. The paper's goal is to define and solve forgetting at the representation level, using the geometric structure that deep classifiers develop at convergence.

Key Contributions

1. Representation-level unlearning formulation and metric. The authors extend weak unlearning to the representation level (Definition 2.1), requiring the unlearned model's feature distribution to be close to a retrained reference model's under a discrepancy measure. They introduce the Representation Unlearning Score (RUS) to quantify both forgetting and retention fidelity in feature space.
2. A three-term decomposition. Proposition 2.2 bounds the feature-distribution discrepancy as a sum of three interpretable terms: class separation, forgotten-class discrepancy, and retained-class discrepancy, exposing a three-way interplay among forgetting efficacy, retention fidelity, and class separation.
3. Two new Neural Collapse properties. Proposition 3.1 shows the simplex ETF structure is a sufficient condition for Bayes optimality under isotropic Gaussian class-conditional assumptions (it maximizes the minimum pairwise angle and the multiclass angular margin of a nearest-class-mean classifier). Proposition 3.2 shows the orthogonal projection of a simplex ETF with one vertex removed remains a simplex ETF of size C−1 in the complementary subspace.
4. The POUR algorithm with an optimality proof. POUR comes in two variants: POUR-P, a closed-form orthogonal projection on the classifier weight direction of the forgotten class, and POUR-D, a teacher-student distillation scheme that uses only the forget set to push the projection into the feature extractor. Theorem 4.2 proves POUR-P is optimal under the paper's definition of representation-level weak unlearning.

Main Findings

• CIFAR-10 (ResNet-18, Table 1): POUR-P reaches Acc_r 94.97 ± 0.16, Acc_f 0.00 ± 0.00, and the highest AUS of 1.01 ± 0.00, above the Retrained Model's AUS of 1.00. POUR-D reaches Acc_r 92.86 ± 1.02, Acc_f 0.37 ± 0.64, AUS 0.97 ± 0.00. For comparison, DELETE reaches AUS 0.92 ± 0.02, Boundary Shrink 0.81 ± 0.01, Gradient Ascent 0.80 ± 0.01, and Random Label 0.75 ± 0.00.
• Representation-level scores on CIFAR-10: POUR-D achieves CKA_f^(r) 0.31 ± 0.01, CKA_r^(r) 0.94 ± 0.00, and RUS^(r) 0.47 ± 0.01, versus Boundary Shrink (CKA_f^(r) 0.28, RUS^(r) 0.42), Random Label (RUS^(r) 0.37), and Gradient Ascent (RUS^(r) 0.29). The Retrained Model's RUS^(r) is 1.00 by construction. Representation-level metrics are omitted for POUR-P because it does not modify the encoder representations.
• CIFAR-100 (Table 2): POUR-P reaches Acc_r 77.65 and Acc_f 0.00 with AUS 1.00, versus the Retrained Model's Acc_r 76.28. POUR-D reaches Acc_r 73.44, Acc_f 1.00, AUS 0.95, CKA_f^(r) 0.57, CKA_r^(r) 0.76, RUS^(r) 0.65, and rMIA 46.00. Baselines are weaker: Gradient Ascent drops to Acc_r 50.46, Boundary Expand 66.47, DELETE 64.67.
• Higher entanglement makes forgetting harder. The authors report that CIFAR-100 classes are more entangled, reflected in a high CKA_f^(r), meaning supervision available from the forget set is lower — consistent with the class-separation term in their Proposition 2.2 bound.
• PathMNIST with a pretrained ViT-S/16 (Table 3): POUR-P achieves Acc_r 87.14 / Acc_f 0.00 / AUS 1.00 on the internal test set and Acc_r 87.44 / Acc_f 0.00 / AUS 1.00 on the external test set. POUR-D achieves 81.09 / 7.88 / AUS 0.87 (internal) and 80.90 / 7.92 / AUS 0.87 (external), with RUS^(o) 0.63 and 0.61 respectively. Finetune scores much lower representation-level RUS^(o) (0.05 internal, 0.06 external) despite high classification accuracy, because it accesses the retain set.
• Domain shift behaviour. Random Label and Gradient Ascent score higher on the internal test set than the external one (for example, Gradient Ascent Acc_r 81.43 internal versus 76.72 external), which the authors interpret as "learning to mask" the forget set rather than erasing the knowledge. DELETE and POUR are reported as consistent across both domains.
• Boundary methods fail on PathMNIST. The paper states that Boundary Shrink and Boundary Expand did not work in this setting.
• Cross-modal unlearning with CLIP-L/14 on ImageNet (Table 4): Removing text embeddings reduces the associated class's accuracy while leaving others largely unchanged. Goldfish: Acc_f falls from 94.00 to 14.00 (−80.00) while Acc_r rises from 67.36 to 69.23 (+1.87). Bison: Acc_f 100.00 to 6.00 (−94.00), Acc_r 67.33 to 67.76 (+0.43). European fire salamander: Acc_f 94.00 to 22.00 (−72.00), Acc_r 67.36 to 67.25 (−0.11). Boa constrictor: Acc_f 60.00 to 6.00 (−54.00), Acc_r 67.53 to 70.33 (+2.80). Centipede: Acc_f 42.00 to 10.00 (−32.00), Acc_r 67.62 to 67.74 (+0.12).
• Segmentation unlearning on VOC2012 (Table 5, DeepLabV3+ with ResNet-101): Forgetting Dog gives IoU_f 90.98 to 0.00 (−90.98) with IoU_r 77.81 to 71.53 (−6.28). Forgetting Cat gives IoU_f 93.55 to 0.00 (−93.55) with IoU_r 77.53 to 70.10 (−7.43). A Bicycle row begins with IoU_r 78.96 before unlearning; the remaining values are not present in the provided content.
• Neural Collapse holds empirically. Figure 5 shows classifier weight angle distributions whose empirical mean angles align almost perfectly with the ideal simplex ETF angle on CIFAR-10, CIFAR-100, and PathMNIST.
• Grad-CAM evidence. On PathMNIST, after applying POUR to the adipose class, its Grad-CAM signal vanishes while retained classes (debris, lymphocytes, mucus) keep clear attention patterns.

Methodology in Plain English

The authors start from the observation that a well-trained image classifier organizes its internal features so that each class sits around its own direction in feature space, and those directions are arranged as symmetrically as possible — a simplex Equiangular Tight Frame. Each class is therefore a direction you can point at.

Forgetting one class then becomes a geometric operation: subtract the component of every feature vector that lies along the forgotten class's direction. This is an orthogonal projection. The paper proves two things that make this sensible. First, when you delete one vertex from a regular simplex and project the rest onto the remaining subspace, you get a smaller regular simplex — so the retained classes keep their optimal angular separation. Second, this projected geometry is the same kind of structure that makes a nearest-class-mean classifier Bayes-optimal, so the projected model remains optimal on the retained classes. Because a forgotten class's features map to the origin, its predictions become uniform over the retained classes, which is exactly what the retrained reference model would produce.

This gives POUR-P: a one-shot, closed-form projection using the classifier weight of the forgotten class (or, if weights are unavailable, the empirical mean of the forgotten class's penultimate features). It requires no gradient updates.

Because POUR-P only acts after the fact, the authors add POUR-D. Here the projected model acts as a teacher and a student feature extractor is trained with an L2 loss to match the projected teacher features on the forget set only. They prove that driving this L2 loss to zero drives CKA similarity to one, so the student's representation converges toward the projected geometry.

Evaluation proceeds under two constraints: the retain set is unavailable during unlearning, and the original training procedure is not modified. Metrics include retain and forget accuracy (Acc_r, Acc_f and training-set variants Acc_tr, Acc_tf), the Adaptive Unlearning Score AUS = (1 − drop_r)/(1 + acc_f), a representation-level membership-inference attack (rMIA) using a five-fold linear regressor, and CKA-based comparisons against both the original model (superscript (o)) and the retrained model (superscript (r)), combined into RUS as the harmonic mean of the retention CKA and the forgetting indicator.

Why This Matters

• Research impact: The paper reframes unlearning from logit matching to geometry manipulation and supplies a provable optimality argument rather than a heuristic. It also contributes two general properties of Neural Collapse — the ETF-as-Bayes-optimality certificate and ETF invariance under orthogonal projection — which are of interest beyond unlearning. It offers a concrete metric (RUS) for a problem that previously lacked a representation-level score, and shows prior methods can look successful on classification metrics while scoring poorly on representation-level ones (for example, Finetune on PathMNIST with RUS^(o) 0.05 internal).
• Real-world applications:
  • Regulatory compliance with data protection laws and the "right to be forgotten," where a deployed model must remove a person's or class's data.
  • Medical imaging: the demonstrated PathMNIST case removes the adipose tissue class while keeping debris, lymphocytes, and mucus recognizable, relevant for privacy-sensitive clinical pipelines.
  • Removing spurious correlations or biased visual concepts from pre-trained models without retraining.
  • Removing concepts from vision-language models by deleting text embeddings, as demonstrated with CLIP-L/14.
• Industry relevance: POUR-P needs no retain set and no retraining, and its projection can be computed from an empirical feature mean, which the authors note makes it applicable when only the encoder is available, as with vision-language models. The distillation variant requires only the forget set, which

Authors’ abstract

In computer vision, machine unlearning aims to remove the influence of specific visual concepts or training images without retraining from scratch. Studies show that existing approaches often modify the classifier while leaving internal representations intact, resulting in incomplete forgetting. In this work, we extend the notion of unlearning to the representation level, deriving a three-term interplay between forgetting efficacy, retention fidelity, and class separation. Building on Neural Collapse theory, we show that the orthogonal projection of a simplex Equiangular Tight Frame (ETF) remains an ETF in a lower dimensional space, yielding a provably optimal forgetting operator. We further introduce the Representation Unlearning Score (RUS) to quantify representation-level forgetting and retention fidelity. Building on this, we introduce POUR (Provably Optimal Unlearning of Representations), a geometric projection method with closed-form (POUR-P) and a feature-level unlearning variant under a distillation scheme (POUR-D). Experiments on CIFAR-10/100 and PathMNIST demonstrate that POUR achieves effective unlearning while preserving retained knowledge, outperforming state-of-the-art unlearning methods on both classification-level and representation-level metrics.