Accurate Forgetting for All-in-One Image Restoration Model

TL;DR

This work tackles privacy protection in image restoration by proposing a model-agnostic forgetting framework that targets and eliminates the influence of a private degradation type on an All-In-One restoration model. It introduces instance-wise unlearning with a forgetting loss $\mathrm{L}_{\mathrm{UL}}$ and an adversarial regularization term $\mathrm{L}_{\mathrm{UL}}^{\mathrm{Adv}}$, enabling accurate forgetting without full retraining, and demonstrates robustness across degradation types and hyperparameters. Evaluations on haze, rain, and noise tasks using base models like PromptIR and AdaIR show effective forgetting of the specified degradation while preserving or even improving performance on retained tasks. The approach opens a new privacy-preservation avenue for low-level vision, offering a practical, low-cost alternative to retraining while highlighting directions for scalability and single-class restoration future work.

Abstract

Privacy protection has always been an ongoing topic, especially for AI. Currently, a low-cost scheme called Machine Unlearning forgets the private data remembered in the model. Specifically, given a private dataset and a trained neural network, we need to use e.g. pruning, fine-tuning, and gradient ascent to remove the influence of the private dataset on the neural network. Inspired by this, we try to use this concept to bridge the gap between the fields of image restoration and security, creating a new research idea. We propose the scene for the All-In-One model (a neural network that restores a wide range of degraded information), where a given dataset such as haze, or rain, is private and needs to be eliminated from the influence of it on the trained model. Notably, we find great challenges in this task to remove the influence of sensitive data while ensuring that the overall model performance remains robust, which is akin to directing a symphony orchestra without specific instruments while keeping the playing soothing. Here we explore a simple but effective approach: Instance-wise Unlearning through the use of adversarial examples and gradient ascent techniques. Our approach is a low-cost solution compared to the strategy of retraining the model from scratch, where the gradient ascent trick forgets the specified data and the performance of the adversarial sample maintenance model is robust. Through extensive experimentation on two popular unified image restoration models, we show that our approach effectively preserves knowledge of remaining data while unlearning a given degradation type.

Figures (3)

• Figure 1: Our method. We start by giving a dataset that needs to be privacy-preserving, other additional degradation information, and an all-in-one neural network that has been pre-trained. Then, we impose an elevated L1 loss on the model's outputs for the targeted degradation type along with their corresponding clean images, thereby encouraging the model to forget the specific degradation pattern during the training process (we employ a gradient-up approach).On the other hand, we introduce adversarial examples and a small number of datasets from other tasks to ensure robust output of the model.
• Figure 2: Eliminating the dehazing capability from the unified model, the visual outcomes illustrate that our approach can significantly eradicate the impact of deleting data, in contrast to a model that has not been trained on such datasets.
• Figure 3: The effects of learning rate and batch size on the model's unlearning and restoration performance across various image processing tasks.