Fast Unlearning at Scale via Margin Self-Correction

2026-06-01

AI Summary

This work addresses the challenge of efficiently unlearning specific data from language models without retraining. It proposes Margin-Aware Self-Correction (MASC), a method that actively narrows the confidence gap between original and substitute tokens in to-be-forgotten texts by dynamically adjusting logit margins during inference. MASC introduces, for the first time, an adaptive stopping mechanism that operates without downstream validation, thereby eliminating redundant computation and the need for multiple checkpoint evaluations. Experiments on TOFU, MUSE News, and MUSE Books demonstrate that MASC achieves forgetting performance comparable to existing approaches while significantly reducing computational and storage overhead. Moreover, it exhibits superior retention of model utility and more stable forgetting efficacy, particularly on larger-scale models.

Abstract

Language-model unlearning updates a trained model to behave as if it had not seen selected training examples, while preserving utility and avoiding costly retraining. Existing approaches typically fine-tune the pretrained model with a fixed training budget and select the final model afterwards by evaluating several saved checkpoints on downstream validation data. Two sources of unnecessary computation limit scalability: training beyond the desired forget-retain trade-off, and checkpoint selection that requires extra storage and repeated evaluations. To address these limitations, we introduce MArgin Self-Correction (MASC), an efficient unlearning method with an online stopping rule that does not require downstream evaluation. Given a text sequence to be forgotten, MASC actively reduces the logit gap between the original next token and the most likely alternatives. It outputs a final model once this gap is small on average over a sufficiently large proportion of token positions across all forget sequences. On TOFU, MUSE News, and MUSE Books, MASC achieves a competitive forget-retain trade-off at a fraction of the computational cost of existing baselines. We further observe that as we increase model size (a.k.a. number of parameters), the trade-offs improve for both MASC and SimNPO -- the forget metrics remain comparable while retain utility increases.

Problem

• unlearning
• language models
• forget-retain trade-off
• scalability
• efficient computation

Innovation

• unlearning
• margin self-correction
• online stopping rule
• logit gap
• forget-retain trade-off