[Paper Review] Interpreting Bias in Large Language Models: A Feature-Based Approach

Nirmalendu Prakash, L. Roy|arXiv (Cornell University)|Jun 18, 2024

Computational and Text Analysis Methods4 citations

TL;DR

This paper proposes a feature-based interpretability framework to analyze and mitigate gender bias in large language models (LLMs) like LLaMA-2-7B, LLaMA-3-8B, and Mistral-7B-v0.3. By hypothesizing the evolution of bias-related features through model layers and validating them via activation and attribution patching, the authors identify specific MLPs and attention heads responsible for bias propagation, demonstrating targeted debiasing via counterfactual fine-tuning that significantly reduces biased outputs without harming overall performance.

ABSTRACT

Large Language Models (LLMs) such as Mistral and LLaMA have showcased remarkable performance across various natural language processing (NLP) tasks. Despite their success, these models inherit social biases from the diverse datasets on which they are trained. This paper investigates the propagation of biases within LLMs through a novel feature-based analytical approach. Drawing inspiration from causal mediation analysis, we hypothesize the evolution of bias-related features and validate them using interpretability techniques like activation and attribution patching. Our contributions are threefold: (1) We introduce and empirically validate a feature-based method for bias analysis in LLMs, applied to LLaMA-2-7B, LLaMA-3-8B, and Mistral-7B-v0.3 with templates from a professions dataset. (2) We extend our method to another form of gender bias, demonstrating its generalizability. (3) We differentiate the roles of MLPs and attention heads in bias propagation and implement targeted debiasing using a counterfactual dataset. Our findings reveal the complex nature of bias in LLMs and emphasize the necessity for tailored debiasing strategies, offering a deeper understanding of bias mechanisms and pathways for effective mitigation.

Motivation & Objective

• To develop a feature-based method for interpreting how bias propagates through LLMs, moving beyond generalized component analysis.
• To validate the role of specific model components—particularly MLPs and attention heads—in generating and replicating gender bias.
• To demonstrate the effectiveness of targeted debiasing by fine-tuning only identified components using a counterfactual dataset.
• To extend the method to different forms of gender bias and sentence structures, proving its generalizability across contexts.
• To provide a white-box, interpretable framework for identifying bias origins and pathways in large-scale LLMs.

Proposed method

• The authors define abstract bias-related features (e.g., 'male', 'female', 'wife') as latent concepts that evolve through model layers, forming a 'feature map' from input to output.
• They use causal mediation-inspired hypotheses to map how features like 'woman' and 'wife' emerge from input prompts such as 'Washing the dishes is the duty of the'.
• Activation patching is applied to isolate the contribution of specific layers and components (e.g., MLPs, attention heads) by replacing activations from a reference input.
• Attribution patching is used to identify attention heads that copy or amplify bias features, with logit reversal used as a validation signal.
• Counterfactual fine-tuning is applied only to identified components (e.g., specific MLPs and attention heads) to reduce bias while preserving model performance.
• The method is validated across three state-of-the-art LLMs: LLaMA-2-7B, LLaMA-3-8B, and Mistral-7B-v0.3, using templates from a professions dataset.

Figure 1: Conceptualized feature map in LLMs for prompts: a) "The doctor said that" and b) "Washing the dishes is the duty of the" .

Experimental results

Research questions

• RQ1 Which specific model components (MLPs or attention heads) are responsible for generating or amplifying gender bias in LLMs?
• RQ2 How do bias-related features such as 'wife' or 'man' evolve through the layers of a transformer-based LLM?
• RQ3 Can targeted debiasing via counterfactual fine-tuning on identified components effectively reduce biased outputs without degrading performance?
• RQ4 Does the proposed feature-based method generalize to different forms of gender bias beyond the original professions dataset?
• RQ5 To what extent do attention heads copy or propagate bias features, and can this be quantitatively measured using attribution patching?

Key findings

• The feature map analysis revealed that the 'wife' feature is primarily derived from the combination of the 'woman' feature and grammatical structure, with the final abstraction occurring in the last token's MLP.
• Patching the last token MLP activations across layers 4 to 15 caused the token 'wife' to appear in the top 10 outputs, confirming its critical role in feature derivation.
• Copying attention heads from the 'dishes' token, particularly with k=100 and up to layer 20, were found to significantly contribute to bias propagation, as shown by logit reversal upon patching.
• Targeted debiasing using counterfactual fine-tuning on identified components reduced the probability of 'woman' and 'wife' in stereotypical contexts, indicating effective mitigation.
• The method demonstrated generalizability by successfully analyzing and reducing bias in different sentence structures and gender bias forms beyond the original dataset.
• The study observed potential compensation effects in mechanistic interpretability, where some heads may take over the role of patched heads, suggesting a need for further validation in debiased models.

Figure 3: (a) Fraction of samples where logit is reversed vs MLP layers patched (b) Fraction of samples where logit is reversed vs top K value.