💻 computer science

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

The paper proposes TEA, a lightweight and model-agnostic framework that achieves robust concept erasure in text-to-image models by fine-tuning the text encoder to align concept-containing prompts with safe anchors, thereby eliminating adversarial vulnerabilities without inference-time overhead or degradation of benign generation quality.

Original authors: Alireza Dehghanpour Farashah, Zhuan Shi, Negar Rostamzadeh, Golnoosh Farnadi

Published 2026-08-18

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Alireza Dehghanpour Farashah, Zhuan Shi, Negar Rostamzadeh, Golnoosh Farnadi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a world where computers can paint pictures from nothing but a sentence. You type "a cat sitting on a moonlit roof," and a machine, trained on millions of images and captions, creates a stunning, realistic scene in seconds. This technology, known as a text-to-image model, has opened doors for artists and designers, turning imagination into visual reality. However, this same power carries a risk. Just as a skilled forger can mimic a style, a clever user can trick these computers into creating harmful or inappropriate images by using sneaky, rephrased instructions that bypass the safety filters built into the system. The computer sees the words, but it misses the intent, generating content that was never meant to be seen.

For years, scientists have tried to fix this by teaching the computer to forget specific bad ideas, a process called concept erasure. Some methods try to scrub the bad data out before the computer learns, while others try to stop the bad images from appearing only at the moment they are being made. But these solutions often have a catch: they are either too slow to be useful on a large scale, they make the computer worse at drawing good pictures, or they require constant, heavy computing power every time someone asks for an image. The goal has always been to find a way to permanently remove the ability to create harmful content without slowing the machine down or ruining its artistic talent.

A team of researchers has now proposed a new approach that tackles this problem by changing how the computer understands language, rather than how it paints. They call their method TEA, which stands for Text Encoder Alignment. In these image-making systems, there are two main parts working together. The first part is the text encoder, which acts like a translator, turning your written sentence into a set of numbers that the computer can understand. The second part is the generative backbone, which takes those numbers and actually draws the picture. Previous attempts to stop bad images often involved tweaking the drawing part, which is like trying to fix a leak in a house by painting over the walls; it is messy and often ineffective because the problem lies in how the instructions were translated in the first place.

The researchers realized that the key to stopping harmful images lies in the translator. They found that the computer's ability to recognize a specific bad concept, such as nudity, is concentrated in the text encoder. To fix this, they created a training process that only adjusts this translator, leaving the drawing part completely untouched. They started with a pair of sentences: one that contained the harmful idea, like "a nude woman in a forest," and another that was almost identical but safe, like "a woman in a forest." They then taught the translator to make the numbers for the bad sentence look exactly the same as the numbers for the safe sentence.

To do this, they used a clever game of cat and mouse. They set up a small computer program, a discriminator, whose job was to tell the difference between the numbers coming from the bad sentence and the safe one. The main translator tried to fool this discriminator by changing its own settings so that the two sentences produced identical numbers. At the same time, the researchers made sure the translator didn't forget how to handle normal, safe sentences. They did this by checking that the safe sentences still looked like they did before the changes were made. This ensured that the computer would still be able to draw beautiful, harmless pictures of forests and women, just without the ability to draw the specific harmful version.

The results of this method were striking. When tested against a variety of tricky attempts to trick the computer, including those designed by experts to bypass safety systems, this new approach stopped the harmful images almost completely. In tests where other methods failed to stop the bad content or made the computer worse at drawing good pictures, this method succeeded in both. It worked on older versions of the technology and also on the newest, most advanced systems, proving that the idea holds up across different types of image generators. Perhaps most importantly, because the drawing part of the computer was never changed, the system did not slow down at all. It took the same amount of time to create an image as it did before, with no extra computing power needed every time a user asked for a picture.

The researchers also tested whether this method could remove other types of unwanted concepts, such as specific artistic styles. They found that it could successfully stop the computer from copying the style of specific artists, including Van Gogh and Kelly McKernan, when asked, while still allowing it to paint in other styles. This suggests that the method is not just a one-trick pony for removing nudity, but a flexible tool that can be used to remove any specific idea the user does not want the computer to generate. By focusing on the translation step and using a simple, efficient training process, the team has shown that it is possible to make these powerful tools safer without sacrificing their speed or their ability to create art.

This work represents a significant step forward in making artificial intelligence safer and more reliable. It moves away from the idea of constantly monitoring and blocking outputs, and instead builds safety directly into the way the computer understands language. The method is lightweight, requiring only a short period of training, and it leaves the rest of the system exactly as it was. For anyone relying on these tools to create content, this means a future where the computer can be trusted to follow instructions without accidentally crossing the line into harmful territory, all while keeping its creative spark intact. The researchers have made their code available, allowing others to test and build upon this approach, ensuring that the path to safer image generation remains open for further discovery.

Technical Summary: TEA (Text Encoder Alignment for Robust Concept Erasure)

Problem Statement

Text-to-image diffusion models are susceptible to misuse, where adversarial or paraphrased prompts can bypass built-in safety mechanisms to generate harmful content (e.g., deepfakes, NSFW imagery). Existing concept erasure methods face significant trade-offs:

Training-free methods (e.g., inference-time projections) avoid upfront costs but introduce persistent per-query computational overhead and lack robustness against strong adversarial attacks.
Training-based methods that fine-tune the generative backbone (e.g., U-Net) are computationally expensive, risk degrading generation quality on benign prompts, and struggle with selective erasure due to the broad distribution of concept-related parameters across network layers.
Existing text-encoder approaches often require large-scale datasets, codebook memory, or computationally heavy adversarial training that compromises generation quality.

The core challenge is to achieve robust concept erasure against both black-box and white-box attacks while preserving generation quality on benign prompts and maintaining zero inference-time overhead.

Methodology: Text Encoder Alignment (TEA)

The authors formulate concept erasure as a domain alignment problem within the text representation space. The method relies on the decomposition of text-to-image models into a text encoder ( $\tau$ ) and a frozen generative backbone ( $G$ ).

1. Theoretical Formulation
The paper posits that suppressing a target concept $c$ is equivalent to aligning the representation distribution of concept-bearing inputs ( $P_c$ ) with that of concept-free "anchor" inputs ( $P_{\neg c}$ ).

Risk Bound: Based on H-divergence theory, the concept risk on concept-bearing inputs is bounded by the concept risk on concept-free inputs plus the divergence between the two representation distributions.
Objective: By minimizing the divergence (making representations indistinguishable) while preserving the risk on concept-free inputs (utility preservation), the model can no longer distinguish between safe and unsafe prompts, thereby suppressing the concept.

2. Framework Architecture
TEA fine-tunes only the text encoder while keeping the generative backbone (U-Net or DiT) fully frozen. The process involves:

Prompt Pair Construction: For each concept prompt $p_c$ (containing the target concept), an anchor prompt $p_a$ is generated (e.g., via LLM) that preserves surrounding semantics but removes the target concept.
Token-Level Alignment: Unlike methods that use pooled sequence representations, TEA operates at the token level. This allows the model to align specific semantic differences between $p_c$ and $p_a$ without collapsing unrelated semantics.
Adversarial Training:
- A lightweight discriminator ( $D$ ) is trained to distinguish between token-level representations of concept prompts ( $h^c_t$ ) and anchor prompts ( $h^a_t$ ).
- The text encoder ( $\tau_\theta$ ) is updated to fool the discriminator, making $h^c_t$ indistinguishable from $h^a_t$ .
- Dynamic Anchoring: Both concept and anchor prompts are passed through the current encoder during training. This prevents the encoder from collapsing onto a fixed reference and encourages tracking a consistently evolving safe distribution.
Loss Functions:
- Adversarial Loss ( $L_{adv}$ ): Minimizes the discriminator's ability to classify concept tokens correctly.
- Preservation Loss ( $L_{pres}$ ): A cosine similarity loss that keeps the representations of anchor prompts ( $h^{a, new}_t$ ) close to those produced by the original frozen encoder ( $h^{a, orig}_t$ ). This ensures that benign semantics are not degraded.
- Total Objective: $L_{enc} = \lambda_{adv} L_{adv} + \lambda_{pres} L_{pres}$ .

Key Contributions

Formulation: The authors provide a principled formulation of concept erasure as a representation-space domain alignment problem, deriving a theoretical bound where erasure is controlled by an alignment term and a utility preservation term.
TEA Framework: They propose a lightweight, adversarial text encoder fine-tuning framework that achieves robust suppression with zero inference-time overhead and requires only a small number of fine-tuning steps.
Model Agnosticism & Novel Architectures: The method is model-agnostic and successfully transfers to modern architectures, specifically Rectified Flow Transformers (e.g., Stable Diffusion v3.5) with T5 conditioning, where prior concept erasure methods have been largely unexplored.
Robustness: The method demonstrates superior robustness against both black-box (prompt engineering) and white-box (optimization-based) adversarial attacks compared to existing baselines.

Experimental Results

The method was evaluated on Stable Diffusion v1.4 and Stable Diffusion v3.5-large for erasing explicit content (nudity) and artistic styles (Van Gogh, Kelly McKernan).

Explicit Content Erasure:
- Black-Box: TEA achieved the lowest average Attack Success Rate (ASR) across benchmarks (MMA, Ring-A-Bell, P4D), outperforming methods like AdvUnlearn and DES.
- White-Box: TEA consistently achieved the lowest ASR across UnlearnDiff, MMA, and Ring-A-Bell white-box attacks.
- Quality Preservation: TEA maintained high generation quality on benign prompts (COCO dataset), showing competitive or superior FID and CLIP scores compared to baselines. Notably, it avoided the quality degradation seen in adversarial training-based methods.
Artistic Style Erasure: TEA achieved high perceptual distance (LPIPSe) for target artists while maintaining low distance (LPIPSu) for other artists, indicating effective erasure without over-erasing unrelated styles.
Representation Analysis: Empirical measurements of Proxy A-Distance (PAD) confirmed that TEA significantly reduced the separability between concept and anchor representations across encoder layers, while maintaining separability between concept and unrelated (COCO) prompts.

Significance and Claims

The paper claims that TEA addresses the critical limitations of current concept erasure methods by offering a computationally efficient, persistent, and robust solution.

Efficiency: By freezing the generative backbone and only fine-tuning the text encoder, TEA introduces no inference-time overhead, making it suitable for large-scale deployment.
Robustness: The adversarial alignment in the representation space provides a stronger defense against adversarial prompts than training-free methods or methods that modify the generative backbone.
Scalability: The method's ability to adapt to Rectified Flow Transformers (SD v3.5) extends the applicability of concept erasure to next-generation diffusion architectures.
Ethical Stance: The authors emphasize that the work is motivated by harm reduction. While the evaluation requires explicit prompts and images to verify efficacy, these are strictly controlled, masked in figures, and used solely for defensive purposes to mitigate unsafe generation.