ICML 2026 T2I Safety and Alignment

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 Hongwei Zheng Faguo Wu Wenjun Wu

Beihang University · University of Science and Technology of China · Shanghai University · Tencent · Beijing Academy of Blockchain and Edge Computing

Paper arXiv Code BibTeX
Z-Erase concept erasure examples across nudity, artistic style, and logo concepts
Z-Erase removes target concepts while preserving image quality. Sensitive regions are pixelated.

Motivation

Why single-stream models need a new erasure recipe

Recent T2I systems such as Z-Image process text and image tokens as one unified sequence. That design is efficient, but it removes the clean cross-attention handle used by earlier concept erasure methods.

The paper finds that naive fine-tuning of shared self-attention projections consistently damages the image pathway, producing noisy collapse instead of controlled erasure. Token-level attention suppression can work in toy cases, but small prompt perturbations quickly bypass it.

No cross-attention handle Shared projections collapse generation Prompt perturbations bypass token tricks
Token zeroing can erase localized concepts but is brittle under prompt variation
Attention intervention is revealing but brittle: changing the prompt spelling can bring the concept back.

Method

Two safeguards, one stable erasure path

01

Disentangle the stream before learning

Z-Erase freezes the visual processing route and injects low-rank updates only into textual hidden states. The resulting subspace lets erasure gradients target concepts without rewriting the image synthesis backbone.

02

Adapt erasure strength under a preservation constraint

A Lagrangian-guided modulation loop increases erasure only when preservation stays within a tolerance. The dynamic weight acts as a gatekeeper between removing a concept and keeping normal generation useful.

Single-stream transformer architecture in Z-Image
Single-stream transformer: text and image tokens share self-attention projections.
Gradient balancing between erasure and preservation
Adaptive modulation steers the erasure update when preservation would be harmed.

Results

Cleaner erasure with less collateral damage

On Z-Image Turbo, Z-Erase is evaluated across NSFW, celebrity identity, entity, artistic style, abstraction, multi-concept erasure, user study, and adversarial prompt attacks.

649 → 161 Detected nudity on I2P prompts
1086 → 324 Detected violence, best among compared methods
26.46 / 31.25 FID / CLIP on MS-COCO after nudity erasure
Ha 4.82 Best celebrity erasure balance score

The comparison captures the core trade-off: aggressive erasure often introduces artifacts, while conservative updates can leave the target concept intact.

Qualitative comparison between Z-Image Turbo, UCE, EraseAnything, and Z-Erase
Across concrete objects, styles, abstract attributes, and identities, Z-Erase avoids the artifacts and semantic drift seen in prior adaptations.

Evidence

Robustness, scale, and human preference

Multi-concept erasure examples with weighted LoRA averaging
Multiple concepts can be erased together through weighted LoRA averaging.
User study radar chart comparing erasure methods
In a five-dimensional user study, Z-Erase receives the strongest overall human preference.
Ablation study on epsilon and beta hyperparameters
The preservation tolerance controls the erasure-preservation trade-off in a monotonic, interpretable way.

Abstract

In brief

Concept erasure removes unwanted concepts from text-to-image models without retraining from scratch. Z-Erase targets the newer single-stream diffusion transformer paradigm, where text and image tokens share the same transformer pathway and existing methods can collapse generation.

The method combines a stream-disentangled LoRA update space with a Lagrangian-guided adaptive objective, then proves convergence toward a Pareto stationary point. Experiments show state-of-the-art balance across safety erasure, image quality, robustness, and human evaluation.

Citation

BibTeX

@article{jiang2026zerase,
  title={Z-Erase: Enabling Concept Erasure in Single-Stream Diffusion Transformers},
  author={Nanxiang Jiang and Zhaoxin Fan and Baisen Wang and Daiheng Gao and Junhang Cheng and Jifeng Guo and Yalan Qin and Yeying Jin and Hongwei Zheng and Faguo Wu and Wenjun Wu},
  journal={Forty-Third International Conference on Machine Learning},
  year={2026}
}