BEM: Learning Bregman Divergences for Concept Erasure in Flow Matching Models

Mohamed-Hicham Leghettas; Markus Püschel

cs.LGOct 17, 2026

BEM: Learning Bregman Divergences for Concept Erasure in Flow Matching Models

Authors: Mohamed-Hicham Leghettas, Markus Püschel

Organizations: Department of Computer Science, ETH Zurich, Switzerland

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.

Figures & tables

Method	I2P	T2I-RP	RAB	MMA	P4D	Basic	FID	min	T2I-RP	Basic	FID	min
	✗ nudity								✗ bloody gore
FLUX	18.05	52.04	66.67	27.70	44.49	76	0.00	0	84.27	88	0.00	0
GEM	6.87	22.11	27.72	4.60	15.07	8	8.64	4:05	53.85	0	4.91	7:56
BEM (Ours)	5.59	26.02 †	10.53 †	3.00	4.78 †	1 †	6.64	211:15	14.36 †	0.33	4.26	107:10

Table 1: Safety after erasing ✗ nudity (left) and ✗ bloody gore (right) on FLUX. Lower is better in every column, and the better of the two erased models is in bold. The gore rates of BEM are a mean over three seeds, every other entry a single run. † marks a rate of BEM that differs from GEM’s with

p<0.05

(two-sided Fisher exact test; App. A.3 gives

95\%

intervals). RAB is Ring-A-Bell.

Panel of Figure 1 — Figure 1: One prompt, six seeds. A T2I-RiskyPrompts gore prompt on which GEM leaks, generated from six noise seeds by the original model, by GEM and by BEM. BEM removes the blood on all six seeds; on two seeds it also changes the subject, rendering the pig as a clothed figure.

Panel of Figure 2 — Figure 2: Rights-protected erasure on FLUX. Above, one generation per concept and per model. Naruto was not erased and comes from the model that erases Son Goku, so it shows collateral damage. Below, erasure (Er., lower is better) and average retention (Ret., higher is better), with the better of the two erased models in bold, † as in Tab. 1 , and the time averaged over the four concepts.

Method	✗ Stitch	P	N	S	Avg.	min
	Erasure $\downarrow$	Retention $\uparrow$				Time
SD3.5	100	100	97	100	99.00	0
GEM	17	100	92	99	97.00	10:09
BEM (Ours)	3 †	97	84	95	92.00 †	438:47

Table 2: Erasing ✗ Stitch on SD3.5, over

100

generations, with in-domain retention on ✓ Pikachu (P), ✓ Naruto (N) and ✓ Snoopy (S). Bold and † are as in Tab. 1 .

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 3 — Figure 3: FLUX samples for the bloody gore setting, on three prompts from the blood subset of T2I-RiskyPrompts and one basic prompt. Rows are the original model, GEM as we reproduce it, and BEM. The rates are in Tab. 1 .

Panel of Figure 4 — Figure 4: SD3.5 samples after erasing Stitch , with the three retention characters. Rows are the original model, GEM as we reproduce it, and BEM. The rates are in Tab. 2 .

Panel of Figure 5 — Figure 5: Qualitative nudity erasure results. Rows are the original model, GEM as we reproduce it, and BEM. The first four columns are one prompt from each of four benchmarks, with nudity censored (NudeNet detections, completed by hand), and the last column probes general utility on an MS-COCO prompt.

Panel of Figure 6 — Figure 6: Qualitative FLUX samples for four erasure settings: ✗ bloody gore, ✗ Albert Einstein, ✗ Son Goku and ✗ Studio Ghibli. Rows are the original model, GEM and BEM. The columns of an erased concept come from the model that erases it, and the ✓ Clinton and ✓ Naruto columns from the models that erase Albert Einstein and Son Goku, so they report what the erasure costs a related concept.

Panel of Figure 7 — Figure 7: Erasing the Studio Ghibli style on FLUX. Each column is one of the ten prompts of this concept, named by its first words (the prompts end in “of Studio Ghibli”, or “featuring Studio Ghibli” for the dramatic scene), and shows the first of its ten seeds. Rows are the original model, GEM and BEM, all from the same noise.

Panel of Figure 8 — Figure 8: Erasure and retention side by side. For each rights-protected concept, the left column is the concept prompt and the right column its nearest retention prompt, both from the model that erases the concept on the left. Rows are the original model, GEM and BEM, all from the same noise.

Panel of Figure 9 — Figure 9: Held-out phrasings. GEM and both levels of BEM train on the literal name alone (first column of each block, the benchmark prompt); the other three columns are phrasings on which neither model was trained: an alias in a scene, a clay figure, and a plush toy. Rows are the original model, GEM and BEM, all from the same noise.

	Unsafe rate % $\downarrow$						Utility	Time
Method	I2P	T2I-RP	RAB	MMA	P4D	Basic	FID $\downarrow$	min $\downarrow$
as reported by Grebe et al. (2026)
FLUX	20.20	51.60	63.86	27.40	47.43	77	0.00	0
ESD	17.62	46.89	62.11	14.00	41.54	56	4.12	32:26
UCE	18.69	49.29	55.44	16.90	38.60	73	2.47	0:12
CA	12.43	32.84	47.19	19.25	25.36	53	8.12	35:14

Appendix

Table 3: Safety after erasing ✗ nudity on FLUX. The unsafe rate is measured with NudeNet and utility with the FID. In each block, the best entry is in bold and the second best underlined; the original model and the time are not ranked.

	Unsafe rate % $\downarrow$		Utility	Time
Method	T2I-RP	Basic	FID $\downarrow$	min $\downarrow$
as reported by Grebe et al. (2026)
FLUX	83.93	100	0.00	0
ESD	73.68	4	5.04	33:04
CA	78.97	66	5.97	32:12
UCE	79.83	50	2.64	0:12

Appendix

Table 4: Safety after erasing ✗ bloody gore on FLUX. The unsafe rate is measured with Q16 and utility with the FID. The rates of BEM are a mean over three seeds. Bold and underlined entries are as in Tab. 3 .

	Erasure $\downarrow$	Retention $\uparrow$	Erasure $\downarrow$	Retention $\uparrow$	Time
Method	✗ Einstein	Average	✗ Merkel	Average	min $\downarrow$
as reported by Grebe et al. (2026)
FLUX	100	100.00	100	100.00	0
ESD	1	54.67	0	68.33	6:27
UCE	0	77.33	0	73.00	0:12
EraseFlow	2	58.00	0	16.67	5:15

Appendix

Table 5: Celebrity erasure on FLUX, evaluated on

100

generations. We erase ✗ Albert Einstein (left) and ✗ Angela Merkel (right), measuring average retention on ✓ Nelson Mandela, ✓ Hillary Clinton and ✓ Barack Obama. Our block uses a local Qwen2.5-VL-7B classifier in place of Gemini 2.5 Flash, so its rates are comparable within the block only. Bold and underlined entries are as in Tab. 3 .

	Erasure $\downarrow$	Retention $\uparrow$	Erasure $\downarrow$	Retention $\uparrow$	Time
Method	✗ Stitch	Average	✗ Son Goku	Average	min $\downarrow$
as reported by Grebe et al. (2026)
FLUX	100	100.00	100	100.00	0
ESD	0	91.67	2	70.67	13:04
UCE	0	95.33	1	95.33	0:13
EraseFlow	8	86.00	2	67.67	7:48

Appendix

Table 6: Copyrighted character erasure on FLUX, evaluated on

100

generations. We erase ✗ Stitch (left) and ✗ Son Goku (right), measuring average retention on ✓ Pikachu, ✓ Naruto and ✓ Snoopy. Our block uses the classifier of Tab. 5 . Bold and underlined entries are as in Tab. 3 .

	Erasure $\downarrow$	Retention $\uparrow$				Time
Method	✗ Stitch	P	N	S	Avg.	min $\downarrow$
as reported by Grebe et al. (2026)
SD3	100	100	100	99	94.67	0
ESD	6	94	78	76	82.67	7:38
EraseFlow	0	61	17	13	30.67	5:03
GEM	2	95	93	82	90.00	2:28

Appendix

Table 7: Copyrighted character erasure on Stable Diffusion, evaluated on

100

generations. We erase ✗ Stitch, measuring in-domain retention on ✓ Pikachu (P), ✓ Naruto (N) and ✓ Snoopy (S). Both blocks use SD3.5 Medium, which Grebe et al. (2026) call SD3. The averages of the SD3 and EraseFlow rows are printed as in Grebe et al. (2026) ; their three retention rates give

99.67

and

30.33

. Bold and underlined entries are as in Tab. 3 .

Setting	$n$	GEM	BEM	$p$
Nudity, I2P	931	6.87 [5.4, 8.7]	5.59 [4.3, 7.3]	$0.29$
Nudity, T2I-RP	1,126	22.11 [19.8, 24.6]	26.02 [23.5, 28.7]	$0.034$
Nudity, Ring-A-Bell	285	27.72 [22.8, 33.2]	10.53 [7.5, 14.6]	$2.2\cdot 10^{-7}$
Nudity, MMA-Diffusion	1,000	4.60 [3.5, 6.1]	3.00 [2.1, 4.3]	$0.079$
Nudity, P4D	272	15.07 [11.3, 19.8]	4.78 [2.8, 8.0]	$7.9\cdot 10^{-5}$
Nudity, Basic	100	8.00 [4.1, 15.0]	1.00 [0.2, 5.4]	$0.035$

Appendix

Table 8: Rates in percent with Wilson

95\%

intervals, over

n

generations per model, and the

p

-value of a two-sided Fisher exact test between GEM and BEM.

Part	Hyperparameter	Value
Context path	noise-level embedding	$64$
	prompt features	$128$
	T5 attention	$1$ head of width $64$
	token and global encodings ${\bm{c}}_{i}$ , ${\bm{c}}$	$128$
	width inside the potential	$64$
	transformer blocks	$2$ , with $4$ heads and MLP ratio $2$

Appendix

Table 9: Hyperparameters of the potential and of its context path, shared by all runs except where noted. The context path is larger on SD3.5, whose pooled text embedding has

2048

dimensions against

768

on FLUX.

Method	Reference	Models	Acts on	Objective
Closed-form edits
TIME	Orgad et al.,2023	SD	K/V projections	least squares, closed form
ReFACT	Arad et al.,2024	SD	text encoder layer	closed-form edit of the text encoder
UCE	Gandikota et al.,2024	SD	K/V projections	least squares, closed form
RECE	Gong et al.,2024	SD	K/V projections	least squares, alternated with adversarial embeddings
MACE	Lu et al.,2024	SD	cross-attention, LoRA	least squares, then a masked attention-map loss

Appendix

Table 10: Erasure methods and attacks: models on which they are evaluated, where they act, and what they minimize.

Diffusion Model Unlearning

Diffusion model unlearning removes specified concepts or training examples from generative models while preserving unrelated generation.

See all 77 papers

FlowErase-RL: Rethinking Concept Erasure as Reward Optimization in Flow Matching Models
May 19, 2026Yi Sun, Zhiqi Zhang, Xinhao Zhong +5Flow Matching T2I Generation
Harbin Institute of Technology, Shenzhen · Jilin University · Tsinghua Shenzhen International Graduate School, Tsinghua University +2
FlowErase-OPD: Multi-Concept Erasure via Anchored On-Policy Distillation in Flow Matching Models
Aug 7, 2026Yi Sun, Yimin Zhou, Xinhao Zhong +3Flow Matching Diffusion Model Unlearning
Harbin Institute of Technology, Shenzhen · Tsinghua Shenzhen International Graduate School, Tsinghua University · Peng Cheng Laboratory
Geometric Erasure by Contrastive Velocity Matching in Rectified Flows
May 29, 2026Jonas Henry Grebe, Tobias Braun, Anna Rohrbach +1Contrastive Flow Matching Rectified Flow
Technical University of Darmstadt & hessian.AI, Germany.

May 19, 2026J/K move · Enter open · S save

FlowErase-RL: Rethinking Concept Erasure as Reward Optimization in Flow Matching Models

Yi Sun, Zhiqi Zhang, Xinhao Zhong, Yimin Zhou+4

Harbin Institute of Technology, Shenzhen · Jilin University · Tsinghua Shenzhen International Graduate School, Tsinghua University+2

Flow Matching T2I Generation Safe T2I Generation Flow-Based Generative Modeling RL for Generative Models Concept Erasure

Open paper AskPDF arXiv

Recent advances in flow matching models have significantly improved text-to-image generation quality, but also introduce growing safety risks due to the generation of harmful or undesirable content. Existing concept erasure methods are either inference-time interventions with limited effectiveness or rely on supervised fine-tuning (SFT), which requires precisely aligned data and struggles with scalability and multi-concept settings. In this paper, we propose FlowErase-RL, the first GRPO-based framework for concept erasure in flow matching models. We reformulate concept erasure as a reward optimization problem and introduce a dynamic dual-path reward mechanism that jointly optimizes (i) a Concept Erasure (CE) reward to suppress target concepts and (ii) a Non-target Space (NS) reward to preserve generative fidelity. The two reward paths are adaptively balanced during training via a performance-driven switching strategy, enabling stable optimization without explicit supervision. Extensive experiments on nudity, object, and artistic style erasure demonstrate that our method achieves state-of-the-art erasure performance while maintaining strong image quality and semantic alignment. Moreover, it exhibits robust resistance to adversarial attacks and scales effectively to multi-concept scenarios. Our results establish a new paradigm for safe and controllable generation in flow matching models.

Abstract

Diffusion Model Unlearning

Explore similar work

FlowErase-RL: Rethinking Concept Erasure as Reward Optimization in Flow Matching Models

FlowErase-OPD: Multi-Concept Erasure via Anchored On-Policy Distillation in Flow Matching Models

Geometric Erasure by Contrastive Velocity Matching in Rectified Flows