CO-ALIGN: bias mitigation in diffusion models using Concept Graphs

3 minute read

In brief

  • Text-to-Image (T2I) diffusion models perpetuate harmful social and demographic stereotypes inherited from web-scraped training datasets.
  • Existing mitigation strategies intervene either purely at the text encoder (causing text-image misalignment) or purely at the visual denoiser (causing image corruption).
  • CO-ALIGN introduces Concept Graphs to capture semantic relationships across both text and visual spaces, achieving coherent, unbiased generation efficiently.

This retrospective explainer is dated to the July 2026 preprint release. It discusses Efficient bias mitigation in T2I diffusion models using Concept Graphs by Avinash Kori and Francesco Leofante.

Prompt a modern text-to-image diffusion model with a simple professional prompt like “a photograph of a CEO” or “a portrait of a software engineer,” and over 90% of the generated images will depict light-skinned men in business suits. Prompt it with “a photograph of a nurse” or “a housekeeper,” and the outputs skew overwhelmingly toward women.

These disparities are not intentional design choices; they are reflections of historical sampling bias present in the billions of web-scraped image-caption pairs used to train foundation models.

While developers have attempted to fix this through prompt-rewriting or post-hoc negative prompting, these heuristic fixes often backfire. Rewriting prompts often produces bizarre historical anachronisms, while fine-tuning the visual denoiser without touching the text encoder leads to grotesque visual distortions.

This paper presents CO-ALIGN, a mathematically grounded debiasing framework that models concepts and their semantic dependencies as Concept Graphs spanning both the language and vision domains.

Why single-stage interventions fail

The authors demonstrate why existing interventions inevitably break down:

Comparison of Text Encoder only mitigation, Denoiser only mitigation, and CO-ALIGN using Concept Graphs.
Figure from the paper. Overview of failure modes and the CO-ALIGN solution. (Left) Intervening only on the text encoder produces inconsistent outputs and semantic hallucinations; (Middle) Intervening only on the visual denoiser via negative guidance creates visual artifacts and blurry faces; (Right) The proposed CO-ALIGN framework synchronises the text concept graph ($CG_{\text{Txt Enc}}$), ground truth graph ($CG_{\text{GT}}$), and denoiser graph ($CG_{\text{Denoiser}}$) to generate coherent, demographic-balanced outputs. Source: Avinash Kori and Francesco Leofante. Extracted from the source manuscript; click to enlarge.
  1. Text Encoder-Only Interventions (panel left): By perturbing text token embeddings to neutralize demographic bias, the model disrupts the delicate semantic bridge between the text encoder and the U-Net. For an occupational prompt like “lawyer,” the text encoder produces corrupted feature vectors, causing the denoiser to render illegible text posters or distorted abstract patterns (inconsistent outputs).
  2. Denoiser-Only Interventions (panel middle): Applying negated guidance directly during diffusion denoising pushes the reverse trajectory away from the biased mode. However, without updating the text representations, the denoiser fights the incoming cross-attention conditioning, producing washed-out skin, missing limbs, or distorted facial features (negated denoiser correction).

Unifying modalities through Concept Graphs

CO-ALIGN resolves this tension by constructing structured Concept Graphs (CG) that explicitly represent entities, demographic attributes, and their semantic co-occurrences.

A Concept Graph is defined as $\mathcal{G} = (\mathcal{V}, \mathcal{E})$, where vertices $v \in \mathcal{V}$ represent visual concepts (e.g., “stethoscope”, “scrubs”, “female”, “male”, “Asian”, “Black”) and edges $e = (u, v) \in \mathcal{E}$ represent conditional co-occurrence probabilities.

The framework defines three interconnected graph representations:

  • $CG_{\text{Txt Enc}}$: The concept hierarchy encoded in the text conditioning layer.
  • $CG_{\text{Denoiser}}$: The concept affinities expressed by the visual cross-attention maps of the U-Net denoiser.
  • $CG_{\text{GT}}$: The desired fair, unbiased target distribution (where occupational roles are uniformly decoupled from demographic attributes).

The learning objective aligns both model components simultaneously:

\[\mathcal{L}_{\text{CO-ALIGN}} = \mathcal{L}_{\text{diffusion}} + \alpha \, \mathcal{D}_{\text{graph}}\left(CG_{\text{Txt Enc}}, CG_{\text{GT}}\right) + \beta \, \mathcal{D}_{\text{graph}}\left(CG_{\text{Denoiser}}, CG_{\text{GT}}\right),\]

where $\mathcal{D}_{\text{graph}}$ measures topological and edge-weight divergence between the model’s empirical graph and the balanced target graph using graph Wasserstein distances:

\[\mathcal{D}_{\text{graph}}(\mathcal{G}_1, \mathcal{G}_2) = \inf_{\gamma \in \Pi(\mu_1, \mu_2)} \sum_{i, j} \gamma_{ij} \, d(v_i, v_j) + \lambda \sum_{(i, k), (j, l)} \gamma_{ij} \gamma_{kl} |A_{1, ik} - A_{2, jl}|.\]

Coherent and unbiased visual generation

By aligning the concept graphs across modalities:

  1. Demographic Balance: For occupational prompts spanning medical, corporate, engineering, and service professions, CO-ALIGN produces near-equal distributions across gender (50/50 balance) and racial diversity matching global demographic targets.
  2. Preservation of Visual Quality: Unlike denoiser-only clipping methods that degrade visual aesthetics, CO-ALIGN preserves fine-grained texture, lighting, and anatomical coherence (demonstrated in the right panel of the figure where male and female nurses across multiple ethnicities are rendered with crisp clinical scrubs and realistic facial features).
  3. Parameter Efficiency: Alignment is achieved by updating lightweight Low-Rank Adaptation (LoRA) matrices within cross-attention blocks, requiring less than 1% of the GPU hours of full retraining.

Paper and sources