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

2026-03-26

AI Summary

This work proposes Z-Erase, the first concept erasure framework tailored for single-stream diffusion Transformers, addressing the challenge that existing methods often induce generation collapse and fail to effectively remove undesirable concepts. Z-Erase decouples the update mechanisms of textual and visual streams and introduces a Lagrangian-guided adaptive erasure modulation strategy, enabling precise removal of target concepts while preserving generation stability and high image quality. Theoretical analysis demonstrates that the proposed approach converges to a Pareto-stable equilibrium. Extensive experiments show that Z-Erase achieves state-of-the-art erasure performance across multiple tasks without compromising the fidelity of generated images.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: Generation

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Web Mining and Content Analysis: Large pretrained models with web data

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.

Problem

• concept erasure
• single-stream diffusion transformers
• text-to-image generation
• generation collapse

Innovation

• concept erasure
• single-stream diffusion transformer
• generation collapse
• adaptive erasure modulation
• Pareto stationary point