Critic-Guided Reinforcement Unlearning in Text-to-Image Diffusion

TL;DR

This work addresses the challenge of selectively unlearning targeted concepts in text-to-image diffusion models without sacrificing utility. It introduces Critic-Guided Reinforcement Unlearning (CGRU), treating the reverse diffusion process as an RL policy and equipping it with a per-timestep critic that estimates the expected terminal reward, enabling dense advantage signals via $A(s_t,a_t)=r(x_0,c)-V_\phi(x_t,c,t)$. The method leverages off-policy importance sampling and a CLIP-based reward predictor to localize forgetting within the denoising trajectory and stabilize optimization. Empirically, CGRU achieves strong unlearning performance on the UnlearnCanvas benchmark (e.g., UA=95.55%) while maintaining competitive retain quality (IRA) and image fidelity (FID), outperforming terminal-reward baselines such as DDPO. The results demonstrate the practical viability of RL-based diffusion unlearning and highlight the value of timestep-aware critics and noisy-conditioned rewards for stable, credit-assigned forgetting.

Abstract

Machine unlearning in text-to-image diffusion models aims to remove targeted concepts while preserving overall utility. Prior diffusion unlearning methods typically rely on supervised weight edits or global penalties; reinforcement-learning (RL) approaches, while flexible, often optimize sparse end-of-trajectory rewards, yielding high-variance updates and weak credit assignment. We present a general RL framework for diffusion unlearning that treats denoising as a sequential decision process and introduces a timestep-aware critic with noisy-step rewards. Concretely, we train a CLIP-based reward predictor on noisy latents and use its per-step signal to compute advantage estimates for policy-gradient updates of the reverse diffusion kernel. Our algorithm is simple to implement, supports off-policy reuse, and plugs into standard text-to-image backbones. Across multiple concepts, the method achieves better or comparable forgetting to strong baselines while maintaining image quality and benign prompt fidelity; ablations show that (i) per-step critics and (ii) noisy-conditioned rewards are key to stability and effectiveness. We release code and evaluation scripts to facilitate reproducibility and future research on RL-based diffusion unlearning.

Figures (9)

• Figure 1: Mean Aesthetic Score reward during training for CGRU and DDPO. CGRU shows superior performance with faster convergence and higher final rewards.
• Figure 2: Concept suppression comparison between CGRU and DDPO on the "Cats" class. CGRU demonstrates superior performance in reducing unwanted object generation.
• Figure 3: Progression of concept forgetting during training for CGRU and DDPO. Each row shows generated images from different training checkpoints (default, early, mid, late), demonstrating how each method handles concept suppression while preserving overall image quality.
• Figure 4: Ablation study comparing the performance of a standard critic vs. our timestep-aware critic on noisy latents. The timestep-aware model demonstrates significantly higher accuracy and macro precision, validating the importance of temporal conditioning for value estimation.
• Figure 5: Trade-off between Unlearning Accuracy (UA) and In-domain Retain Accuracy (IRA). The green dashed line indicates the trend of existing baselines. CGRU (red star) lies above this trend, demonstrating a superior trade-off point favoring effective erasure while maintaining competitive retention.