Flow matching models learn flexible concept erasure with bregman divergences
BEM: Learning Bregman Divergences for Concept Erasure in Flow Matching Models
Machine LearningComputer Vision and Pattern Recognition
Summary
Some AI models that create images from text can accidentally produce unwanted content like nudity or violence. The authors studied a new way to help these models forget specific unwanted concepts without harming their ability to make other images. They use a special kind of mathematical loss called a Bregman divergence, which lets them better control what is erased and what stays. Their method, called BEM, learns this loss automatically to remove harmful content more effectively than existing methods.
What this means in practice
- •For ai product teams: Improve safety controls in text-to-image AI by selectively removing harmful content while preserving other image qualities.$Commercial implications: Enables development of safer commercial AI image generation tools by effectively filtering unsafe or restricted imagery.
- •For content moderation engineers: Design better filters that erase sensitive or copyrighted concepts from AI-generated images without globally degrading model performance.
Authors
Mohamed-Hicham Leghettas, Markus Püschel
Abstract
Concept erasure removes harmful, private, or copyrighted concepts from text-to-image models while preserving their other abilities. Leading models such as FLUX and Stable Diffusion 3 are trained with flow matching, and erasure methods for them fine-tune the model toward a guided teacher's velocities under a Euclidean loss. Yet the squared distance is only one member of a larger family: Bregman divergences are exactly the losses under which flow matching recovers the correct velocity field, and each member shapes differently what is erased and what is preserved. We introduce BEM (Bregman erasure matching), which is, to our knowledge, the first to learn this divergence through bi-level optimization, so that the fine-tuned model forgets the concept on unseen prompts while its unrelated generations stay unchanged. In the experiments we target the suppression of nudity, violence, real identities, and copyrighted characters. Under the evaluation protocol of GEM, BEM removes more of the concept than GEM on most settings, cutting the unsafe rate on bloody gore from 53.9% to 14.4% and the FID on nudity from 8.64 to 6.64, while retention holds on some concepts and degrades on others.