Z-Erase: Enabling Concept Erasure in Single-Stream Diffusion Transformers

Nanxiang Jiang; Zhaoxin Fan; Baisen Wang; Daiheng Gao; Junhang Cheng; Jifeng Guo; Yalan Qin; Yeying Jin

cs.CVMar 26, 2026

Z-Erase: Enabling Concept Erasure in Single-Stream Diffusion Transformers

Authors: Nanxiang Jiang, Zhaoxin Fan, Baisen Wang, Daiheng Gao, Junhang Cheng, Jifeng Guo, Yalan Qin, Yeying Jin,

Organizations: Beijing Advanced Innovation Center for Future Blockchain and Privacy Computing, School of Artificial Intelligence, Beihang University. · University of Science and Technology of China · Shanghai University · Tencent ·

Abstract

Concept erasure serves as a vital safety mechanism for removing unwanted concepts from text-to-image (T2I) models. While extensively studied in U-Net and dual-stream architectures (e.g., Flux), this task remains under-explored in the recent emerging paradigm of single-stream diffusion transformers (e.g., Z-Image). In this new paradigm, text and image tokens are processed as a single unified sequence via shared parameters. Consequently, directly applying prior erasure methods typically leads to generation collapse. To bridge this gap, we introduce Z-Erase, the first concept erasure method tailored for single-stream T2I models. To guarantee stable image generation, Z-Erase first proposes a Stream Disentangled Concept Erasure Framework that decouples updates and enables existing methods on single-stream models. Subsequently, within this framework, we introduce Lagrangian-Guided Adaptive Erasure Modulation, a constrained algorithm that further balances the sensitive erasure-preservation trade-off. Moreover, we provide a rigorous convergence analysis proving that Z-Erase can converge to a Pareto stationary point. Experiments demonstrate that Z-Erase successfully overcomes the generation collapse issue, achieving state-of-the-art performance across a wide range of tasks.

Figures & tables

Figure 1: We introduce Z-Erase , a concept erasure method tailored for single-stream models. First row : Prior methods adapted via our Stream Disentangled Concept Erasure Framework are tested with the “nudity” concept, showing under-erasure (AC, EraseAnything) or severe artifacts (UCE). Blue boxes reveal the generation collapse caused by naive fine-tuning without our adaptation. Second row : Z-Erase effectively removes target concepts while preserving quality. Original outputs are in yellow boxes. Sensitive content pixelated.

Figure 2: Single-stream attention analysis. Given the prompt “A girl with a wristwatch on her hand amid cherry blossoms” , the attention maps reveal distinct localized responses for text tokens.

Figure 3: Left: Attention erasure is brittle. Concept can be erased by zeroing its attention columns $\mathbf{A}_{attn}[:,:,idx]=0$ once the target token is localized from the prompt. However, this trick shows poor robustness to real-world prompt variations. Right: Naive fine-tuning collapses. Directly fine-tuning shared projections in single-stream models collapses generation with noisy outputs.

Figure 4: Dynamic balancing of erasure and preservation gradients. Left: When $\nabla\mathcal{L}_{er}$ does not conflict with $\nabla\mathcal{L}_{pr}$ , the erasure gradient itself is the update direction. Right: When the gradients conflict, the update direction $\mathbf{d}^{*}_{t}$ is obtained by steering $\nabla\mathcal{L}_{er}$ toward $\nabla\mathcal{L}_{pr}$ until the preservation constraint $\varepsilon$ is satisfied.

Figure 5: Single-concept erasure. Our method removes both concrete and abstract concepts while preserving image quality with minimal collateral changes. In contrast, UCE severely distorts images and introduces strong artifacts, whereas EraseAnything may fail to erase or causes noticeable semantic shifts.

Figure 6: User study. An interface is built to compare generated images from various erasure methods on Z-Image Turbo (see Appendix E for details). Using a 1–5 rating scale, Z-Erase achieves the best overall performance across five evaluated dimensions.

Method	Common	Female	Male	Total $\downarrow$	FID $\downarrow$	CLIP $\uparrow$	Detected Violence $\downarrow$	FID $\downarrow$	CLIP $\uparrow$
	Detected Nudity				MS-COCO 10K			MS-COCO 10K
AC (Model-Based) ( Kumari et al., 2023 )	245	35	53	333	29.30	29.05	402	32.47	29.89
AC (Noise-Based) ( Kumari et al., 2023 )	270	38	55	363	30.07	28.16	496	33.12	28.73
ESD-x ( Gandikota et al., 2023 )	279	42	60	381	31.44	27.65	777	33.79	27.48
ESD-u ( Gandikota et al., 2023 )	260	41	61	362	31.96	27.13	668	34.22	26.80
EAP ( Bui et al., 2024 )	190	29	49	268	27.32	30.71	447	28.30	30.15

Table 1: Evaluation on NSFW erasure. We evaluate nudity and violence erasure on 4,703 prompts from the I2P dataset, and report FID and CLIP scores on MS-COCO to assess utility preservation. Results of the original Z-Image Turbo are presented for reference.

Method	Acc e $\downarrow$	Acc ir $\uparrow$	H a $\uparrow$	FID $\downarrow$	CLIP $\uparrow$
AC (Model-Based)	27.95	27.10	-0.85	31.85	28.22
AC (Noise-Based)	27.34	26.88	-0.46	32.46	27.93
ESD-x	27.66	27.20	-0.46	33.15	27.41
ESD-u	25.13	26.68	1.55	34.06	26.95
MACE	24.06	28.52	4.46	30.79	30.17
UCE	21.91	25.70	3.79	52.28	23.70

Table 2: Evaluation on celebrity erasure. CLIP-based classification accuracy are reported on the erased identities ( Acc e , efficacy) and unaffected identities ( Acc ir , specificity). The combined performance is summarized by H a = Acc ir

-

Acc e , where higher values indicate a better balance between erasure and preservation. FID and CLIP are results based on MS-COCO 10K dataset.

	Entity		Art Style		Abstraction
Method	Acc e $\downarrow$	H a $\uparrow$	Acc e $\downarrow$	H a $\uparrow$	Acc e $\downarrow$	H a $\uparrow$
AC (Model-Based)	24.6	2.0	28.6	-4.2	26.9	-1.1
AC (Noise-Based)	25.5	0.9	29.3	-5.5	27.3	-1.1
ESD-x	23.3	5.3	28.5	-3.8	26.5	-0.8
ESD-u	21.8	6.0	27.8	-3.3	25.7	-0.4
MACE	23.2	4.3	26.6	-0.8	25.1	3 .4

Table 3: Evaluation on erasing specific category : Entity ( e.g., church), Artistic Style ( e.g., Claude Monet) and Abstraction ( e.g., color). We report erasure efficacy ( Acc e ) and balance score (H a ) here due to space limit. The full metrics are in Appendix D.2 .

				UnlearnDiffAtk
Method	w/o Attack	R-A-B	ReFlux	Step = 0	Steps = 0,1,2
DiT-Localization	57.89	58.94	74.73	59.64	6 0.35
EraseAnything	1 4.38	2 6.67	4 3.85	23.50	25.61
Token-Zeroing	21.75	82.81	92.85	72.50	76.49
Ours	10.52	24.91	42.81	2 4.56	25.61
Z-Image Turbo	83.15	92.63	96.84	82.45	85.96

Table 4: Evaluation against adversarial attacks. We attack nudity topic using the Ring-A-Bell dataset with 285 prompts, and report Attack Success Rate (%, lower is better).

Figure 7: Multi-concept erasure. Multiple concepts can be erased simultaneously via weighted LoRA averaging (see Appendix F ).

Config	Detected Nudity $\downarrow$	FID $\downarrow$	CLIP $\uparrow$
$\mathcal{L}_{erase}$	206	29.65	29.86
$\mathcal{L}_{erase}+\mathcal{L}_{attn}$	148	30.04	29.12
$\mathcal{L}_{erase}+\mathcal{L}_{pr}$	235	26.42	31.28
$\mathcal{L}_{attn}+\mathcal{L}_{pr}$	314	27.15	31.16
Full	161	26.46	31.25
Z-Image Turbo	649	26.33	31.47

Table 5: Ablation study on erasing nudity. We ablate three loss terms used in our method on 4,703 prompts from the I2P dataset.

Figure 8: Ablation study on hyperparameters

\varepsilon

and

\beta

. With our Lagrangian-Guided Adaptive Erasure Modulation algorithm, erasure and preservation become monotonically controllable.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 9: Abstract illustration of the Single-Stream DiT architecture. The workflow begins with two parallel pathways: the prompt is processed by a text encoder, while the image latent is patched and projected similarly. These tokens are concatenated into a Unified Input Sequence , which serves as the input to a stack of

N

Single Stream Blocks . Global conditioning information ( e.g. , Timestep) is injected into each block, influencing the layer normalization and attention scaling. Inside a Single Stream Block, the core Self-Attention mechanism processes the sequence using shared Query (

Q

), Key (

K

), and Value (

V

) projections for both text and image tokens.

Experiment / Setting	Recommended $\varepsilon$	Rationale
Nudity, Violence	$\varepsilon=1\times 10^{-3}$	Sensitive concepts, require tight preservation
Artistic Style	$\varepsilon=1\times 10^{-3}$	Global style is easy to remove, require tight preservation
Entity	$\varepsilon=1\times 10^{-2}$	Concrete objects harder to disentangle
Abstract Concepts	$\varepsilon=1\times 10^{-2}$	Abstract priors harder to disentangle
Celebrity Identity	$\varepsilon=1\times 10^{-2}$	Identity preservation more robust to slack
Downstream editability / utility (post-erase LoRA use)	$\varepsilon=5\times 10^{-3}$	Maintain editability & style generalization

Appendix

Table 6: Recommended

\varepsilon

values used for different settings in our experiments

Category	# Number	Prompt template	Concepts
Entity	10	‘A photo of [ Entity ]’	‘ Church ’, ‘ Soccer ’, ‘ Car ’, ‘ Airplane ’, ‘ Tower ’, ‘ Pencil ’, ‘ Pizza ’, ‘ Shoes ’, ‘ Cat ’, ‘ Clownfish ’
Artistic Style	10	‘An art in the style of [ Artistic Style ]’	‘ Pablo Picasso ’, ‘ Salvador Dali ’, ‘ Claude Monet ’, ‘ Vincent Van Gogh ’, ‘ Rembrandt van Rijn ’, ‘ Frida Kahlo ’, ‘ Edvard Munch ’, ‘ Leonardo da Vinci ’, ‘ Nicholas Roerich ’, ‘ Gustave Dore ’
Abstraction	10	‘A close scene of [ Abstraction ]’	‘ Shake Hand ’, ‘ Kiss ’, ‘ Hug ’, ‘ Time ’, ‘ Explosion ’, ‘ Green ’, ‘ Origami ’, ‘ Fried ’, ‘ Amidst ’, ‘ Sad ’

Appendix

Table 7: Complete list of concepts of Entity, Artistic Style, and Abstraction used in our experiment.

	Entity			Artistic Style			Abstraction
Method	Acc e $\downarrow$	Acc ir $\uparrow$	H a $\uparrow$	Acc e $\downarrow$	Acc ir $\uparrow$	H a $\uparrow$	Acc e $\downarrow$	Acc ir $\uparrow$	H a $\uparrow$
AC (Model-Based)	24.6	26.6	2.0	28.6	24.4	-4.2	26.9	25.8	-1.1
AC (Noise-Based)	25.5	26.4	0.9	29.3	23.8	-5.5	27.3	26.2	-1.1
ESD-x	23.3	28.6	5.3	28.5	24.7	-3.8	26.5	25.7	-0.8
ESD-u	21.8	27.8	6.0	27.8	24.5	-3.3	25.7	25.3	-0.4
MACE	23.2	27.5	4.3	26.6	25.8	-0.8	25.1	2 8.5	3 .4

Appendix

Table 8: Evaluation on erasing the specific category : Entity ( e.g., church), Artistic Style ( e.g., Claude Monet) and Abstraction ( e.g., color). CLIP classification accuracy are reported for each category on the erased category itself ( Acc e , efficacy), and the remaining unaffected categories ( Acc ir , specificity). The combined performance is summarized by H a = Acc ir

-

Acc e , where higher values indicate a better balance between erasure and preservation.

Category	# Number	Mapping Concept	Celebrities
Erasure Group	50	‘ A person ’	‘ Adele ’, ‘ Albert Camus ’, ‘ Angelina Jolie ’, ‘ Arnold Schwarzenegger ’, ‘ Audrey Hepburn ’, ‘ Barack Obama ’, ‘ Beyoncé ’, ‘ Brad Pitt ’, ‘ Bruce Lee ’, ‘ Chris Evans ’, ‘ Christiano Ronaldo ’, ‘ David Beckham ’, ‘ Dr Dre ’, ‘ Drake ’, ‘ Elizabeth Taylor ’, ‘ Eminem ’, ‘ Elon Musk ’, ‘ Emma Watson ’, ‘ Frida Kahlo ’, ‘ Hugh Jackman ’, ‘ Hillary Clinton ’, ‘ Isaac Newton ’, ‘ Jay-Z ’, ‘ Justin Bieber ’, ‘ John Lennon ’, ‘ Keanu Reeves ’, ‘ Leonardo Dicaprio ’, ‘ Mariah Carey ’, ‘ Madonna ’, ‘ Marlon Brando ’, ‘ Mahatma Gandhi ’, ‘ Mark Zuckerberg ’, ‘ Michael Jordan ’, ‘ Muhammad Ali ’, ‘ Nancy Pelosi ’,‘ Neil Armstrong ’, ‘ Nelson Mandela ’, ‘ Oprah Winfrey ’, ‘ Rihanna ’, ‘ Roger Federer ’, ‘ Robert De Niro ’, ‘ Ryan Gosling ’, ‘ Scarlett Johansson ’, ‘ Stan Lee ’, ‘ Tiger Woods ’, ‘ Timothee Chalamet ’, ‘ Taylor Swift ’, ‘ Tom Hardy ’, ‘ William Shakespeare ’, ‘ Zac Efron ’
Retention Group	50	-	‘ Angela Merkel ’, ‘ Albert Einstein ’, ‘ Al Pacino ’, ‘ Batman ’, ‘ Babe Ruth Jr ’, ‘ Ben Affleck ’, ‘ Bette Midler ’, ‘ Benedict Cumberbatch ’, ‘ Bruce Willis ’, ‘ Bruno Mars ’, ‘ Donald Trump ’, ‘ Doraemon ’, ‘ Denzel Washington ’, ‘ Ed Sheeran ’, ‘ Emmanuel Macron ’, ‘ Elvis Presley ’, ‘ Gal Gadot ’, ‘ George Clooney ’, ‘ Goku ’,‘ Jake Gyllenhaal ’, ‘ Johnny Depp ’, ‘ Karl Marx ’, ‘ Kanye West ’, ‘ Kim Jong Un ’, ‘ Kim Kardashian ’, ‘ Kung Fu Panda ’, ‘ Lionel Messi ’, ‘ Lady Gaga ’, ‘ Martin Luther King Jr. ’, ‘ Matthew McConaughey ’, ‘ Morgan Freeman ’, ‘ Monkey D. Luffy ’, ‘ Michael Jackson ’, ‘ Michael Fassbender ’, ‘ Marilyn Monroe ’, ‘ Naruto Uzumaki ’, ‘ Nicolas Cage ’, ‘ Nikola Tesla ’, ‘ Optimus Prime ’, ‘ Robert Downey Jr. ’, ‘ Saitama ’, ‘ Serena Williams ’, ‘ Snow White ’, ‘ Superman ’, ‘ The Hulk ’, ‘ Tom Cruise ’, ‘ Vladimir Putin ’, ‘ Warren Buffett ’, ‘ Will Smith ’, ‘ Wonderwoman ’

Appendix

Table 9: Complete list of celebrities used in our experiment.

Figure 10: User study interface on Erasing Cleanliness and Irrelevant Preservation.

Figure 11: User study interface on Image Quality, Prompt Adherence and Output Diversity.

	Nudity			Entity			Artistic Style			Abstraction
Method	Total $\downarrow$	FID $\downarrow$	CLIP $\uparrow$	Acc e $\downarrow$	Acc ir $\uparrow$	H a $\uparrow$	Acc e $\downarrow$	Acc ir $\uparrow$	H a $\uparrow$	Acc e $\downarrow$	Acc ir $\uparrow$	H a $\uparrow$
DiT-Localization ( Zarei et al., 2025 )	306	36.7	26.3	16.8	23.4	6.6	30.7	23.1	- 7.6	27.9	26.4	-1.5
EraseAnything ( Gao et al., 2025a )	2 31	2 4.0	2 9.9	24.6	2 7.5	2.9	25.6	2 6.4	0.8	2 5.3	28.9	3 .6
Ours	145	23.6	30.8	2 4.3	27.9	3 .6	2 6.2	27.0	0.8	24.9	2 8.8	3.9
HunyuanImage-3.0	628	23.1	31.2	28.6	31.5	-	32.1	28.3	-	29.6	29.8	-

Appendix

Table 10: Additional quantitative results on HunyuanImage-3.0 model.

Figure 12: Visualizations on erasing the concept “Pablo Picasso” . With the retention prompt kept as “painting” in LoRA training, our method preserves the overall painting modality while effectively removing Picasso’s signature abstract-geometric style. The generated results remain high-quality and faithful, exhibiting minimal semantic or compositional shift.

Figure 13: Visualizations on erasing the concept “Vincent van Gogh” . With the retention prompt kept as “realistic” in LoRA training, our method generates realistic scenarios while effectively removing van Gogh’s signature abstract style. The generated results remain high-quality and faithful, exhibiting minimal semantic or compositional shift.

Figure 14: Visualizations of ablation study on $\varepsilon$ .

Figure 15: Visual ablation of different fine-tuning configurations. We compare various parameter update strategies to demonstrate the necessity of our method. “Top K” and “Top 1” layers are selected based on the highest attention scores for the target concept, following the logic of DiT-Localization ( Zarei et al., 2025 ) . The results reveal critical insights: fine-tuning the FFN weights has negligible impact on erasure; meanwhile, any standard LoRA fine-tuning on shared projections (even when restricted to a single layer) leads to immediate generation collapse and severe visual artifacts. Only our Stream Disentangled Concept Erasure Framework (bottom row) achieves high-quality erasure while preserving the structural integrity of the image. Therefore, our Stream Disentangled Concept Erasure Framework is a prerequisite for stable erasure in pure single-stream models.

Figure 16: Comparison with mainstream concept erasing methods. Here, ESD, AC and EraseAnything are adapted with our Stream Disentangled Concept Erasure Framework to prevent generation collapse.

Figure 17: Visualization on LoRA Disentanglement. The left side of the blue dashed line delineates the original image ( yellow box) and the erased image by our Z-Erase. The right side illustrates the result on unrelated concepts upon incorporating the LoRA associated with the erased concept. Top rows: Celebrity ; Mid rows: Entity ; Last rows: Art .

Concept Erasure

Concept erasure removes specified knowledge or attributes from models while preserving their ability to perform unrelated tasks.

See all 61 papers

Erasing Thousands of Concepts: Towards Scalable and Practical Concept Erasure for Text-to-Image Diffusion Models
Apr 12, 2026Hoigi Seo, Byung Hyun Lee, Jaehyun Cho +2Safe T2I Generation Diffusion Model Unlearning
Dept. of Electrical and Computer Engineering, Seoul National University, Republic of Korea · IPAI, Seoul National University, Republic of Korea · INMC, Seoul National University, Republic of Korea
EraseAnything++: Enabling Concept Erasure in Rectified Flow Transformers Leveraging Multi-Object Optimization
Mar 1, 2026Zhaoxin Fan, Nanxiang Jiang, Daiheng Gao +2Concept Erasure Diffusion Model Unlearning
Beijing Advanced Innovation Center for Future Blockchain and Privacy Computing, School of Artificial Intelligence, Beihang University. · University of Science and Technology of China.
TEA: Text Encoder Alignment for Robust Concept Erasure in Text-to-Image Models
Aug 15, 2026Alireza Dehghanpour Farashah, Zhuan Shi, Negar Rostamzadeh +1Safe T2I Generation Adversarial Robustness
Mila – Quebec AI Institute · McGill University · Google Research

Apr 12, 2026J/K move · Enter open · S save

Erasing Thousands of Concepts: Towards Scalable and Practical Concept Erasure for Text-to-Image Diffusion Models

Hoigi Seo, Byung Hyun Lee, Jaehyun Cho, Sungjin Lim+1

Dept. of Electrical and Computer Engineering, Seoul National University, Republic of Korea · IPAI, Seoul National University, Republic of Korea · INMC, Seoul National University, Republic of Korea

Safe T2I Generation Diffusion Model Unlearning T2I Diffusion Models

Open paper AskPDF arXiv

Large-scale text-to-image (T2I) diffusion models deliver remarkable visual fidelity but pose safety risks due to their capacity to reproduce undesirable content, such as copyrighted ones. Concept erasure has emerged as a mitigation strategy, yet existing approaches struggle to balance scalability, precision, and robustness, which restricts their applicability to erasing only a few hundred concepts. To address these limitations, we present Erasing Thousands of Concepts (ETC), a scalable framework capable of erasing thousands of concepts while preserving generation quality. Our method first models low-rank concept distributions via a Student's t-distribution Mixture Model (tMM). It enables pin-point erasure of target concepts via affine optimal transport while preserving others by anchoring the boundaries of target concept distributions without pre-defined anchor concepts. We then train a Mixture-of-Experts (MoE)-based module, termed MoEraser, which removes target embeddings while preserving the anchor embeddings. By injecting noise into the text embedding projector and fine-tuning MoEraser for recovery, our framework achieves robustness to white-box attack such as module removal. Extensive experiments on over 2,000 concepts across heterogeneous domains and diffusion models demerate state-of-the-art scalability and precision in large-scale concept erasure.

Abstract

Concept Erasure

Explore similar work

Erasing Thousands of Concepts: Towards Scalable and Practical Concept Erasure for Text-to-Image Diffusion Models

EraseAnything++: Enabling Concept Erasure in Rectified Flow Transformers Leveraging Multi-Object Optimization

TEA: Text Encoder Alignment for Robust Concept Erasure in Text-to-Image Models