EraseFlow: Learning Concept Erasure Policies via GFlowNet-Driven Alignment

We are encouraged by your review! We thank you for your comprehensive evaluation of our paper. We are grateful that you found EraseFlow intuitive and performs well across the tasks. And we are glad to hear that our paper is clear and easy to follow. Please find the clarification requested below:

Ablation study on Beta

To understand the impact of the $\beta$ parameter, we conduct another focused ablation on removing the “nudity” concept by varying the $\beta$. As shown in the table below, having $\beta \leq 1$ performs very poorly on I2P evaluations. $\beta > 1$ makes the EraseFlow work really well and especially when $\beta \in [2-3]$. We attribute this behaviour to the term $log(\beta)$ in Eq. (8), which causes the smaller $\beta$ values to create training instability.

Multiconcept erasure

We thank the reviewer for their thoughtful comment regarding multiconcept unlearning. To explore this, we extended our experiments to settings involving a growing number of artist concepts to be unlearned. We evaluate performance using CLIP-based metrics for both unlearning and retention, along with the trade-off score Ha =Retain−Unlearn.

Our results show that EraseFlow retains its effectiveness across multiconcept scenarios, with the model continuing to forget the target concepts while maintaining overall generation quality. To support this, we introduce a retaining artists dataset and apply the EraseFlow loss to trajectories derived from anchor prompts. This encourages the model to focus on intended removals while preserving broader capabilities.

While the unlearning setting becomes more challenging as the number of concepts grows, EraseFlow maintains a balanced approach that emphasizes control and stability. This makes it especially useful in applications where targeted erasure with minimal disruption is desired. We will include these results and their implications in the final version of the paper.

Code release

We humbly apologize for this inconvenience. It seems that we missed uploading the scripts as part of the supplementary. We have reached out to AC to confirm if we can still share the anonymous github link. If we receive the permission, we are more than happy to make it available ASAP. Hopefully, we will get back to you on this soon. We again appreciate your careful consideration.

Minor comments

[typo] We thank the reviewer for sharing these typos. We will ensure that we incorporate these suggestions and look for any other remaining ones.

[references] We thank the reviewer again for sharing the missing references. As most of them are orthogonal works (especially training-free), they can be merged with EraseFlow, like we did for SAFREE and AdvUnlearn. We will ensure that we provide proper comparisons and add them to our related works. Additionally, we evaluated another method, TraSCE, in a plug-and-play fashion. When combined with EraseFlow, we observed improved erasure performance while retaining general generation quality, suggesting their complementary strengths. This is consistent with our experiment with SAFREE.

We hope this clarifies all the remaining concerns and questions. Additionally, we offer a summary of responses to other reviews for your reference in the global response.

Thank you very much for your effort and comprehensive response. I can easily imagine that conducting additional experiments was very tough. I really appreciate it.

I have read all the review comments and the authors' responses. I think that the authors have addressed the concerns of the reviewers, including mine, well. Let me ask follow-up questions.

Regarding the ablation study on $\beta$, I recommend including not only the quantitative result but also qualitative results (generated images) in a revised version if there is a tendency in generated images when $\beta$ is too small/large (I know showing visual results in the rebuttal phase is prohibited). Also, I have an additional question. I am wondering why the performance changes so drastically between $\beta=1.00$ and $\beta=2.00$, not gradually across a wider range. Do the authors have any thoughts on this? This question is out of curiosity, not for assessing the paper.

Regarding the code release, thank you for asking the area chairs. Regardless of whether the code can be shared in the rebuttal phase, do you plan to make your code publicly available upon acceptance?

Thank you for your extensive experiments. These results and discussions are very informative and insightful. I hope to see the same ablation study on another benchmark (another erasure task) to know if the $\beta$ and $z_{\phi}$ values are largely affected by what concept we want to erase, but this discussion period will end very soon. I do not require it this time.

I will raise my rating. Thank you for your work.

We are thrilled by your review and truly appreciate the time and consideration given to our work. We are glad that EraseFlow is recognized as an original approach and well-justified. We are happy to hear that our method achieves a superior trade-off in performance and efficiency while being able to work along with orthogonal methodologies to further boost the performance. In response to your inquiries, please find our clarifications below:

Variance of the trajectory balance objective

We agree with the reviewer that as diffusion steps T increases, the system could become more unstable. And especially, the memory requirement will also explode. To mitigate these limitations, we proposed to leverage the fixed time steps during the training. Specifically, we only select 20 time steps (10 initial denoising steps and 10 of the denoising process). These particular choices build upon the prior findings that initial denoising steps are highly stochastic, and later, they become deterministic. This allows us to bypass the potential instability. We will make this clear in the final draft for future readers.

Discussion on training solely on safe trajectories

We thank the reviewer for this thoughtful comment. In our preposition, we assumed that $Z_\phi$ will become the same post-training, hence they cancel out. However, this may not be ideal all the time, as the reviewer pointed out, and could lead to potential under-estimation. Firstly, we analyzed our pretrained checkpoints and found that $Z_\phi$ post-training values converged to about 0.83 for all tasks, which shows that it does not explode. However, to mitigate this potential issue for future adaptation, there could be several strategies:

• As the reviewer pointed out, making ZP(t,c)=0 is a good direction. However, as shown in the table below, we observe this leads to mode collapse – causing the model to generate noise for all prompts (see table below). We further ablate over possible design choices.
• Another option could be encouraging the optimal behavior of P(t,c) by leveraging the diffusion SDE objective on random prompts as a regularizer to avoid such mode collapse.

We trust that our response adequately addresses your concerns. We look forward to the discussion.

We are grateful for the time and consideration given to our paper, EraseFlow. We are happy to see that our work is well-founded in theory and provides a novel approach for lightweight concept erasure along with well-written paper.

In response to your inquiries, please find our clarifications below:

Ablations of the anchor prompt

To understand the impact of the anchor prompt choice, we conduct a focused study on removing the “nudity” concept using various anchor prompts. We generate 9 random prompts using ChatGPT, ranging from those semantically close to the target prompt → to random → to opposite. We observe a clear trend: as the anchor prompt becomes more semantically opposite to the target concept, the effectiveness of unlearning improves. Notably, using an empty prompt yields the best results when a suitable anchor is hard to define. For instance, when removing specific brand logos like the Nike logo from images, it is difficult to define a clear semantic “opposite”—in such cases, using an empty prompt serves as a practical fallback.

Extension to multiconcept erasure

In the 100-celebrity multiconcept erasure task, we observe that EraseFlow shows relatively lower forgetting performance compared to MACE. We believe this stems from EraseFlow’s design to apply only marginal updates to denoising trajectories, which contributes to its robustness—particularly against adversarial or off-distribution prompts. However, in cases involving multiple visually similar concepts, such as human faces, this conservative behavior can lead to interference across targets.

In contrast, for domains like artistic style unlearning, where concepts are more globally distinct, EraseFlow scales well and maintains consistent performance, as shown in the artist multiconcept experiments. We will explore ways to further improve scalability while retaining EraseFlow’s precision, and include the detailed discussion in the final version of the paper.

Memory Cost

We appreciate this valuable comment. To provide a fair comparison, we now report peak GPU memory for all methods on the same A100 and batch size. EraseFlow consumes more memory than schemes that update only key–value weights (e.g., UCE, MACE) because it back-propagates through every trajectory step, yet it still converges markedly faster than prior gradient-based baselines such as DUO, R.A.C.E., and AdvUnlearn. Overall, EraseFlow offers a compelling trade-off: while more computationally involved than ultra-lightweight methods, it is orders of magnitude more efficient than previous gradient-based unlearning approaches, without sacrificing effectiveness as shown in Table 1.

We will include GPU utilization in Table 1 in the camera-ready version to improve the completeness and transparency of our evaluation.

Generalization to SD3.5/Flux

We thank the reviewer for this suggestion. Firstly, we would like to share that such an alignment process generally depends upon stochasticity to calculate the log probability from Eq. (8). However, SD3.5 and Flux are deterministic flow models, which makes calculating the log prob a trivial case. Therefore, EraseFlow does not extend to such models.

Having said that, we leverage recent advancements and incorporate ODE to SDE conversion process first to calculate the log probabilities [1] and apply the EraseFlow on SD3.5/Flux models. Below, we compare EraseFlow with baselines, and we can observe that EraseFlow consistently achieves better performance even on Flux.

Additional Related Works

We are truly grateful to the reviewer for referring us to missing references. And we humbly apologize for the same. We will incorporate them in the final draft. While we recognize that most are concurrent works due to NeurIPS policy, we took the liberty and made a quick comparison with EraseAnything, and in the above experiments, we outperformed the new baseline on SD3.5/Flux.

We trust that our response adequately addresses your concerns and encourages you to reevaluate our submission. We look forward to the discussion.

---

[1] Jie Liu, Gongye Liu, Jiajun Liang, Yangguang Li, Jiaheng Liu, Xintao Wang, Pengfei Wan, Di Zhang, & Wanli Ouyang. (2025). Flow-GRPO: Training Flow Matching Models via Online RL. arXiv:2505.05470.

Thank you for your thorough review. We’re pleased that you found our model-agnostic approach—offering provable guarantees and fresh perspectives on unlearning—well supported. Please find our clarifications below:

Providing intuitive explanations for future readers.

We thank the reviewer for providing the feedback. To improve the readability for future readers, we will introduce a nice, intuitive explanation in the final version as follows:

Imagine the diffusion process as navigating a maze from noise to a final image, where each path represents a denoising trajectory. Some trajectories lead to undesired generations, while others lead to safe outputs. EraseFlow assigns a constant reward only to anchor-like (safe) trajectories, and uses the trajectory balance constraint in Eq. (6) to align the reverse sampling process under the target prompt with the reward-weighted forward process from the anchor. The partition function $Z_\phi$ serves as a normalizer over all trajectory rewards and enables this global alignment. Minimizing this loss increases flow through safe trajectories and suppresses unsafe ones, causing the model to favor anchor-like generations—even when given unsafe prompts. This is how trajectory-based alignment manifests in sample space: by reweighting denoising paths to avoid the target concept while preserving generation quality.

Ablations over the several design choices.

We agree that having these ablations would further strengthen our work. Therefore, we conducted the following ablations as requested.

Anchor prompt choice. To understand the impact of the anchor prompt choice, we conduct a focused study on removing the “nudity” concept using various anchor prompts. We generate 9 random prompts using ChatGPT, ranging from those semantically close to the target prompt → to random → to opposite. We observe a clear trend: as the anchor prompt becomes more semantically opposite to the target concept, the effectiveness of unlearning improves. Notably, using an empty prompt yields the best results when a suitable anchor is hard to define. For instance, when removing specific brand logos like the Nike logo from images, it is difficult to define a clear semantic “opposite”—in such cases, using an empty prompt serves as a practical fallback.

$\beta$ parameter ablation. To understand the impact of the $\beta$ parameter, we conduct another focused ablation on removing the “nudity” concept by varying the $\beta$. As shown in the table below, having $\beta \leq 1$ performs very poorly on I2P evaluations. $\beta > 1$ makes the EraseFlow work really well and especially when $\beta \in [2-3]$. We attribute this behaviour to the term $log(\beta)$ in Eq. (8), which causes the smaller $\beta$ values to create training instability.

Trajectory sampling strategy. We appreciate the reviewer’s suggestion. To study the effect of trajectory sampling, we conducted ablations on two key aspects: the switch epoch and the selection of timesteps used in the training loss.

• Switch Epoch Ablation: A larger switch epoch—sampling new anchor trajectories more frequently—gives the model a richer stream of safe examples, improves credit assignment, and yields stronger concept erasure. Conversely, a small switch epoch limits trajectory diversity and weakens erasure. This underscores the need for consistent guidance early in training. |Switch Epoch|I2P ↓|Ring-a-Bell ↓|CLIP ↑|FID ↓| |:--|--:|--:|--:|--:| |0|43.66|67.07|25.86|18.99| |5|15.19|34.17|25.12|17.94| |10|12.67|29.11|24.40|18.21| |15|9.10|41.77|24.50|17.83| |20|2.80|0.00|25.67|17.93|
• Timestep Selection Ablation: We also examined how the choice of timesteps affects performance. Sampling only early or only late timesteps weakens erasure, while uniform sampling adds noisy gradients. A small mix of early + late steps balances uncertainty and alignment, producing steadier training and stronger concept removal. These results reinforce EraseFlow’s core idea: stable, effective alignment requires sampling across the trajectory’s evolution, not just its endpoints. | Timesteps | I2P ↓ | Ring-a-Bell ↓ | CLIP ↑ | FID ↓ | |:------------------------|--------:|----------------:|---------:|--------:| | First 20 | 7.00 | 26.50 | 24.15 | 18.09 | | Last 20 | 6.30 | 24.00 | 24.30 | 17.78 | | Random 20 | 14.00 | 39.20 | 24.54 | 17.70 | | 10 random (from first 40) + last 10 | 2.80 | 0.00 | 25.67 | 17.93 |

Impact of the choice of anchor prompts.

We thank you for this valuable comment. We have provided a small ablation over the choice of ablation study grounded in semantic similarity above. Here, we further extend the discussion.

First, we agree that the choice of anchor prompt is critical, as also demonstrated in recent work [1]. As shown in the table above, our results indicate a clear pattern: anchor prompts that are semantically opposite to the target concept consistently yield better unlearning performance compared to random choices. Furthermore, in scenarios where defining a meaningful anchor prompt is inherently difficult—such as in our company logo removal benchmark—our method remains effective. For example, when removing the Nike logo from shoes, there is no obvious “opposite” prompt to guide unlearning, yet our approach still performs reliably.

Sufficiency of constant reward.

The goal of EraseFlow is straightforward: make the target prompt trajectory behave exactly like a safe anchor prompt trajectory. Because we sample every trajectory from this anchor, all trajectories are “safe” by design, so we never have to score or inspect those from the unsafe prompt. The reward can therefore be held constant—it merely normalizes the trajectory-balance loss rather than steering it. Minimizing this loss on anchor-derived trajectories is enough to copy the anchor’s behavior onto the target prompt, achieving concept erasure without any separate reward model. This also aligns with the next query below. We will add this clearer explanation to the main paper.

Constant vs. Learned reward.

We have provided this ablation in Figure (3) of the main paper, where we specifically analyzed the effect of the NSFW classifier as the reward model. We observed that TB with learned reward itself is better than the DB objective. However, constant reward further improves the performance of TB objectives. Importantly, we observed that the NSFW classifier introduced some adversarial effects into the system. If we take a look at our TB formulation more carefully, Eq. (8), we can see that the reward is calculated for the anchor trajectory, which is on the “safe” prompt. If we know it is a safe image, then there is no need for the learned reward model because it is safe all the time. Notably, this is not possible for standard RL alignment methods as the reward is being calculated for the potentially generated unsafe image.

Extension to multiconcept erasure

We scaled EraseFlow to progressively larger sets of artist concepts and evaluated performance using CLIP-based unlearning and retention scores, along with the trade-off metric $H_a$ = Retain − Unlearn. Across all settings, EraseFlow consistently forgot the target concepts while preserving overall image quality. We also introduced a small retaining-artists set and applied the same loss on anchor-prompt trajectories, ensuring removals stay focused without degrading other capabilities. Even as the task grew harder with more concepts, EraseFlow remained stable and controllable, demonstrating practical targeted erasure with minimal disruption. Full results and analysis will appear in the camera-ready version.

We hope these clarifications address all your concerns and enhance the comprehensiveness of our work. We look forward to the discussion.

---

[1] Anh Bui, Trang Vu, Long Vuong, Trung Le, Paul Montague, Tamas Abraham, Junae Kim, & Dinh Phung. (2025). Fantastic Targets for Concept Erasure in Diffusion Models and Where To Find Them. arXiv:2501.18950.