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

2026-08-15

AI Summary

This study addresses the poor robustness, quality degradation, and high computational overhead associated with concept erasure in text-to-image generation. We propose TEA, a lightweight Text Encoder Alignment framework that pioneers a domain alignment paradigm by fine-tuning only the text encoder. Utilizing a token-level discriminator to render concept and safety prompt representations indistinguishable, TEA achieves zero inference overhead while keeping the generative backbone frozen. As the first robust erasure method extended to Rectified Flow Transformers, TEA demonstrates cross-architectural effectiveness. It achieves state-of-the-art adversarial robustness while preserving generation quality on Stable Diffusion v1.4 and records the lowest attack success rate on SD v3.5, validating its efficacy across diverse model architectures.

Abstract

Text-to-image diffusion models can be misused to generate harmful content through adversarial or paraphrased prompts that bypass built-in safety mechanisms. Existing concept erasure methods often suffer from limited robustness against adversarial prompts, degradation of benign generation quality, or reliance on inference-time interventions that introduce persistent computational overhead. To address these limitations, we formulate concept erasure as a domain alignment problem in the text representation space. We propose a lightweight Text Encoder Alignment framework (TEA) that fine-tunes only the text encoder while keeping the generative backbone fully frozen. Given concept--anchor prompt pairs, our method trains a discriminator to distinguish token-level representations of concept-containing prompts from those of safe anchor prompts, while updating the text encoder to make these representations indistinguishable. TEA introduces zero inference-time overhead and requires only a small number of fine-tuning steps, making it highly efficient to deploy at scale. Despite this efficiency, TEA achieves state-of-the-art erasure robustness against black-box and white-box adversarial attacks on Stable Diffusion v1.4, while preserving generation quality on benign prompts. Furthermore, TEA is model-agnostic and achieves the lowest attack success rate on Stable Diffusion v3.5, extending concept erasure to a Rectified Flow Transformer architecture with T5 conditioning where prior methods remain largely unexplored. Code is available at \href{https://github.com/alirezafarashah/TEA.git}{https://github.com/alirezafarashah/TEA.git}

Problem

• Concept Erasure
• Text-to-Image Diffusion Models
• Adversarial Robustness
• Safety Mechanisms
• Generation Quality

Innovation

• Text Encoder Alignment
• Concept Erasure
• Domain Alignment
• Zero Inference Overhead
• Adversarial Robustness