Unleashing Uncertainty: Efficient Machine Unlearning for Generative AI

TL;DR

This paper presents SAFEMax, an innovative method for machine unlearning in diffusion models that enhances privacy by increasing entropy in generated images, selectively forgetting specific classes efficiently.

Contribution

SAFEMax introduces an information-theoretic approach to unlearning in diffusion models, focusing on early diffusion steps for improved efficiency and control.

Findings

SAFEMax effectively increases entropy to forget specific classes.

The method achieves substantial efficiency gains over existing approaches.

SAFEMax successfully halts denoising for impermissible classes.

Abstract

We introduce SAFEMax, a novel method for Machine Unlearning in diffusion models. Grounded in information-theoretic principles, SAFEMax maximizes the entropy in generated images, causing the model to generate Gaussian noise when conditioned on impermissible classes by ultimately halting its denoising process. Also, our method controls the balance between forgetting and retention by selectively focusing on the early diffusion steps, where class-specific information is prominent. Our results demonstrate the effectiveness of SAFEMax and highlight its substantial efficiency gains over state-of-the-art methods.

Equations

$$q(x_t\mid x_0)=\mathcal{N}(x_t;\sqrt{{\bar{\alpha }}_t}{x}_0,(1-{\bar{\alpha }}_t)\mathbf{I}),{\bar{\alpha }}_t=\prod _{i=1}^t{\alpha }_i$$

$$\mathcal{L}={\mathbb{E}}_{t\in [1,T],ϵ\sim \mathcal{N}(0,1)}[\mid \mid {ϵ}_t-{ϵ}_{\theta }(x_t,c,t)\mid {\mid }_2^2]$$

$${\mathcal{L}}_f={\mathbb{E}}_{t\in [1,T],ϵ\sim \mathcal{N}(0,1)}[\psi (t)\mid \mid {ϵ}_T-{ϵ}_{\theta }(x_t,c_f,t)\mid {\mid }_2^2]$$

$$P_e\ge \frac{H(x\mid \hat{x})-1}{log\mid \mathcal{X}\mid }$$

$$\psi (t)=\exp (-\lambda \frac{t}{T}),\text{for}t\in [0,T]$$

Full text

Machine Learning, ICML, Machine Unlearning

1 Introduction

Machine Unlearning (MU) for generative AI aims to prevent the generation of impermissible content, such as samples from a specific class—referred to as the forget samples or forget class. The goal of MU is to correct pre-trained models efficiently, without retraining from scratch, thereby minimizing computational overhead. Efficiency is crucial to reducing both the cost and latency of model correction.

Most state-of-the-art methods in generative MU rely heavily on unlearning strategies originally developed for discriminative tasks, either by directly adapting existing methods to generative settings or by incorporating techniques, such as the Fisher Information Matrix (FIM) (Golatkar et al., 2020; Foster et al., 2024) and weight masking (Jia et al., 2023), originally introduced for discriminative MU.

For example, Selective Amnesia (Heng & Soh, 2023) uses FIM to guide unlearning through elastic weight consolidation in generative models. Saliency Unlearning (Fan et al., 2024) applies weight masks on top of discriminative MU methods—Random Labeling (Graves et al., 2021) and Gradient Ascent (Thudi et al., 2022)—which have been extended in generative settings. Other works (Wu & Harandi, 2024; wu2024scissorhands; Ko et al., 2024; Patel & Qiu, 2025) adapt Gradient Ascent or NegGrad+ (Kurmanji et al., 2023) to generative models and apply multi-objective optimization techniques.

Despite progress in bridging discriminative and generative MU, insights from state-of-the-art approaches in discriminative tasks remain underexplored. Among recent discriminative MU methods, LoTUS (Spartalis et al., 2025) stands out for its strong unlearning performance and computational efficiency. It is an entropy-based method that increases the model’s uncertainty on forget samples by smoothing the corresponding prediction probabilities up to an information-theoretic bound, thereby controlling the entropy increase. Motivated by this approach, we investigate the following research questions in the context of generative MU:

1. Can we efficiently train a generative model to forget specific samples by increasing the entropy on those samples?

2. Can we control this process to better balance the trade-off between unlearning and retention of useful knowledge?

To this end, we introduce a Simple And Fast Entropy Maximization method for efficient MU in diffusion models. SAFEMax, for short, maximizes the entropy on the forget samples by training the model to generate Gaussian noise when conditioned on the forget class, ultimately halting the denoising process. Also, it leverages the progressive loss of information in the diffusion process, as shown in Figure 1, to control unlearning and balance forgetting and retention. Specifically, it applies stronger unlearning in early diffusion steps, where class-specific details are formed and refined, and weaker unlearning in the later steps, where latent states across different classes converge due to the dominant influence of Gaussian noise. SAFEMax significantly improves efficiency compared to existing unlearning methods while preserving state-of-the-art performance, establishing a scalable and cost-effective solution for MU in generative AI.

2 Simple and Fast Entropy Maximization for Machine Unlearning in Generative AI

The training of Denoising Diffusion Probabilistic Models (DDPMs) (Ho et al., 2020) consists of two processes. First, during the forward diffusion process, the original image $x_{0}\!\sim\!q(x)$ is progressively corrupted with Gaussian noise over $T$ steps, such that $x_{T}\!\sim\!\mathcal{N}(0,\mathbf{I})$. The noise level at each step is determined by a schedule $\alpha_{t}$, and the latent state $x_{t}$ remains tractable and can be sampled directly:

[TABLE]

In the early diffusion steps, the Gaussian noise has lower variance (indicated by $\lim_{t\to 0}\bar{\alpha}_{t}\!=\!1$), making the latent states $x_{t}$ more similar to the original image $x_{0}$. These early states contain rich semantic information, which we refer to as class-specific information. As diffusion progresses, the latent states become more entropic (i.e., noisy) because of the dominance of cumulative Gaussian noise. In later diffusion steps, the variance of the injected Gaussian noise approximates its maximum (indicated by $\lim_{t\to T}\bar{\alpha}_{t}\!=\!0$) and the distribution $q(x_{t})$ becomes increasingly broad. Consequently, the latent states gradually lose the specific structure of the data and approach a state of maximum entropy. Specifically, for large $T$, the final latent state of the diffusion process, $x_{T}$, approximates the mean of the training data distribution for all inputs $x_{0}$ (Zhong et al., 2024).

Then, in the denoising process, a denoiser network $\epsilon_{\theta}$ is trained to predict the noise $\epsilon_{t}$ at any arbitrary step t:

[TABLE]

where $c$ denotes the class of the original image $x_{0}$.

Unlearning with SAFEMax leverages the inherent Gaussian noise of the diffusion process, particularly the noise $\epsilon_{T}$, which corresponds to the maximum entropy latent state $x_{T}$. To prevent the denoiser from reconstructing samples of an impermissible class $c_{f}$, we effectively halt the denoising process and maximize entropy in the model’s output by fine-tuning the model with the following forget loss:

[TABLE]

where $\epsilon_{T}$ is the cumulative Gaussian noise in the final latent state, $c_{f}$ is the forget class, and $\psi(t)=\exp(-t/T)$ is a monotonically decaying function within the range $[0,1]$ that emphasizes unlearning in the early diffusion steps, where class-specific information is most prominent. This design encourages selective model updates, aiming to better balance unlearning and retention of useful prior knowledge. For the remaining classes, we perform a regular fine-tuning.

Information-Theoretic Analysis of MU via Entropy Maximization.

Consider a training sample $x_{0}$ and its reconstruction $\hat{x}$ generated by a diffusion model $g(x)$. This process can be viewed as a Markov chain: $x\mapsto g(x)\mapsto\hat{x}$. Let $P_{e}=\text{Pr}\{x\neq\hat{x}\}$ denote the probability of reconstruction error. Within the context of diffusion models, we can define that the equality $x=\hat{x}$ applies when the reconstructed image $\hat{x}$ contains the same semantic information as the original image $x$ (e.g., belonging to the same class). According to Fano’s inequality (thomas2012elements), we have:

[TABLE]

where $\mid\mathcal{X}\mid$ is the cardinality of the input space. Equation 4 indicates that as the reconstructed image $\hat{x}$ becomes less informative about the original image $x$ (i.e., as the conditional entropy $H(x\mid\hat{x})$ increases), the lower bound of the error probability $P_{e}$ increases, implying that the generated image $\hat{x}$ does not share the same semantic information as the original input $x$.

To prevent generating images from the forget class, we train the model to generate images similar to the ideal latent state $x_{T}$, which contains pure Gaussian noise (as $T\to\infty$ in the diffusion process) and thus no semantic information from $x$ (denoted as $x_{0}$ in the diffusion process). Leveraging the entropy increase that is inherent in the diffusion trajectory, this design maximizes $H(x\mid\hat{x})$, thereby increasing the lower bound of reconstruction error $P_{e}$ and guiding unlearning.

Balancing forgetting and retention with $\psi(t)$.

Inspired by delineating information for targeted MU in discriminative tasks (Spartalis et al., 2025), and loss scheduling for transfer learning in diffusion models (Zhong et al., 2024), we introduce a loss scheduling strategy for targeted MU in generative tasks. Motivated by the observation that class-specific information is most prominent in the early diffusion steps and becomes increasingly obscured by Gaussian noise as the diffusion evolves, as shown in Figure 1, we propose a scheduling function that progressively reduces the influence of the forget loss throughout the diffusion process, thereby targeting unlearning in the semantically richer early stages:

[TABLE]

where $\lambda$ is a hyperparameter to controls the rate of decay. A larger $\lambda$ causes $\psi(t)$ to decay rapidly, concentrating the unlearning effect more narrowly on the early diffusion steps.

3 Experiments & Discussion

We follow the experimental setup and evaluation framework of Selective Amnesia (SA) and Saliency Unlearning (SalUn). We aim to forget a specific class within a DDPM trained on CIFAR-10 (krizhevsky2009learning). We apply SalUn on top of the Random Labeling strategy using saliency masks that update 50% of the weights, as proposed by the authors for DDPMs. For SA, we follow the author’s configuration and perform unlearning for 20,000 iterations. For SAFEMax, we adopt the same hyperparameter values as SalUn and perform unlearning for 1,000 iterations, consistent with SalUn.

To evaluate unlearning, we use a ResNet34 classifier, pre-trained on ImageNet (deng2009imagenet) and fine-tuned on CIFAR-10 for 20 epochs. We report two key metrics: (i) Unlearning Accuracy (UA), defined as 100% minus the accuracy of the classifier on the forget class, and (ii) the average entropy ($H$) of the classifier’s prediction for the forget class, which captures the uncertainty introduced by the unlearning method. To evaluate retention (i.e., the model’s performance on the remaining classes), we compute the Fréchet Inception Distance on generated images for retained classes. Finally, we measure efficiency by reporting each method’s Runtime Estimation (RTE) and peak GPU memory usage.

Effectiveness.

SAFEMax demonstrates strong unlearning performance by excelling in both forgetting and retention of the respective information. Overall, our method achieves better unlearning accuracy (UA) and retention quality (FID) than state-of-the-art approaches, as shown in Tables 1 and 2.

In terms of unlearning, SAFEMax consistently achieves the maximum UA score (100%) in all-but-one case. The only exception is Class 2 (birds), where the UA drops to 0%, indicating that the classifier predicts all generated images for the unlearned class as birds. However, this anomaly is not due a failure of our method. Instead, the classifier mistakenly identifies Gaussian noise as birds, as evidenced in Figure 2. This observation highlights that UA scores can be misleading in isolation and underscores the importance of examining both quantitative and qualitative results.

Notably, SAFEMax consistently achieves the highest increase in entropy, directly aligning with its design goal of maximizing uncertainty for the forget class. As shown in Figure 2, it generates high-entropy noise more reliably than SA, leading to greater classifier uncertainty, as reflected in the entropy scores ($H$) in Table 1. In contrast, SalUn undermines the entropy-increase objective of Random Labeling in discriminative tasks. While Random Labeling originally assigned random incorrect labels to forget samples at each unlearning iteration, its adaptation in SalUn maps the forget class to a fixed alternative class (as shown in Figure 2). This leads the evaluation model to make incorrect yet high-confidence predictions, as shown in Table 1.

In terms of retention—preserving high image quality for the remaining classes as measured by FID—SAFEMax generally achieves state-of-the-art results as shown in Table 2. Overall, SAFEMax delivers effective forgetting and retention while requiring significantly less time and memory.

Efficiency.

As shown in Table 2, a key advantage of SAFEMax is achieving a strong balance between forgetting and retention without relying on computationally intensive techniques (e.g., the use of FIM, weight masks, or regularization terms for multi-objective optimization) that introduce significant overhead before or during unlearning. Compared to SalUn, which runs for the same number of iterations, SAFEMax is $2\times$ faster. Against SA, our method achieves a $30\times$ speed-up due to its ability to unlearn effectively in far fewer iterations. When also accounting for the time to compute the FIM for SA (1226.98 minutes, not included in the RTE metric), SAFEMax offers a total $230\times$ improvement in runtime.

In terms of GPU memory usage, SAFEMax is 59% more efficient than SalUn and 45% more efficient than SA. The high memory demands of SalUn and SA result from their reliance on storing the saliency masks and FIM, respectively. In contrast, SAFEMax does not require complex auxiliary structures to balance forgetting and retention, and instead uses a simple, predefined scheduler. This simple yet effective design makes SAFEMax not only faster but also more practical for large-scale or resource-constrained applications.

Ablation Study.

To assess the role of our scheduling function $\psi(t)$, we conducted an ablation study, presented in Figure 3. The results show that SAFEMax achieves strong unlearning performance even without the scheduler, underscoring the robustness of our core method. However, introducing the decaying scheduler $\psi(t)$ improves the trade-off between unlearning and retention, as evidenced by improved image quality in the generated outputs. We further analyze the impact of the decay parameter $\lambda$ and verify that a faster decay (i.e., larger $\lambda$) improves retention—even for the forget class. In our experiments, we used a moderate value of $\lambda\!=\!1$, without additional tuning. Tuning $\lambda$ may yield better results.

4 Conclusion

In this paper, we introduced SAFEMax, an effective and significantly more efficient Machine Unlearning strategy for diffusion models. Motivated by the information-theoretic analysis of LoTUS (Spartalis et al., 2025) in discriminative tasks, we maximized the entropy of the forget samples by leveraging the natural entropy increase inherent to diffusion models. We further proposed a simple scheduling mechanism that targets class-specific information, enhancing the balance between unlearning and the retention of useful knowledge. We compared SAFEMax against the most prominent state-of-the-art approaches for DDPM unlearning and showed that our method not only achieves strong unlearning performance but also offers substantial improvements in computational efficiency. These findings suggest that SAFEMax is a promising, scalable, and cost-effective unlearning strategy.

Limitations & Future Work.

A more comprehensive evaluation could further strengthen the validity and applicability of our approach, even though SAFEMax has already demonstrated strong performance across key benchmarks. We plan to compare against more recent DDPM unlearning methods, and to extend our evaluation to additional datasets and Stable Diffusion models.

5 Acknowledgment

This work was partially supported by the EU funded project ATLANTIS (Grant Agreement Number 101073909).

Appendix A Broader Social Impact

Methods like SAFEMax align machine learning models with privacy regulations and ethical standards, by particularly preventing the generation of impermissible or sensitive content. However, they can also be misused by adversaries to deliberately degrade the performance of otherwise well-functioning models, suggesting careful consideration of deployment practices.

Appendix B Study Note on the Noise Incorporated in SAFEMax and Selective Amnesia

In Selective Amnesia, the per-pixel noise is drawn from a uniform distribution $\mathcal{U}(-1,1)$, which has a differential entropy of $H=\log(2)\approx 0.6931$ per dimension (i.e., channel). In contrast, in SAFEMax, the per-pixel noise is sampled from a Gaussian distribution $\mathcal{N}(0,1)$, which has a higher differential entropy of $H=\frac{1}{2}\log(2\pi e)\approx 1.4189$ per dimension.

Appendix C More Visualizations