Unmerge: Efficient Machine Unlearning via Task Arithmetic

2026-09-29

AI Summary

This study addresses key challenges in machine unlearning, including the hyperparameter sensitivity of gradient-based methods, knowledge entanglement, and limited interpretability, by proposing an efficient model editing framework based on task arithmetic. The core idea formulates unlearning as subtracting a low-rank forgetting vector, achieving precise erasure through activation-space constrained optimization. Furthermore, this work introduces low-rank forgetting basis decomposition and inter-layer geometric diagnostics, endowing the unlearning process with structured interpretability and rigorously establishing its theoretical boundaries. Extensive experiments demonstrate that the proposed method significantly outperforms existing baselines across ResNet, ViT, and Llama architectures, approaching full retraining performance while improving computational efficiency several-fold.

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.

Problem

• Machine Unlearning
• Task Arithmetic
• Knowledge Entanglement
• Hyperparameter Search
• Model Interpretability

Innovation

• Machine Unlearning
• Task Arithmetic
• Low-Rank Factorization
• Activation Subspace
• Unmerge