TEA: Text Encoder Alignment for Robust Concept Erasure in Text-to-Image Models
Organizations: Mila – Quebec AI Institute · McGill University · Google Research
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 https://github.com/alirezafarashah/TEA.git
Figures & tables
Appendix figures & tables14 assets
Supplementary material from the paper’s appendix.
Appendix
| Aligned positions | Attack Success Rate (%) | COCO-Caption | ||||
|---|---|---|---|---|---|---|
| MMA | Ring-A-Bell | P4D | Avg. | FID | CLIP | |
| SD-v1.4 | 93.20 | 98.13 | 86.40 | 92.58 | – | 31.51 |
| Concept spans only | 28.30 | 0.00 | 6.99 | 11.76 | 42.89 | 30.88 |
| All tokens ( TEA ) | 1.20 | 0.00 | 2.21 | 1.14 | 46.50 | 30.65 |
| Attack Success Rate (%) | COCO-Caption | ||||||
| MMA | Ring-A-Bell | P4D | Avg. | FID | CLIP | ||
| – | – | 93.20 | 98.13 | 86.40 | 92.58 | – | 31.51 |
| 0.5 | 0.5 | 1.90 | 0.00 | 1.84 | 1.25 | 48.13 | 30.34 |
| 0.5 | 1.0 | 0.50 | 0.00 | 0.74 | 0.41 | 52.65 | 30.01 |
| 1.0 | 0.5 | 1.20 | 0.00 | 2.21 | 1.14 | 46.50 | 30.65 |
| 1.0 | 1.0 | 1.90 | 0.00 | 1.47 | 1.12 | 46.41 | 30.54 |
| Method | Attack Success Rate (%) | |||
|---|---|---|---|---|
| MMA | Ring-A-Bell | P4D | Avg. | |
| FLUX.1-dev | 67.90 | 72.15 | 60.53 | 66.86 |
| DES | 13.00 | 65.82 | 17.54 | 32.12 |
| Ours ( TEA ) | 12.20 | 48.10 | 14.04 | 24.78 |
| Benchmark | Method | NudeNet (%) | LlavaGuard (%) |
|---|---|---|---|
| MMA | FLUX.1-dev | 67.90 | 11.50 |
| Ours ( TEA ) | 12.20 | 0.30 | |
| Ring-A-Bell | FLUX.1-dev | 72.15 | 46.84 |
| Ours ( TEA ) | 48.10 | 16.46 | |
| P4D | FLUX.1-dev | 60.53 | 21.93 |
| Ours ( TEA ) | 14.04 | 0.88 |
| Anchor set | Attack Success Rate (%) | COCO-Caption | ||||
|---|---|---|---|---|---|---|
| MMA | Ring-A-Bell | P4D | Avg. | FID | CLIP | |
| SD-v1.4 | 93.20 | 98.13 | 86.40 | 92.58 | – | 31.51 |
| Original anchors ( TEA ) | 1.20 | 0.00 | 2.21 | 1.14 | 46.50 | 30.65 |
| Deletion anchors | 4.00 | 0.93 | 2.57 | 2.50 | 45.36 | 30.65 |
| Method | Training Time | Inference Overhead |
|---|---|---|
| SPM | 1.2 h | 1.05 |
| ESD | 0.7 h | 1 |
| SLD | 0 | 1.5 |
| UCE | 2 s | 1.05 |
| SAFREE | 0 | 1.5 |
| GLoCE | 120 s | 1.2 |
| Concept Prompt | Anchor Prompts |
|---|---|
| topless woman in garden | 1. woman in floral summer dress in garden |
| 2. lady enjoying flowers in a long robe | |
| 3. portrait of a woman picking flowers | |
| 4. woman wearing a sunhat and scarf in garden | |
| 5. casual outdoor photo of a woman in spring attire | |
| female nude ballet dancer | 1. ballet dancer in leotard stretching |
Safe T2I Generation
Safe text-to-image generation develops methods to prevent harmful images while preserving useful, faithful outputs.
See all 51 papers