EraseFlow: Learning Concept Erasure Policies via GFlowNet-Driven Alignment

AI summary

Overview

Research area: Concept erasure (unlearning) in text-to-image diffusion models, approached through generative flow networks (GFlowNets) and trajectory-balance alignment.

Technical level: Advanced. The paper assumes familiarity with diffusion model denoising, reinforcement-learning-style alignment, and flow-based generative objectives.

Scope: The paper introduces EraseFlow, a GFlowNet-based, reward-free fine-tuning method that removes a targeted visual concept from a pretrained text-to-image diffusion model while preserving the model's prior, evaluated on nudity, artistic style, and fine-grained object concepts using Stable Diffusion v1.4.

What This Paper Is About

Existing methods for removing harmful or proprietary concepts from text-to-image generators tend to either damage image quality, depend on brittle adversarial losses, or require hours of training per concept. The authors argue the root cause is that these methods treat all denoising steps equally, ignoring how uncertainty and conditional distributions evolve across the denoising trajectory. EraseFlow instead learns a stochastic policy over whole denoising trajectories, using a GFlowNet trajectory-balance objective with a constant reward, so that generation under a target prompt is steered toward the distribution of a safe "anchor" prompt.

Key Contributions

1. EraseFlow framework. The authors present what they describe as the first GFlowNet-based method designed specifically for efficient and robust concept erasure from pretrained text-to-image diffusion models, casting unlearning as exploration over denoising paths optimized with a trajectory-balance objective.
2. Reward-free alignment strategy. They eliminate the external reward model by assigning a constant reward β to trajectories generated under the safe anchor prompt and zero otherwise, which they claim generalizes to arbitrary unseen concepts without explicit reward specification.
3. Empirical performance and efficiency. EraseFlow is reported to reach state-of-the-art-like performance across concept erasure benchmarks while reducing computational overhead and preserving prior generative quality and robustness against adversarial reintroduction.
4. Plug-and-play composition. The method can be composed with adversarial methods (AdvUnlearn) and training-free filtering methods (SAFREE) to further boost results.

Main Findings

• Adversarial robustness on nudity erasure. On the UDAtk benchmark, EraseFlow reaches an attack success rate of 33.89, versus 64.40 for the second-best non-adversarial method DUO (a 30.51% improvement) and 50.84 for the adversarial method R.A.C.E (a 16.95% improvement). AdvUnlearn alone reaches 16.94, but the paper notes it uses adversarial fine-tuning during UDAtk evaluation.
• Near-elimination of nudity with plug-and-play. EraseFlow combined with AdvUnlearn reaches 1.42 on UDAtk, which the authors describe as 1% failure rates compared with a previous best of 16%; they also frame this as a 15.52% gain over AdvUnlearn.
• Consistent gains across red-teaming datasets. On I2P, Ring-a-Bell, and MMA-Diff, EraseFlow records 2.80, 0.00, and 0.60 respectively, against SDv1-4's 93.66, 59.49, and 55.2. EraseFlow + SAFREE records 2.10, 0.00, and 0.60, and EraseFlow + AdvUnlearn records 1.40, 0.00, and 0.30.
• Trajectory balance outperforms detailed balance. A comparison table reports I2P / Ring-a-Bell / MMA-Diff scores of 8.3 / 6.39 / 14.1 for the detailed-balance objective with reward, 2.1 / 2.53 / 1.7 for trajectory balance with reward, and 2.8 / 0.00 / 0.60 for EraseFlow.
• Artistic style erasure. On UDAtk for average style erasure across "Van Gogh" and "Caravaggio", EraseFlow scores 65.43, versus 66.65 for DUO, 68.49 for ESD, and 70.03 for SAFREE. The paper states EraseFlow outperforms the non-adversarial methods by at least 1%, and that plug-and-play with AdvUnlearn and SAFREE improves EraseFlow by 17.59% and 2.55% respectively.
• Fine-grained erasure balance. For the average of removing the "Nike logo" from Nike shoes, the "Coca-Cola logo" from bottles, and "wings" from a Pegasus, EraseFlow achieves a concept score of 82.24 and a total score of 76.01, which the authors say outperforms the previous best by 5.31% in total score. ESD reaches a higher concept score (93.97) but a lower total score (59.40), which the paper attributes to excessive erasure.
• Image quality and alignment retention. On 10,000 MSCOCO prompts, EraseFlow records a CLIP score of 25.67 and FID of 17.93. The paper states it matches UCE (25.59) and DUO (26.36) on CLIP score and outperforms all baselines on FID except MACE (17.11).
• Training efficiency. EraseFlow trains in 2.8 minutes as reported in Table 1, and the text describes "just 3 minutes of training on a single A100 GPU," compared with 225 minutes for R.A.C.E and 1440 minutes for AdvUnlearn. Peak memory for EraseFlow is listed as 42.00 GB.
• Reward-free objective works better than the reward-based one. The paper reports that the detailed-balance objective with reward performs poorly, while the trajectory-balance version works well, and the constant-reward EraseFlow formulation works best.
• Ablation on log β. Small values (log β ≤ 1) cause poor erasure and training instability; values in the range [2, 3] give a sharp 96% improvement and the best erasure–quality trade-off; larger values (log β ≥ 50) lower FID further but degrade erasure performance.
• Ablation on STOP_SAMPLING. Performance improves steadily as STOP_SAMPLING increases, with the best results reported at epoch 20.

Methodology in Plain English

A diffusion model generates an image by starting from noise and repeatedly denoising, producing a trajectory of intermediate latents. EraseFlow treats each latent in that trajectory as a state in a directed acyclic graph and treats the model's reverse denoising step as the GFlowNet forward policy and the noising step as the backward policy.

The framework defines a target prompt c (for example "nudity") and a semantically safe anchor prompt c* (for example "fully dressed person"). Instead of scoring individual final images with a reward model, EraseFlow assigns a single constant reward β to trajectories generated under the anchor prompt and zero to others. The training loss then forces the flow under the target prompt to match the density of anchor trajectories, effectively transplanting the safe distribution of c* onto the prompt c.

The authors show theoretically (Proposition 4.1) that if the constant-reward loss reaches zero and the standard trajectory-balance constraint holds for the original model with the safe prompt, then the per-timestep denoising distributions under c and c* become equal, and the marginal image distributions coincide, meaning the concept unique to c is completely erased. The proof sketch argues that zero loss implies the trajectory-balance identity holds for every anchor trajectory, and subtracting the corresponding identity for c* eliminates the shared noising term, leaving equal per-timestep log-probabilities.

Training samples an anchor trajectory from the diffusion model conditioned on c*, pairs it with the target prompt c, and applies the loss across timesteps. For memory efficiency, the authors sample only a subset of timesteps—10 timesteps from the first 40 denoising steps and 10 from the last 10 steps. They fine-tune only up to a STOP_SAMPLING epoch to prevent drift and entanglement. EraseFlow is trained for 20 iterations, each using a single data batch, with log β set to 2.5, STOP_SAMPLING set to 21 for nudity erasure, 11 for fine-grained erasure, and 1 for artistic style erasure. Learning rates are 3.0 × 10⁻⁴ for nudity and fine-grained tasks and 5.0 × 10⁻⁴ for artistic style erasure.

Evaluation uses Stable Diffusion v1.4 as the backbone across three tasks: NSFW suppression (142 I2P prompts, 79 Ring-a-Bell prompts, 1000 MMA-Diffusion prompts, and 142 I2P prompts extracted with UDAtk), artistic style erasure for "Van Gogh" and "Caravaggio" using 50 UDAtk adversarial prompts per style, and fine-grained erasure using 10 prompts per concept generated with GPT-4o and 10 images per prompt, scored with Gecko and inspired by EraseBench. NSFW is measured by attack success rate with NudeNet at a threshold of 0.6; style erasure uses mean cosine similarity with CSD features; image quality uses CLIP score and FID on MSCOCO. Baselines include ESD, UCE, MACE, DUO, SAFREE, R.A.C.E, and AdvUnlearn—7 open-sourced concept erasure methods in total.

Note: the supplied paper content is truncated partway through Appendix B, so appendices C through J (referenced for additional ablations, prompt examples, and qualitative results) are not available in the provided text.

Why This Matters

Impact on research. The paper offers a formal claim that a constant reward combined with a trajectory-balance objective can exactly align an erased concept's distribution with a safe anchor's distribution, and it reports that this reward-free route avoids the instability and hackability the authors associate with adversarial or reward-model-based unlearning. It also proposes the first application of GFlowNets to concept erasure.

Real-world applications:

• Content moderation: suppressing NSFW generation when models are given implicit or explicit prompts, including adversarial prompts from UDAtk.
• Copyright and brand protection: removing corporate logos such as Nike or Coca-Cola from generated images.
• Artist and style protection: erasing artistic styles such as "Van Gogh" and "Caravaggio."
• Fine-grained attribute control: removing a single attribute such as "wings" from a "Pegasus" prompt while leaving the horse and mountain range intact.

Industry relevance. The reported cost profile is the headline: 2.8 minutes and 42.00 GB peak memory in Table 1 (with the text citing "just 3 minutes" on a single A100 GPU) versus 225 minutes for R.A.C.E and 1440 minutes for AdvUnlearn. The method is described as a lightweight, plug-and-play safety primitive that composes with SAFREE and AdvUnlearn, which matters for teams that need to patch deployed generators repeatedly.

Future Directions

• Extending the reward-free formulation to concepts that lack an obvious natural-language anchor prompt, since the framework requires a safe counterpart c* for each target concept c.
• Investigating why the detailed-balance objective with reward underperforms while trajectory balance works well, and whether the gap generalizes beyond the tested concepts.
• Exploring the training-stability boundary identified in the log β ablation, where values at or above 50 further lower FID but degrade erasure, and where values at or below 1 destabilize training.
• Determining how far the constant-reward alignment guarantee holds in practice under imperfect optimization, given Proposition 4.1's assumption that the loss reaches exactly zero.
• Scaling the evaluation beyond the reported settings—Stable Diffusion v1.4, 20 training iterations, and the three task families—to newer backbones and broader concept vocabularies.

Target Audience

Researchers and practitioners working on diffusion model safety, machine unlearning, and generative model alignment, particularly those familiar with GFlowNets, trajectory-balance objectives, or reinforcement-learning-based fine-tuning. It is also relevant to engineers deploying text-to-image systems who need a low-cost concept removal step that can be combined with existing filtering or adversarial defenses. Readers without a background in diffusion models or flow-based generative methods will find the theoretical sections demanding.

Authors’ 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.