A Representation-Level Assessment of Bias Mitigation in Foundation Models

This paper demonstrates that bias mitigation in foundation models effectively reduces gender-occupation disparities in their embedding spaces through interpretable geometric transformations, a finding validated across both encoder-only and decoder-only architectures using a newly introduced dataset called WinoDec.

The Gist

Imagine you have two very smart, but slightly prejudiced, librarians. One librarian (let's call him BERT) reads books by looking at the whole page at once, while the other (Llama) reads books sentence by sentence, from left to right.

Both librarians have learned from a massive library of human writing. Unfortunately, because human history has stereotypes (like "nurses are usually women" or "firefighters are usually men"), these librarians have started to believe those stereotypes too. If you ask them to guess a job based on a gender, they might guess "nurse" for a woman and "firefighter" for a man, even if that's not fair or accurate.

This paper is like an internal audit to see what happens when we try to "de-bias" these librarians. The researchers didn't just check if the librarians gave better answers at the end; they looked inside the librarians' brains to see how their understanding of words actually changed.

Here is the breakdown of their study using simple analogies:

1. The "Mental Map" (The Embedding Space)

Think of the librarian's brain as a giant, invisible 3D map.

• In this map, words that are related are placed close together.
• Normally, the word "Man" might be floating very close to "Firefighter," and "Woman" might be floating very close to "Nurse."
• The distance between them represents how strongly the model associates them. If they are too close, the model is biased.

2. The Experiment: Moving the Furniture

The researchers took two versions of these librarians:

• The "Before" Version: The standard, biased model.
• The "After" Version: A model that went through a special training process (using techniques like counterfactual data or human feedback) to unlearn those stereotypes.

They then asked a simple question: "Did the furniture move?"

They measured the distance between gender words (Man/Woman) and job words (Plumber/HR) in that mental map.

• The Result: Yes, the furniture moved!
• In the "Before" version, "Man" and "Plumber" were best friends (very close).
• In the "After" version, "Man" and "Plumber" moved apart, and "Woman" and "Plumber" moved closer to the center. The map became more balanced. The gap between "Man" and "Woman" regarding these jobs shrank.

3. The Two Types of Librarians

The researchers tested two different types of models to see if the "de-biasing" worked the same way for both:

• The "All-Seeing" Librarian (Encoder-only/BERT): This one sees the whole sentence at once.
• The "One-Step-At-A-Time" Librarian (Decoder-only/Llama): This one predicts the next word based only on what came before.

The Surprise: Even though these two librarians think very differently, the "de-biasing" training moved the furniture in their mental maps in almost the exact same way. This is a huge deal because it means we can use the same "fairness" rules for different types of AI, and we can trust that the internal changes are consistent.

4. The New Tool: "WinoDec"

The researchers realized that the "One-Step-At-A-Time" librarian (Llama) is tricky to test with old methods. It's like trying to test a person's memory by asking them to remember a sentence they haven't finished reading yet.

So, they built a new test called WinoDec.

• The Analogy: Imagine a game of "Simon Says" with sentences. Instead of just saying "The man is a firefighter," they created pairs like: "The firefighter is a man. The man is a firefighter."
• This forces the AI to look at the connection between the job and the gender from both directions, ensuring the test is fair and accurate for this specific type of model. They released this new test kit to the public so others can use it.

5. Why This Matters

Before this study, we mostly checked if AI was fair by looking at its final answers (e.g., "Did it hire the right person?"). But that's like judging a chef only by the taste of the soup, without knowing if they used fresh ingredients or old ones.

This paper proves that checking the ingredients (the internal map) is just as important.

• They found that when the AI becomes fairer, its internal "mental map" actually becomes more neutral.
• This gives us a way to see fairness happening inside the machine, not just guess it from the outside.
• It acts as a "truth detector" to ensure that when we tell an AI to be fair, it actually changes its mind, rather than just pretending to be fair.

The Bottom Line

The researchers successfully showed that when we fix bias in AI, we aren't just patching the surface; we are actually rearranging the AI's internal understanding of the world. The "Man" and "Woman" words are no longer glued to specific jobs in the AI's brain. They have been gently nudged apart, creating a more balanced and fair mental map for everyone.

Technical Summary

1. Problem Statement

Foundation models (FMs) often inherit and amplify societal biases present in their training data, particularly regarding gender and occupation. While various bias-mitigation techniques exist (pre-processing, in-processing, post-processing), their effectiveness is typically measured at the output level (e.g., classification accuracy or generation quality).

There is a significant gap in understanding how these mitigation strategies alter the model's internal mechanics. Specifically, it is unclear:

• Whether successful bias mitigation results in interpretable, geometric changes within the model's embedding space.
• Whether these internal representational shifts are consistent across different model architectures, specifically comparing encoder-only (bidirectional context) and decoder-only (autoregressive) models.

2. Methodology

The authors conducted a comparative analysis of baseline and bias-mitigated variants of two distinct foundation model families: BERT-large (encoder-only) and Llama2-7B (decoder-only).

A. Models and Data

• Encoder-Only (BERT):
  • Baseline: BERT-large fine-tuned on a real-world job shortlisting dataset (predicting candidate selection based on job descriptions and profiles).
  • Mitigated: A variant of BERT-large trained with Counterfactual Data Augmentation (CDA) targeting gender bias.
  • Evaluation Data: Real-world job samples (HR, Plumber) and a synthetic dataset of 40 sentence pairs varying only by gender terms.
• Decoder-Only (Llama2):
  • Baseline: Llama2-7B.
  • Mitigated: Llama2-7B-Chat (aligned via Reinforcement Learning with Human Feedback for safety).
  • Evaluation Data: The authors introduced WinoDec, a new dataset of 4,000 sentence pairs designed specifically for decoder-only models. Unlike static datasets, WinoDec uses structured pairs (e.g., "The firefighter is a man. The man is a firefighter.") to create bidirectional mutual influence between gender and occupation terms, overcoming the unidirectional nature of decoder attention.

B. Experimental Design

• Metric: Cosine Similarity was used to measure the association strength between gender terms (e.g., "man," "woman") and occupation terms (e.g., "plumber," "nurse") within the final hidden layer embeddings.
• Analysis:
  • Extracted contextual embeddings for key terms from the final hidden layer.
  • Visualized distributions of similarity scores using Kernel Density Estimation (KDE).
  • Applied the Kolmogorov-Smirnov (KS) test to statistically assess whether the distributions of similarity scores differed significantly between baseline and mitigated models, and between gender groups.

3. Key Contributions

1. Representation-Level Validation: The study provides empirical evidence that successful bias mitigation corresponds to measurable, interpretable geometric shifts in the embedding space, reducing disparities in gender-occupation associations.
2. Cross-Architecture Consistency: It demonstrates that bias mitigation induces similar internal representational changes in both encoder-only and decoder-only models, suggesting that fairness improvements are not architecture-specific.
3. WinoDec Dataset: The authors released WinoDec, a dataset of 4,000 sequences specifically designed to evaluate gender-occupation associations in decoder-only models by accounting for their masked self-attention mechanisms.

4. Key Results

Encoder-Only Models (BERT)

• Baseline Bias: The baseline BERT model showed significant gender bias. For example, "HR" was strongly associated with "female," and "Plumber" showed unexpected but significant bias patterns in the embedding space.
• Mitigation Impact:
  • HR: The disparity in cosine similarity between "male/HR" and "female/HR" decreased significantly (KS D-statistic dropped from 0.2600 to 0.1290).
  • Plumber: The mitigation strategy primarily shifted the embeddings of "male" to reduce the disparity, resulting in non-significant differences between male/female associations with "Plumber" post-mitigation.
  • Synthetic Data: Visualizations showed a convergence of distributions for male and female terms across stereotypical occupations, indicating a more neutral embedding space.

Decoder-Only Models (Llama2)

• Baseline Bias: The baseline Llama2 exhibited distinct gendered associations. In the "Gender 2 – Occupation 2" configuration, female terms were significantly more aligned with stereotypically female occupations (KS D = 0.2580).
• Mitigation Impact:
  • Convergence: Post-mitigation, the distributions of similarity scores for male and female terms converged, reducing the disparity (KS D dropped to 0.1510).
  • Distribution Shape: Mitigation reduced the "multi-modal" nature of the distributions, suggesting that gender terms no longer clustered into subgroups associated with specific roles.
  • Statistical Significance: Significant shifts were observed in the embedding distributions for both male and female terms when comparing baseline vs. mitigated models, confirming that the RLHF process altered the internal representations.

5. Significance and Conclusion

• Interpretability: The paper establishes that embedding analysis is a powerful tool for auditing foundation models. It moves beyond "black box" output metrics to provide a grounded explanation for why mitigation works: it corrects the relative positioning of terms in the vector space.
• Transparency: By demonstrating that bias reduction manifests as geometric transformations, the study offers an additional layer of validation for debiasing methods, enhancing trust in AI systems.
• Early Detection: Analyzing internal representations allows for the detection of residual or hidden biases that might not yet be visible in output-level metrics.
• Future Directions: The authors suggest extending this analysis to other bias dimensions (race, age) and correlating embedding shifts with downstream task performance to ensure that internal fairness translates to practical application fairness.

In summary, the paper confirms that bias mitigation strategies effectively reshape the internal semantic landscape of foundation models, making them more equitable and providing a new framework for auditing model fairness.