Forget-It-All: Multi-Concept Machine Unlearning via Concept-Aware Neuron Masking
Abstract
The widespread adoption of text-to-image (T2I) diffusion models has raised concerns about their potential to generate copyrighted, inappropriate, or sensitive imagery. As a practical solution, machine unlearning aims to erase unwanted concepts without retraining from scratch. While most existing methods are effective for single-concept unlearning, they often struggle when removing multiple concepts, causing significant challenges in unlearning effectiveness, generation quality, and sensitivity to hyperparameters and datasets. We take a unique perspective on multi-concept unlearning by leveraging model sparsity and propose the Forget It All (FIA) framework. FIA first introduces Contrastive Concept Saliency to quantify each weight connection's contribution to a target concept. It then identifies Concept Sensitive Neurons by combining temporal and spatial information, ensuring that only neurons consistently responsive to the target concept are selected. Finally, FIA constructs masks from the identified neurons and fuses them into a unified multi-concept mask, where Concept Agnostic Neurons that broadly support general content generation are preserved while concept-specific neurons are pruned to remove the targets. FIA is training-free and requires minimal hyperparameter tuning for new tasks, enabling plug-and-play use. Extensive experiments across three distinct unlearning tasks demonstrate that FIA achieves more reliable multi-concept unlearning, improving forgetting effectiveness while maintaining generation fidelity and quality. Code is available at https://github.com/kaiyuan02415/Forget-It-All
Lay Summary
AI image generators can create realistic pictures from text descriptions, but they sometimes produce images that raise concerns, such as copyrighted artwork, inappropriate content, or sensitive material, because they learned from vast amounts of internet data. One way to address this is to make the model "forget" specific unwanted content without retraining it from scratch, a process known as machine unlearning. However, current approaches work well for removing one type of content at a time but struggle when asked to forget multiple types simultaneously. They often degrade the quality of the images the model can still produce. We propose a framework called Forget It All (FIA) that takes a different approach inspired by how neural networks store information. FIA identifies the specific internal components responsible for generating unwanted content and removes them, while carefully preserving the components needed for normal image generation. Because it works by analyzing the structure of the model rather than retraining it, FIA requires no additional training and minimal manual tuning. Our experiments show that FIA can reliably remove multiple types of unwanted content at once while maintaining the ability to generate high quality images for legitimate uses.