UniBias: Unveiling and Mitigating LLM Bias through Internal Attention and FFN Manipulation

TL;DR

This work tackles prompt brittleness in in-context learning by uncovering internal biases arising from FFN vectors and attention heads in decoder-only LLMs. It introduces UniBias, an inference-only method that identifies biased components via three criteria (relatedness, bias, low variance) using a logit-lens interpretation and then masks them during inference. Across 12 NLP datasets and multiple model scales, UniBias yields consistent performance gains and substantially reduces sensitivity to prompt design, surpassing existing calibration approaches. The results suggest that internal component manipulation can robustify ICL and hint at global biased-component strategies for cross-task bias mitigation.

Abstract

Large language models (LLMs) have demonstrated impressive capabilities in various tasks using the in-context learning (ICL) paradigm. However, their effectiveness is often compromised by inherent bias, leading to prompt brittleness, i.e., sensitivity to design settings such as example selection, order, and prompt formatting. Previous studies have addressed LLM bias through external adjustment of model outputs, but the internal mechanisms that lead to such bias remain unexplored. Our work delves into these mechanisms, particularly investigating how feedforward neural networks (FFNs) and attention heads result in the bias of LLMs. By Interpreting the contribution of individual FFN vectors and attention heads, we identify the biased LLM components that skew LLMs' prediction toward specific labels. To mitigate these biases, we introduce UniBias, an inference-only method that effectively identifies and eliminates biased FFN vectors and attention heads. Extensive experiments across 12 NLP datasets demonstrate that UniBias significantly enhances ICL performance and alleviates prompt brittleness of LLMs.

Figures (9)

• Figure 1: illustrates the prompt brittleness of ICL and the effectiveness of our method in mitigating this issue. Experiments are conducted in one-shot setting, using SST2 socher-etal-2013-recursive dataset for experiments on example selection and prompt formatting and AGnews NIPS2015_250cf8b5 dataset for example order experiment due to more diverse combination of orders.
• Figure 2: Unveiling vanilla label bias by uncontextual accumulated FFN logits.
• Figure 3: The internal mechanism of the recency bias.
• Figure 4: The internal mechanism of the selection bias.
• Figure 5: The performance comparison under different numbers of ICL shots using Llama-2-7b.