Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

EraseFlow: Learning Concept Erasure Policies via GFlowNet-Driven Alignment

Abhiram Kusumba, Maitreya Patel, Kyle Min, Changhoon Kim, Chitta Baral, Yezhou Yang

TL;DR

EraseFlow reframes concept erasure in text-to-image diffusion models as a trajectory-level alignment problem solvable with Generative Flow Networks under a trajectory-balance objective. By employing reward-free alignment and leveraging entire denoising trajectories, it achieves robust erasure of concepts (e.g., NSFW content, logos, artistic styles) while preserving the model's prior and image quality, with markedly lower training cost. The method demonstrates strong robustness to red-teaming, maintains image-text alignment, and generalizes to recent architectures, while offering plug-and-play compatibility with other safety tools. Theoretical guarantees via the constant-reward TB proposition underpin stable distributional alignment, and extensive experiments show EraseFlow attaining state-of-the-art-like performance across diverse erasure tasks with high efficiency and practical impact for safety in diffusion-based generation.

Abstract

Erasing harmful or proprietary concepts from powerful text to image generators is an emerging safety requirement, yet current "concept erasure" techniques either collapse image quality, rely on brittle adversarial losses, or demand prohibitive retraining cycles. We trace these limitations to a myopic view of the denoising trajectories that govern diffusion based generation. We introduce EraseFlow, the first framework that casts concept unlearning as exploration in the space of denoising paths and optimizes it with GFlowNets equipped with the trajectory balance objective. By sampling entire trajectories rather than single end states, EraseFlow learns a stochastic policy that steers generation away from target concepts while preserving the model's prior. EraseFlow eliminates the need for carefully crafted reward models and by doing this, it generalizes effectively to unseen concepts and avoids hackable rewards while improving the performance. Extensive empirical results demonstrate that EraseFlow outperforms existing baselines and achieves an optimal trade off between performance and prior preservation.

EraseFlow: Learning Concept Erasure Policies via GFlowNet-Driven Alignment

TL;DR

EraseFlow reframes concept erasure in text-to-image diffusion models as a trajectory-level alignment problem solvable with Generative Flow Networks under a trajectory-balance objective. By employing reward-free alignment and leveraging entire denoising trajectories, it achieves robust erasure of concepts (e.g., NSFW content, logos, artistic styles) while preserving the model's prior and image quality, with markedly lower training cost. The method demonstrates strong robustness to red-teaming, maintains image-text alignment, and generalizes to recent architectures, while offering plug-and-play compatibility with other safety tools. Theoretical guarantees via the constant-reward TB proposition underpin stable distributional alignment, and extensive experiments show EraseFlow attaining state-of-the-art-like performance across diverse erasure tasks with high efficiency and practical impact for safety in diffusion-based generation.

Abstract

Erasing harmful or proprietary concepts from powerful text to image generators is an emerging safety requirement, yet current "concept erasure" techniques either collapse image quality, rely on brittle adversarial losses, or demand prohibitive retraining cycles. We trace these limitations to a myopic view of the denoising trajectories that govern diffusion based generation. We introduce EraseFlow, the first framework that casts concept unlearning as exploration in the space of denoising paths and optimizes it with GFlowNets equipped with the trajectory balance objective. By sampling entire trajectories rather than single end states, EraseFlow learns a stochastic policy that steers generation away from target concepts while preserving the model's prior. EraseFlow eliminates the need for carefully crafted reward models and by doing this, it generalizes effectively to unseen concepts and avoids hackable rewards while improving the performance. Extensive empirical results demonstrate that EraseFlow outperforms existing baselines and achieves an optimal trade off between performance and prior preservation.

Paper Structure

This paper contains 50 sections, 2 theorems, 25 equations, 16 figures, 13 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Works
  3. Concept Erasure for Text-to-Image Diffusion Models.
  4. Adversarial Robustness Text-to-Image Diffusion Models.
  5. Diffusion Alignment and GFlowNets.
  6. Preliminaries
  7. Concept Erasure for Diffusion Models
  8. Generative Flow Networks (GFlowNets)
  9. Proposed Methodology: EraseFlow
  10. Trajectory balance (TB).
  11. Reward‑free alignment.
  12. Proof sketch.
  13. Plug and Play.
  14. EraseFlow training algorithm
  15. Experimental Results
  16. ...and 35 more sections

Key Result

Proposition 1

Let the noising kernel $q(\cdot \mid \cdot)$ be fixed and non-degenerate. Assume there exist parameters $(\theta^*,\phi^*)$ such that the constant-reward loss equation eq:eraseflow_loss satisfies $\mathcal{L}_{c \leftarrow c^*}^\text{EraseFlow}=0$ and, for the original model with safe prompt $c^*$, and consequently the marginal image distributions coincide: Hence the visual concept unique to $c$

Figures (16)

  • Figure 1: EraseFlow (ours) achieves effective concept erasure when compared with various concept erasure methods across diverse tasks—removing NSFW content (top), artistic styles like “Van Gogh” (middle), and fine-grained elements such as the “Nike” logo from shoes (bottom)—while preserving image quality and fidelity.
  • Figure 2: Comparison of EraseFlow variants and baselines on erasing the nudity concept in Stable Diffusion v1.4. Robustness is measured by adversarial attack success rate (↓) and utility by FID (↓). Lower is better on both axes. Circle size reflects training cost.
  • Figure 3: Qualitative (Top) and Quantitative (Bottom) Comparison of DB Vs. TB vs EraseFlow(Ours) on an NSFW Prompt.
  • Figure 4: Illustration of probability redistribution during EraseFlow optimization. (a) In the pretrained model, probability mass is concentrated around target regions, increasing the likelihood of generating target concepts. (b) During training, the trajectory-balance objective redistributes probability mass from target regions toward anchor regions. (c) After optimization, the learned distribution aligns with anchor regions, effectively suppressing target concepts while maintaining visual and semantic fidelity.
  • Figure 5: Image generations by SDv1.4 and concept-erasure methods on different prompts. (top) A nudity prompt attacked by UDAtk; EraseFlow effectively suppresses NSFW content while most methods fail. (middle) A Van Gogh-style prompt attacked by UDAtk; EraseFlow removes the artistic style successfully. (bottom) A prompt with fine-grained elements; EraseFlow removes “wings” from “Pegasus” while preserving all other details like the horse and the mountain range.
  • ...and 11 more figures

Theorems & Definitions (3)

  • Proposition 1: Concept erasure via constant-reward TB
  • Proposition 2: Concept erasure via constant-reward TB
  • proof