🤖 AI

Efficient bias mitigation in T2I diffusion models using Concept Graphs

The paper introduces CO-ALIGN, a novel bias mitigation framework for Text-to-Image diffusion models that leverages concept-graph alignment within internal ontologies to significantly reduce harmful bias and semantic incoherence while preserving image quality and enhancing concept unlearning robustness.

Original authors: Mansi, Avinash Kori, Francesco Leofante

Published 2026-07-07

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Mansi, Avinash Kori, Francesco Leofante

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 you have a very talented artist who can draw anything you describe. If you say "a nurse," this artist usually draws a woman. If you say "a doctor," they almost always draw a man. This isn't because the artist is mean; it's because they learned from a massive library of old pictures and stories where these roles were mostly filled by specific genders. This is what we call bias.

The paper introduces a new method called CO-ALIGN (Concept Ontology Alignment) to fix this problem without ruining the artist's ability to draw good pictures.

Here is how it works, using simple analogies:

The Problem: Two Artists, One Mismatch

Current ways to fix this bias are like trying to fix a broken machine by only adjusting one gear.

The Text Encoder (The Translator): This part listens to your words and turns them into a "recipe." If you say "nurse," the translator currently thinks, "Oh, that means 'female nurse'."
The Denoiser (The Painter): This part takes the recipe and actually draws the picture. It has learned to expect the translator to always say "female nurse" when you say "nurse."

The Mistake of Old Methods:

If you only fix the Translator to say "nurse" is gender-neutral, the Painter gets confused. It receives a recipe it doesn't understand, and the result is a blurry, nonsensical mess.
If you only fix the Painter to ignore gender, the Translator keeps sending biased instructions. The painter tries to follow the new rules but gets pulled back by the old instructions, so the bias remains.

The Solution: CO-ALIGN (The Team Huddle)

CO-ALIGN realizes that the Translator and the Painter are a team. You can't fix one without talking to the other.

Mapping the Relationships (The Concept Graph):
Imagine the artist's brain has a giant web of sticky notes. One note says "Nurse," another says "Male Nurse," and another says "Female Nurse." In the biased model, the "Nurse" note is glued very tightly to "Female Nurse" and loosely to "Male Nurse."
CO-ALIGN maps out this entire web for both the Translator and the Painter.
Aligning the Team:
Instead of just ripping a note off the wall, CO-ALIGN gently moves the "Nurse" note so it is equally close to "Male Nurse" and "Female Nurse."
- It does this for the Translator first.
- Then, it immediately updates the Painter's web to match the new arrangement.
The Result:
Now, when you say "nurse," the Translator sends a balanced recipe, and the Painter understands it perfectly. The result is a clear, high-quality picture of a nurse that could be a man or a woman, without the image looking blurry or broken.

The Magic Side Effect: The "Neighborhood" Pull

The paper discovered a cool side effect. Imagine you have a group of friends standing in a circle. If you pull one friend (the "Nurse" concept) toward the center, their closest friends (related concepts like "Male Nurse") get pulled along with them, even if you didn't touch them directly.

The authors used this to fix a different problem: Unlearning.
Sometimes, you want the artist to forget how to draw something harmful (like nudity). If you just tell the artist to forget "Nudity," a clever trickster might ask for a "Nymphette" or "Creampie" (words that are neighbors to "Nudity" in the artist's brain) to get the same result.

CO-ALIGN uses the "Neighborhood Pull" to drag those trickster words ("Nymphette," etc.) right next to "Nudity" before you tell the artist to forget. Now, when the artist forgets "Nudity," they automatically forget the trickster words too, making the system much safer against bypasses.

What Did They Prove?

The team tested this on a popular AI model (Stable Diffusion) and found:

Fairness: It reduced bias by 30% compared to the best previous methods.
Quality: The pictures stayed sharp and realistic (improving image quality scores).
Coherence: It stopped the "blurry mess" problem. It reduced nonsensical images by 88%.
Safety: It made the model much harder to trick into drawing harmful content, even without changing the safety tools themselves.

In short, CO-ALIGN fixes bias by making sure the "brain" that understands words and the "brain" that draws pictures are on the same page, working together in harmony.

Technical Summary: CO-ALIGN – Efficient Bias Mitigation in T2I Diffusion Models using Concept Graphs

1. Problem Statement

Text-to-Image (T2I) diffusion models systematically reproduce and amplify demographic biases (gender, race, age) inherited from their training corpora (e.g., LAION-5B). These biases are not merely superficial artifacts but structural properties of the model's learned concept associations, where biased concepts share disproportionate neural substrate with stereotypical attributes.

Existing bias mitigation techniques typically intervene at only one level of the generation pipeline:

Text Encoder Only: Methods fine-tune embeddings or steer inference-time signals to remove biased directions. However, this leaves the denoiser's internal concept organization unchanged. Consequently, the denoiser receives out-of-distribution conditioning signals it cannot interpret coherently, leading to semantically incoherent outputs.
Denoiser Only: Methods edit cross-attention weights or latent spaces to reroute prompts. However, the text encoder continues to route underspecified prompts (e.g., "a nurse") toward biased attribute embeddings, partially negating the denoiser-side correction.

The core problem identified is that treating the text encoder and denoiser as independent correction targets is suboptimal. Correcting one without the other disrupts their geometric alignment, resulting in a trade-off where improved fairness leads to a sharp increase in semantic incoherence.

2. Methodology: CO-ALIGN

The authors propose CO-ALIGN (Concept Ontology Alignment), a framework that mitigates bias by aligning the internal concept ontologies of both the text encoder and the denoiser toward a target topology. The central hypothesis is that bias resides in the relationships between concepts rather than individual representations, and effective debiasing requires joint geometric harmony between components.

A. Concept Graph Extraction

CO-ALIGN constructs weighted concept graphs ( $G$ ) for both components, where nodes represent concepts and edges represent shared neural substrate strength.

Text Encoder Graph ( $G_{TE}$ ): Defined by the pairwise cosine similarity of $L_2$ -normalized mean-pooled last hidden state embeddings for a concept vocabulary $C$ .
Denoiser Graph ( $G_U$ ): Defined using "concept neuron fingerprints." The authors identify sparse subsets of parameters (concept neurons) in the cross-attention projection matrices ( $W^K, W^Q, W^V$ ) that are critical for a specific concept. The graph edges are computed as the average cosine similarity of the masked gradient profiles (fingerprints) of these neurons across layers and heads.

B. Alignment Objective

The framework aligns the current concept graph ( $A$ ) toward a target adjacency matrix ( $A_{GT}$ ) that encodes a desired topology (e.g., equal proximity of "Nurse" to "Male Nurse" and "Female Nurse").

Differentiable Ranking Loss ( $L_{rank}$ ): Operates row-wise on a supervised subset of concepts (target concepts). It uses soft ranks derived from differentiable sigmoid pairwise comparisons to align the current similarity rankings to hard rank targets derived from $A_{GT}$ .
Static Anchor Loss ( $L_{anchor}$ ): Prevents the drift of "anchor" concepts (e.g., specific attribute variants) by penalizing deviations from their pre-alignment representations.
Preservation Loss ( $L_{pres}$ ): A standard noise-prediction loss on reference pairs ensures the model does not suffer from catastrophic forgetting of concepts outside the editing scope.

The total objective is $L = \delta \cdot L_{rank} + \gamma \cdot L_{pres} + L_{anchor}$ . The alignment is performed sequentially: first on the text encoder, then on the denoiser (recalibrated to the updated geometry).

C. Emergent Property: Neighborhood Propagation

A key finding is that aligning a supervised concept pulls its semantic neighborhood toward the target topology without direct supervision. This "label propagation" effect, analogous to graph-based semi-supervised learning, allows the method to address adversarial bypass concepts (concepts semantically adjacent to a target that evade unlearning) by bringing them within the unlearning target's locality radius.

3. Key Contributions

Diagnosis of Single-Component Failure: The paper demonstrates that single-component debiasing techniques (encoder-only or denoiser-only) inevitably trade fairness for generation coherence, leaving residual bias or producing incoherent outputs.
CO-ALIGN Framework: A novel approach that extracts and jointly aligns concept graphs from both the text encoder and denoiser, preserving the geometric harmony between them.
Superior Performance: CO-ALIGN outperforms state-of-the-art methods by achieving significant bias reduction while maintaining high image quality and semantic coherence.
Robust Unlearning Application: The paper leverages the neighborhood propagation effect to significantly improve the adversarial robustness of post-hoc concept unlearning techniques, making them effective against bypass prompts without modifying the unlearning procedure itself.

4. Experimental Results

Experiments were conducted on Stable Diffusion v1.5 across four occupations (nurse, doctor, firefighter, receptionist) and three demographic axes (gender, race, age).

Fairness: CO-ALIGN achieved the lowest Fairness Discrepancy (FD) scores, showing a 30% improvement in gender fairness and 75.4% improvement in race fairness over the next best baselines.
Image Quality: The method improved image fidelity by 2.8% (CLIP-I) and reduced the Fréchet Inception Distance (FID) by 11.4 compared to the base model, outperforming all baselines.
Coherence: Crucially, CO-ALIGN reduced semantically incoherent outputs by 88% compared to the base model. In contrast, single-component methods showed a sharp increase in incoherence as fairness improved.
Unlearning Robustness: When applied as a pre-alignment step for unlearning (targeting "Nudity"), CO-ALIGN significantly reduced the success rate of adversarial bypass prompts. For example, the "Creampie" prompt, which generated nudity in 88–100% of cases with baselines, dropped to 0–32% after CO-ALIGN pre-alignment.

5. Significance and Claims

The paper claims that CO-ALIGN represents a shift from treating bias as a property of individual parameters to viewing it as an asymmetry in the model's internal concept topology. By enforcing joint alignment, the method achieves a Pareto-optimal balance between fairness and generation quality that single-component methods cannot reach.

Beyond immediate bias mitigation, the authors highlight the neighborhood propagation effect as a significant emergent property. This allows the framework to enhance the robustness of existing safety mechanisms (like concept unlearning) against adversarial attacks, suggesting that structural alignment of the model's ontology is a powerful tool for responsible AI deployment.

The authors remain modest regarding limitations, noting that the denoiser graph extraction is computationally intensive and that the method is currently validated only on Stable Diffusion v1.5, requiring adaptation for transformer-based architectures like DiT or SDXL. They also acknowledge that while CO-ALIGN reduces bias, it cannot fully eliminate the structural biases inherent in web-scraped training data.