Title: EraseFlow: Learning Concept Erasure Policies via GFlowNet-Driven Alignment
Venue: NeurIPS 2025
Accepted: yes
Keywords: Safe generation

Text-to-image generation

Flow models

GFlowNets
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 a 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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a concept erasure framework, EraseFlow, that frames concept erasure as a reward-free distributional alignment task. EraseFlow evolves conditional marginal distributions through a trajectory balance (TB) formulation and adopts a reward-free alignment strategy to remove the dependency on manually designed and potentially vulnerable reward models. Experiments on the NSFW, artistic style, and fine-grained tasks demonstrate that EraseFlow performs well in terms of a trade-off between performance and prior preservation. EraseFlow can be combined with adversarial and filtering-based methods to boost their performance.\n\nSTRENGTHS AND WEAKNESSES:\n**Streagths**\n\n1. The paper is well written. I could understand their motivation and the core idea of the proposed method.\n2. The proposed method is simple, but the experimental results are good across several erasure tasks.\n\n**Weaknesses**\n\n1. An ablation study about the hyperparameter $\\beta$ would be desired. It would be great to both quantitatively and qualitatively show what happens when we set a too small/large value to $\\beta$.\n2. Can EraseFlow do multi-concept erasing? If EraseFlow cannot do multi-concept erasing, it should be included in the Limitations section.\n3. The submitted checklist says, \"We release the code for reviewers\", but their code cannot be found.\n\n**Minor comments**\n\n1. Suggestions\n    - L.197: \"alignment based\" -> \"alignment-based\"\n      - A hyphen should be inserted\n    - L.243: \"In this way EraseFlow\" -> \"In this way, EraseFlow\"\n      - A comma should be inserted\n    - L.275: \"STOP_SAMPLING\" should be non-italic.\n    - L.292: \"atleast\" -> \"at least\"\n      - A whitespace should be inserted\n    - Table 3: It is better to embed this table in the previous paragraph \"Fine-Grained Erasure Analysis\" (if possible)\n2. Including recent works as concurrent works will make the paper more valuable. See the \"Contemporaneous Work\" section of the [Call for Papers](https://neurips.cc/Conferences/2025/CallForPapers).\n    - Training-based\n      - Wang et al., \"ACE: Anti-Editing Concept Erasure in Text-to-Image Models\", https://arxiv.org/abs/2501.01633\n      - Bui et al., \"Fantastic Targets for Concept Erasure in Diffusion Models and Where To Find Them\", https://arxiv.org/abs/2501.18950\n    - Training-free\n      - Wang et al., \"Precise, Fast, and Low-cost Concept Erasure in Value Space: Orthogonal Complement Matters\", https://arxiv.org/abs/2412.06143\n      - Jain et al., \"TraSCE: Trajectory Steering for Concept Erasure\", https://arxiv.org/abs/2412.07658\n      - Lee et al., \"Localized Concept Erasure for Text-to-Image Diffusion Models Using Training-Free Gated Low-Rank Adaptation\", https://arxiv.org/abs/2503.12356\n\nQUESTIONS:\nThe topic is important, and the quantitative and qualitative results are good. However, I have several concerns about the paper, which I provided in \"Weaknesses\". I would appreciate it if the authors could address them. If their response is convincing, I will raise my rating.\n\nLIMITATIONS:\nyes\n\nA Limitations section is included in the supplementary material.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes EraseFlow, a novel framework for concept erasure in text-to-image diffusion models. Unlike existing approaches that rely on adversarial fine-tuning, or step-wise training, EraseFlow formulates unlearning as trajectory alignment over the full denoising process using trajectory balance. Specifically, the authors newly introduce a reward-free training strategy to eliminate the risks of reward hacking in auxiliary reward models. Extensive experiments demonstrate that EraseFlow outperforms baselines across multiple benchmarks.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths.**\n\n- The concept erasure with a reward-free trajectory balance objective is original (in terms of both conceptually and methodologically) and well-justified (in terms of both theoretically and empirically).\n- EraseFlow achieves superior trade-offs between erasure strength, preservation of performance, and computational cost.\n- The proposed method is bounded in modifying the fine-tuning objective and shows consistent performance improvements when combined with other orthogonal methods (adversarial methods or inference-time intervention).\n\n---\n\n**Weaknesses.**\n\n- The paper proposes shifting from existing transition-wise to trajectory-wise training. However, as the diffusion step $T$ increases, the variance of the trajectory balance objective may become large and unstable. Is there any mitigation strategy (or discussion) for this trade-off?\n- The current training uses only anchor trajectories $\\tau^{\\ast}$ sampled from the anchor prompt $c^{\\ast}$ (i.e., safe trajectories), and optimizes $ZP(\\tau^{\\ast}|c) = RP(\\tau^{\\ast}|c^{\\ast})$ to increase the likelihood of sampling safe trajectories under $c$. However, training solely on safe trajectories may lead to a large increase in $Z$ while under-estimating $P(\\tau^{\\ast}|c)$. Can authors clarify this potential issue and discuss possible mitigation strategies (e.g., incorporating risky trajectories $\\tau \\sim c$ to directly enforce $ZP(\\tau|c) = 0$)?\n\nQUESTIONS:\nSee weaknesses.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-06-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors tackle concept erasure for text-to-image diffusion models. Their central idea is to view unlearning as distribution matching along the entire denoising trajectory instead of at the final image only. They instantiate this idea with a Generative Flow Network (GFlowNet) trained with the trajectory-balance (TB) objective, but—crucially—with a constant reward. They prove that, if the TB loss is driven to zero, the edited model’s conditional marginals exactly match those of a safe “anchor” prompt (Proposition 4.1). A lightweight anchor-trajectory training procedure (Algorithm 1) updates the diffusion model for ≈3 minutes on one A100 GPU.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths \n1. First application of GFlowNets to concept erasure.\n2. Proposition 4.1 formally links constant-reward TB minimisation to exact distribution alignment.\n3. The paper is well written.\n\n## Weaknesses\n1. The method utilizes a single safe prompt (e.g., “fully-dressed person”) as the reference distribution, without conducting an ablation study to assess how the choice of anchor impacts the results.\n2. The evaluation does not include multi-concept erasure. It remains unclear how EraseFlow performs in scenarios involving multi-concept erasure, such as the 100-celebrity erasure tested in MACE.\n3. Memory cost reporting is expected to be included in Table 1.\n4. The experiments primarily focus on outdated UNet-based diffusion models. Evaluations using modern MMDiT-based diffusion models, such as SD3.5 or Flux, as conducted in EraseAnything [1], are anticipated.\n5. The discussion of related works is somewhat dated. Incorporating more recent studies [2,3,4] would provide readers with a more comprehensive and current perspective.\n\n[1] EraseAnything: Enabling Concept Erasure in Rectified Flow Transformers\n\n[2] Sculpting Memory: Multi-Concept Forgetting in Diffusion Models via Dynamic Mask and Concept-Aware Optimization\n\n[3] SPEED: Scalable, Precise, and Efficient Concept Erasure for Diffusion Models\n\n[4] Set You Straight: Auto-Steering Denoising Trajectories to Sidestep Unwanted Concepts\n\nQUESTIONS:\nView Weaknesses. If my concerns are fully resolved, I’m open to raising my score further. However, if the evaluation remains insufficient, I may reduce the score.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n- This paper introduces EraseFlow, a framework for concept erasure in pretrained text-to-image diffusion models. The goal is to remove specific undesired or harmful visual concepts (e.g., nudity, copyrighted logos) without degrading overall generative quality or requiring expensive retraining.\n- The paper leverages Generative Flow Networks (GFlowNets) and optimizes a trajectory balance (TB) objective to align the diffusion model’s denoising paths conditioned on these prompts\n- It introduces a reward-free training approach, where trajectories from the safe anchor are assigned a constant reward. This avoids the need for adversarial or classifier-based reward signals, which are often brittle or hard to define for arbitrary visual concepts.\n- Italso proposes a memory-efficient training procedure by subsampling timesteps from the early and late stages of denoising and demonstrate that EraseFlow is compatible with other unlearning techniques like SAFREE and AdvUnlearn.\n- Experiments show that EraseFlow achieves superior performance in erasing visual concepts across diverse categories (e.g., NSFW content, artistic styles, brand logos), while maintaining image quality and reducing attack success rates.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- The motivation for using full denoising trajectories (instead of final samples) is well explained, and the link to evolving conditional marginals is compelling.\n- Figures like 1 and 4, and the tables in Sections 5.1–5.2, support claims well. The table highlighting efficiency vs. erasure quality trade-offs (Figure 2, Table 1) is especially compelling.\n- To the best of my knowledge, this is the first paper to apply GFlowNet trajectory balance to the problem of unlearning in diffusion models. The reframing of concept erasure as a distributional trajectory matching task with provable guarantees is novel and technically interesting.\n- The approach is model-agnostic, requires only minimal fine-tuning, and is compatible with other methods like SAFREE or AdvUnlearn, showing potential for deployment in modular safety pipelines.\n\nWeaknesses:\n- While the mathematical derivation looks sound, it is dense and notation-heavy (e.g., Eq. 5–8), which may limit accessibility for broader audiences. The paper could benefit from more intuition, especially around how the trajectory-based alignment manifests in sample space.\n- While the paper includes comparisons with prior work, ablation studies isolating the effect of the anchor prompt choice, constant-reward value or trajectory sampling strategy are missing. These would strengthen understanding of what drives EraseFlow’s success.\n- The method assumes that the anchor prompt  is free of any visual overlap with the target concept. In practice, defining such a semantically \"safe\" anchor may not be trivial for all domains. The paper could discuss limitations or robustness to imperfect anchor prompts.\n\nQUESTIONS:\n- Could the authors conduct a sensitivity analysis or ablation on different anchors? What properties (semantic distance, visual dissimilarity) make a good anchor?\n- The constant reward formulation is elegant and simplifies the training pipeline. However, it's unclear why it is sufficient for effective erasure and whether there are conditions under which it leads to underfitting or ineffective alignment.\n- Can the authors provide more empirical evidence or ablations comparing constant vs. learned reward functions? What failure modes arise (if any) when the concept is highly entangled with surrounding distributions?\n- Can the approach generalize to multiple concepts or composite erasure?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-04"}