ReTrack: Data Unlearning in Diffusion Models through Redirecting the Denoising Trajectory

Abstract

Diffusion models excel at generating high-quality, diverse images but suffer from training data memorization, raising critical privacy and safety concerns. Data unlearning has emerged to mitigate this issue by removing the influence of specific data without retraining from scratch. We propose ReTrack, a fast and effective data unlearning method for diffusion models. ReTrack employs importance sampling to construct a more efficient fine-tuning loss, which we approximate by retaining only dominant terms. This yields an interpretable objective that redirects denoising trajectories toward the $k$-nearest neighbors, enabling efficient unlearning while preserving generative quality. Experiments on MNIST T-Shirt, CelebA-HQ, CIFAR-10, and Stable Diffusion show that ReTrack achieves state-of-the-art performance, striking the best trade-off between unlearning strength and generation quality preservation.

Figures (12)

• Figure 1: An intuitive schematic of our method. By changing the diffusion model's noise prediction on $\boldsymbol{x}_t$, the denoising trajectory that originally pointed to the unlearning sample $\boldsymbol{a}_u$ is redirected to its $k$-nearest neighbors, thus achieving fast unlearning while preserving generation quality.
• Figure 2: Construction of the Stable Diffusion dataset. For each text-image pair, in order to forget the corresponding image in this text condition, the pretrained model is used to sample images in the given text condition and clustered to construct the unlearning set and remaining set required by the data unlearning method. After unlearning, the model no longer generates the corresponding image in the same text condition.
• Figure 3: Schematic of the SSCD metric calculation method on CelebA-HQ and CIFAR-10. Instead of adding noise up to the full $T$ steps and generating image from the pure Gaussian noise, we inject noise for only $t<T$ steps and use the diffusion model to denoise and reconstruct the clean image. The SSCD is then calculated to measure the similarity between the original training image and the reconstructed image, thereby indicating how much of the training image's information is retained by the model.
• Figure 4: Results on Stable Diffusion.
• Figure 5: Partial sampling results from pretrained and unlearned models on the MNIST T-Shirt dataset.

Theorems & Definitions (6)

• Proposition 1
• Proposition 2
• Proposition 2
• proof
• Proposition 2
• proof