Machine unlearning improved by subtracting forget signals efficiently
Unmerge: Efficient Machine Unlearning via Task Arithmetic
Machine LearningArtificial Intelligence
Summary
Sometimes we want a trained AI model to "forget" certain data without retraining it from scratch, which is slow and costly. This paper introduces a new method called Unmerge that views unlearning as subtracting a learned 'forget' part from the combined knowledge the model has. By focusing on important directions where the forget data influences the model, Unmerge can remove that influence quickly and accurately while keeping useful knowledge intact. This approach works well on popular image and language models and gives insights into where models struggle to unlearn.
What this means in practice
- •For ml engineers: Remove specific data influences from trained image and language models efficiently without costly retraining.
- •For privacy compliance teams: Ensure that sensitive or outdated user data can be unlearned from deployed AI models with membership inference risk reduced.
Authors
Haoran Tang, Andrew Tan, Rajiv Khanna
Abstract
Approximate machine unlearning seeks to remove the influence of a forget set from a trained model without full retraining. Existing gradient-based methods require data-dependent hyperparameter search, struggle when forget and retain knowledge are entangled, and offer little insight into where unlearning actually happens inside the network. We recast unlearning through the lens of task arithmetic: if finetuning produces a merged task vector $τ_m$ that combines learning on forget and retain sets, unlearning is the inverse operation that subtracts a learned forget component $τ_F$ to recover the retain task vector $τ_R$. The forget signal is concentrated: at every layer, forget activations lie in a subspace spanned by a handful of dominant directions, so we factorize $τ_F$ in a low-rank forget basis, which is faithful up to a small tail-eigenvalue residual and limits how far the correction can perturb retain. We then optimize three intuitive goals (match the merged vector inside the forget span, suppress leakage into the retain span, and bound the correction size) that provably bound forget leakage and retain damage in activation space. The resulting algorithm, Unmerge, is fast and powerful: on class-level unlearning with ResNet-50 on CIFAR-100 and Tiny ImageNet, it improves Tug-of-War by up to ~24% over a baseline of comparable runtime and by up to ~18% over stronger baselines that run ~5x slower, keeps membership-inference exposure at the level of retraining, and shrinks the feature-distribution gap to the retrained model, where relabeling methods leave forget features cleanly separable. Further studies show that Unmerge also applies to ViT-S/16 and scales to Llama-3.2-3B. The per-layer basis geometry that drives the algorithm also serves as a layerwise diagnostic for when and where unlearning becomes structurally hard.