TokenErase: Robust Concept Erasure via Visual-Injected Token Optimization

Concept erasure is an essential task in text-to-image generation, aiming to prevent diffusion models from producing copyrighted or unsafe content. However, existing approaches often suffer from limited robustness and are vulnerable to adversarial attacks, which can cause erased concepts to reappear. While existing approaches enhance robustness by leveraging visual information, most rely on tuning-based techniques that project image features back into text space or on pre-trained adapters that fuse visual cues through additional cross-attention layers, both of which substantially increase training complexity. To address these challenges, we propose TokenErase, a framework that employs a Visual-Injected Self-Attention mechanism to seamlessly integrate image information into the erasure process in a training-free manner. Building upon the training-free visual injection, TokenErase introduces a Token-Optimizing Cross-Attention mechanism for lightweight concept-erasure fine-tuning, where the $U$-Net is frozen and only the text token is optimized to achieve precise concept removal without degrading generation quality. Furthermore, our learnable prompt mechanism enables TokenErase to serve as a plug-and-play module, compatible with and capable of enhancing existing concept erasure methods. Extensive experiments across multiple benchmarks demonstrate that TokenErase achieves superior robustness while minimizing degradation of the model's general-purpose generation performance. Our code is available at https://github.com/xszz666/TokenErase.