EraseFlow: Learning Concept Erasure Policies via GFlowNet-Driven Alignment

Naga Sai Abhiram kusumba, Maitreya Patel, Kyle Min, Changhoon Kim, Chitta Baral, Yezhou Yang

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.