English

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

Machine Learning 2026-02-17 v3

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.

Keywords

Cite

@article{arxiv.2601.03213,
  title  = {Critic-Guided Reinforcement Unlearning in Text-to-Image Diffusion},
  author = {Mykola Vysotskyi and Zahar Kohut and Mariia Shpir and Taras Rumezhak and Volodymyr Karpiv},
  journal= {arXiv preprint arXiv:2601.03213},
  year   = {2026}
}

Comments

Preprint

R2 v1 2026-07-01T08:52:58.185Z