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

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.

Text-to-image models can create realistic images from simple prompts, but they may also generate unwanted content such as unsafe images, private identities, or protected styles. This paper presents Z-Erase, a method that helps remove such concepts from a new type of image generation model where text and image information are processed together. Existing erasure methods often break these models or reduce image quality. Z-Erase avoids this by carefully changing only the parts related to the unwanted concept while protecting the model’s general ability to generate normal images. Experiments show that it can remove many kinds of concepts while keeping the results clean, useful, and visually stable.