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

Dear Reviewer CnNd,

We deeply appreciate your thorough review and valuable feedback. Below is a summary of our answers (A) to the weaknesses (W) you raised.

---

[W1]: There seems to be some disconnect between the problem it aims to solve and two of the objectives of the preliminary analysis.

[A1]: Thank you for raising this important issue! We apologize for the lack of clarity in linking these aspects in our paper.

The connection between prompt sensitivity of LLMs and their bias towards labels is well-established in the LLM bias mitigation literature and we mentioned this in Line 21 of our paper. However, thanks to your insightful feedback, we realize the need to more thoroughly strengthen and elucidate this connection to enhance understanding of our method's effectiveness. A detailed analysis is provided below:

The Connection Established by Literature: Firstly, the literature demonstrates that the prompt sensitivity arises from LLMs' inherent bias towards predicting certain answers. Therefore, the prompt brittleness can be addressed by mitigating LLMs' bias towards labels. For example, in the classic study of "Calibrate Before Use" [1] it explicitly states that:

We demonstrate that this instability of prompt arises from the bias of language models towards predicting certain answers.

They detailed this analysis in Section 4 of their paper. Building on this foundation, following studies address the bias of LLMs toward to address prompt brittleness. For example, a domain-label bias is identified and mitigated to alleviate prompt brittleness in [2].

Additional Analysis on Our Method: We further analyze why mitigating model's bias towards labels can alleviate prompt brittleness in our method. Due to the inherent bias of LLMs, different prompts can lead to varying biases towards labels. For example, due to recency bias, placing a negative sentiment analysis sample at the end of a prompt can make LLMs tend to predict 'negative', incorrectly classifying positive samples and thus degrading ICL performance. Various bias factors lead to different direction and extend of bias, resulting in different changes in ICL performance and leading to the prompt brittleness. In contrast, our UniBias method effectively mitigates various potential biases inherent in LLMs. By doing so, it minimizes the introduction of bias towards labels regardless of the difference in prompts, leading to more stable and accurate ICL performance across different prompt configurations.

Experimental Validations: We further provide experimental results to demonstrate how mitigating bias towards labels can alleviate prompt brittleness in ICL. We analyze the average prediction probabilities of 'negative' and 'positive' across all test samples in various prompt formats on the SST2 dataset. The ICL performance corresponding to these formats is depicted in Figure 1 of our paper. As shown in the table below, the average prediction probabilities for each format are represented as $(P_{\text{ave}}(\text{neg}), P_{\text{ave}}(\text{pos}))$. Our results reveal that different prompt formats introduce varying biases, leading to different levels of label bias which, in turn, can destabilize ICL performance. Conversely, our UniBias method produces balanced probabilities across labels and significantly mitigates label bias, resulting in more stable and accurate ICL performance and reducing prompt brittleness.

We appreciate your constructive comments, which have significantly enhanced the clarity of our paper. We will make sure to incorporate the above discussions into the revised manuscript.

---

Moreover, we would like to highlight other enhancements made during the rebuttal process:

• Expanded LLM Evaluation: We conduct additional experiments to evaluate UniBias on more LLMs, including GPT-J (6B) and GPT2-XL (1.5B). Detailed results can be found in Table 1 of the attached one-page PDF.
• Exploration of Common Biased Components Across Tasks: We identify and eliminate the common biased components within LLMs and evaluate its performance on multiple tasks. Results are detailed in Table 4 of the new PDF. This experiment demonstrates the potential of our work in stimulating diverse bias mitigation methods in a novel perspective of LLM inner structure manipulation.
• Enhancement in Grid Search Process: We optimize the grid search process used in our method by providing alternatives of implementing one-time grid search and using unlabeled samples in place of labeled ones. The results of these optimizations are depicted in Table 2 and Figure 1 of the rebuttal PDF.
• In-depth Component Analysis: We further visualize and analyze the distribution of identified biased attention heads, as shown in Figure 2 in the PDF.

---

We are encouraged to find that these additioanl analysis and experiments, prompted by your valuable insights, have substantially strengthened our paper. We look forward to hearing your feedback!

---

References:

[1] Calibrate before use: Improving few-shot performance of language models. ICML 2021.

[2] Mitigating Label Biases for In-context Learning. ACL 2023.

Dear Reviewer CnNd,

Thank you for your prompt response and for recognizing the discussions presented in our rebuttal! We greatly appreciate your insightful feedback regarding the broad spectrum of prompt variations that can lead to fluctuations in model performance. We are glad to further discuss this important issue.

We completely agree with your insight about the significant room for variation across different parts of a prompt. Actually, this insight has been a key motivation for the development of our UniBias method. We believe that merely identifying and addressing a limited number of specific bias factors based on superficial observations is not optimal or entirely effective for mitigating LLM bias. Indeed, due to the significant room in variations in prompts/models/data corpus, it is nearly impossible to identify and address all bias factors from just analyzing external effects of prompt variations on model performance. Therefore, the key challenge and limitation of current bias mitigation methods lie in designing a method that can universally address all potential factors contributing to LLM bias and fluctuations ICL performance.

This challenge motivates the rationale for our UniBias method. UniBias is designed to address all potential biases, whether arising from prompt variations, data corpus discrepancies, or model variations. This is feasible because, fundamentally, all these superficial variations lead to biased behaviours internally in either the FFN vectors or the attention heads—components where nearly all LLM parameters reside. By directly identifying and mitigating biases within these FFN vectors and attention heads, UniBias provides a foundational approach to address all forms of bias, representing a poineering work in achieve this objective and a significant advancement in LLM bias mitigation.

The three types of bias factors discussed in the paper are intended to demonstrate the internal mechanisms of these biases observed superficially. Given the vast potential for variations in prompt-induced biases, as you noted, it would be impractical to discuss the internal mechanisms for every possible bias factor.

Thank you for highlighting this crucial point, which not only validates a very important strength of our method but also aligns perfectly with our motivation for developing it. We intend to include a discussion on this issue in our revised manuscript.

Dear Reviewer 8E5c,

We deeply appreciate your thorough review and valuable feedback. Below is a summary of our answers (A) to the weaknesses (W) and questions (Q) you raised.

---

[W1]: Show shared biased components across tasks.

[A1]: We greatly value your suggestion! Inspired by your insights, we not only list the common biased components, but also explore eliminating common biased components for bias mitigation.

Listing of Common Biased Components: The shared biased attention heads with their frequency of occurrence (>=3) are detailed in Table 3 of the rebuttal PDF. The shared biased FFN vectors are not discussed because they are associated with label words, which typically vary across tasks.

Eliminating Common Biased Components: Inspired by your feedback, we explore the potential of eliminating common biased components and apply it to all tasks, rather than handling each task individually. We conduct additional experiments on multiple tasks to assess the effectiveness of directly elinimate these components. Experimental results in Table 4 of the new PDF indicate that although not as effective as our full Unibias method, it outperforms the vanilla ICL by a large margin. Notably, this is a free gain in performance as it involves merely the direct masking of the biased components identified in our work and is applicable to diverse tasks.

---

[W2]: Result analysis could be expanded for Table 2.

[A2]: We deeply appreciate your suggestion and provide a detailed analysis here.

As an ablation experiment, we find that both FFN-only and attention-only methods outperform the standard ICL, demonstrating their effectiveness. However, as you suggested, there is no consistent pattern on when these two ablation methods perform better. The removal of biased FFN vectors predominantly addresses uncontextual biases associated with label words, while the removal of biased attention heads targets contextual biases, such as recency bias. Therefore, the relative effectiveness of these methods depends on the dominance of either type of bias in the given task and the specifics of the prompt examples, which can vary unpredictably.

Despite these variations, it is notable that our comprehensive UniBias method consistently yields the best performance, underscoring the effectiveness and robustness of UniBias.

---

[W3]: It would be interesting to see how bias exists across different layers of the model.

[A3]: Thanks for your insightful suggestion! We visualize the identified biased attention heads in Figure 2 of the new PDF and provide analysis below. Biased FFN vectors are not visualized due to the high dimensionality of FFN layers.

In Figure 2, the intensity of color corresponds to the frequency with which an attention head is identified as biased across 12 datasets, each repeated 5 times. The visualization reveals that biased attention heads mostly occur in the middle layers. This observation aligns with the established understanding of LLM architecture: early layers tend to encode shallow and discrete concepts such as lexical features, while higher layers are responsible for high-level semantics and complex reasoning. Therefore, discrete concepts at early layers are less likely to directly contribute to the bias. In contrast, middle layers, where simple reasoning begins to form are more likely to introduce biases (e.g., "A sentence is always positive rgardless of its content"). The higher layers are less likely to produce this kind of bias associated with simpler, shallow reasoning.

---

[W4]: Some mathematical formulations can be more detailed.

[A4]: We appreciate the valuable suggestion! The formulations for identifying biased components involve measuring the sum of label logits, the bias of label logits, and the coefficient variance of FFN vector coefficients or attention head label logits. They are corresponding to the relatedness, biased, and low variance criterions. Based on your feedback, we will include more comprehensive mathematical derivations and examples in the revised paper to clearly illustrate how these bias criteria are identified and applied.

---

We apologize for the brevity of the following responses due to space limitations.

---

[Q1]: Unibias performance

[A5]: Thank you for this important question. To address it, it's essential to consider that a more appropriate comparison for UniBias would be between standard ICL and UniBias, as this comparison contrasts ICL based on vanilla LLMs with LLMs after the elimination of biased components. In this case, our UniBias consistently improve performance across all 12 datasets we evaluated. On the other hand, calibration methods are under different methodology and perform better on a limited number of datasets. This may be attributed to the specific characteristics of each dataset.

---

[Q2]: Alternative label names for vanilla label bias

[A6]: They are generated by GPT-4.

---

[Q3]: Discrepancies in logit distribution between FFN and attention head

[A7]: Thank you for the insightful question. Generally, they are consistent across the middle and higher layers. However, we observe discrepancies in some individual components.

---

[Q4]: Organization of Sections 2 and 3.

[A8]: Thank you for the detailed suggestions! We will revise Sections 2 and 3 carefully. This includes merging Sections 3.3 and 3.4, enhancing coherence between the criteria and conditions, and strengthening connections between 3.3 and 3.4. We will also ensure a thorough review for typos and writing issues as suggested in Q5.

---

We wish to highlight that we also extend the evaluation of UniBias to include GPT-J and GPT2-XL, which are detailed in Table 1 of the new PDF.

---

We are encouraged to find that these experiments and enhancements, prompted by your valuable insights, have substantially strengthened our paper. We look forward to hearing your feedback!

Dear reviewer LH5i,

We deeply appreciate your thorough review and valuable feedback. We appreciate your acknowledgment of the significance of the problem we investigated and motivation behind our approach. Below is a summary of our answers (A) to the weaknesses (W) you raised.

---

[W1]: Grid search with 20 labeled samples per class can lead to concerns in fairness in comparison with baselines, generalizability of the method, and increased cost when the number of classes is large.

[A1]: Thank you for raising this important concern! It prompted us to further enhance our method.

Initially, we choose a limited number of labeled samples for grid search because it is typically manageable given the small size of labeled data. Additionally, it is only needed for the initial set up to identify biased LLM components. During inference stage, the computation cost of UniBias is completely identical to that of the vanilla LLMs. We believe this is a reasonable trade-off for the substantial improvements in ICL performance and the novel perspective Unibias offers on mitigating bias through LLM inner components.

However, we understand the concern that using labeled samples can be challenging when labeled data are scarce or the number of classes is large. Therefore, we have optimized our method in light of your valuable feedback:

Replacing Labeled Samples with Unlabeled Ones: To Address the potential challenge in accessing labeled samples and ensure a fair comparison with baselines, we further explore the alternative of using unlabeled samples during grid search. In our method, labeled samples are used to ensure each class is represented proportionally in the grid search, without direct use of the specific label information. Therefore, for balanced datasets, it is equally effective to employ a slight larger pool of unlabeled samples.

Our experimental findings, illustrated in Figure 1 of the rebuttal PDF, indicate that approximately $40 \times \text{number of classes}$ unlabeled samples achieves performance comparable to that obtained with labeled samples. This number is significantly lower than the quantity required by the PC baseline, which requires at least $200 \times \text{number of classes}$ unlabeled samples.

Optimization of Grid Search: Inspired by your feedback, we streamline the grid search process. We now provide an alternative to perform grid search on a single dataset to obtain a set of threshold values, which are then applied universally across other datasets. This alternative eliminates the need for repeated grid searches across each dataset, significantly enhancing efficiency and scalability of the method. By directly adopting the fixed threshold, it also addesses the scalability issues associated with an increasing number of classes.

Our experimental results in Table 2 of the rebuttal PDF confirm that using a fixed set of thresholds maintains good performance, slightly below the original UniBias method but still superior to vanilla ICL and other baseline methods. This also illustrate the robustness of our method in terms of threshold selection.

Exploration of Eliminating Common Biased Components: In response to your feedback, we have delved deeper into the potential of eliminating common biased components and apply it to all tasks, rather than handling each task individually. In Section 4.4 of our paper, we discuss the presence of shared biased LLM components across different tasks. Inspired by this finding, we conduct additional experiments to directly elinimate these common biased components and evaluate on different tasks.

Experimental results in Table 4 of the rebuttal PDF indicate that although not as effective as our full Unibias method, it outperforms the vanilla ICL by a large margin. Notably, this is a free gain in performance as it involves merely applying a direct masking to the biased components identified in our work and is applicable to diverse tasks.

Finally, following your suggestion, we evaluate the performance of fine-tuning LLMs using 20 labeled samples per class and find that their performance is significantly lower than that of UniBias. This results further validates the effectiveness of our method.

The enhancements suggested by your feedback have markedly improved the efficiency and scalability of our approach. Furthermore, the third experiment underscores the potential of our work in stimulating diverse bias mitigation methods in a novel perspective of LLM inner structure manipulation.

---

During rebuttal, we also extend the evaluation of UniBias to include GPT-J and GPT2-XL, which are detailed in Table 1 of the new PDF.

---

We would like to further highlight the contributions of our work.

• A New Insight for LLM Bias Mitigation: Unlike existing works based on external adjustments of LLM outputs, we mitigate LLM bias through manipulation of LLM internal structure. This novel perspective could potentially stimulate future research on LLM bias mitigation from inner structure manipulation, offering a new direction for the field.
• Unveil Internal Mechanisms of LLM Bias: We conduct a thorough exploration of the internal mechanisms underlying biases in LLMs, providing deep insights that go beyond superficial observations, revealing the inner causes of these biases.
• Extensive Evaluation: UniBias is evaluate across 12 datasets and 4 LLMs, consistently demonstrating its effectiveness.
• Addressing Prompt Brittleness: Our experimental results show that our method effectively mitigates prompt brittleness, a critical issue in ICL.

---

We are encouraged to find that these experiments and enhancements, prompted by your valuable insights, have substantially strengthened our paper. We look forward to hearing your feedback!

Dear Reviewer BhZE,

We deeply appreciate your thorough review and valuable feedback. Below is a summary of our answers (A) to the weaknesses (W) and questions (Q) you raised in the review.

---

[W1]: Assess how well the findings generalize across different types of LLMs and datasets.

[A1]: We greatly value your insightful suggestion and we have conducted additional experiments to validate our method.

Model Generalization: Following your suggestion, we have conducted additional experiments to evaluate more LLMs, including GPT-J (6B) and GPT2-XL (1.5B). The experimental results are detailed in Table 1 of the rebuttal PDF. Together with the Llama2-7B and Llama2-13B models evaluated in our paper, we rigorously evaluate our method across four popular LLMs of different sizes and demonstrated the effective of our method across these models.

Task Coverage: Our work is currently focus on classification tasks, consistent with the prevailing LLM bias mitigation literature. Our work evaluates extensive array of tasks compared to the exsiting studies. To provide context, current works typically focus either on mitigating bias in general classification tasks [1-2] or exclusively address bias in multiple-choice questions (MCQs) in reasoning tasks [3-4]. In contrast, our study covers 12 datasets spanning 5 distinct tasks, including both general classification tasks and MCQs in reasoning tasks. This breadth not only aligns with but surpasses the scope of tasks examined in the literature. Moving forward, we are enthusiastic about adapting our method to a wider variety of tasks as you suggested.

We appreciate your insights, which significantly strengthened our experimental validation.

---

[W2]: Grid search with a limited number of labeled samples may limit scalability and efficiency

[A2]: Thank you for raising this concern, which prompted us to further enhance our method. Although using a small set of labeled samples is typically manageable, we recognize the potential for improvement. In response, we conduct experiments to address your concern from three aspects.

Optimization of Grid Search: To streamline the grid search process, we now provide an alternative to perform grid search on a single dataset to obtain a set of threshold values, which are then applied universally across other datasets. This alternative eliminates the need for repeated grid searches across each dataset, significantly enhancing efficiency and scalability of the method.

Our experimental results in Table 2 of the rebuttal PDF confirm that using a fixed set of thresholds maintains good performance, slightly below the original UniBias method but still superior to vanilla ICL and other baseline methods. This experiment also illustrate the robustness of our method in terms of threshold selection.

Replacing Labeled Samples with Unlabeled Ones: To Address the potential challenge in accessing labeled samples, we further explore the alternative of using unlabeled samples during grid search. In our method, labeled samples are used to ensure each class is represented proportionally in the grid search, without direct use of the specific label information. Therefore, for balanced datasets, it is equally effective to employ a larger pool of unlabeled samples.

Our experimental findings, illustrated in Figure 1 of the rebuttal PDF, indicate that approximately $40 \times \text{number of classes}$ unlabeled samples can achieve performance comparable to that obtained with labeled samples.

Exploration of Eliminating Common Biased Components: In response to your feedback, we have delved deeper into the potential of eliminating common biased components and apply it to all tasks, rather than handling each task individually. In Section 4.4 of our paper, we discuss the presence of shared biased LLM components across different tasks. Inspired by this finding, we conduct additional experiments to directly elinimate these common biased components and evaluate on different tasks.

Experimental results in Table 3 of the rebuttal PDF indicate that although not as effective as our full Unibias method, it outperforms the vanilla ICL by a large margin. Notably, this is a free gain in performance as it involves merely applying a direct masking to the biased components identified in our work and is applicable to diverse tasks.

The enhancements suggested by your feedback have markedly improved the efficiency and scalability of our approach. Furthermore, the third experiment indicates the potential of our work in stimulating future bias mitigation methods in a novel perspective of LLM inner structure manipulation.

---

[Q1]: Compatibility with SFT

[A3]: Thank you for the interesting question. We would like to respond in two aspects: the performance of our method on fully fine-tuned models and the impact of SFT on biased components.

Compatibility with SFT: UniBias is fully compatible with LLMs that have undergone SFT. We applied UniBias to Llama2-7b fine-tuned on the SST2 and MNLI datasets, with following experimental results:

Impact of SFT on Biased LLM Components: Algorithmically, the processes for identifying and mitigating biases in both SFT and non-SFT models are identical. However, we observed a reduction in the coefficients of biased FFN vectors and the magnitudes of label logits of biased attention heads post-SFT in some cases, suggesting that the SFT process may suppress the biased components in LLMs.

---

[Q2]: Can grid search process be automated?

[A4]: Yes, the grid search process in our method is automated. We will release our code to facilitate reproduction and further exploration.

---

We are encouraged to find that these experiments and enhancements, prompted by your valuable insights, have substantially strengthened our paper. We look forward to hearing your feedback!

Dear Reviewers,

We greatly appreciate your thorough review, insightful feedback, and recognition of our work. In response to your valuable comments, we have diligently provided detailed explanations and conducted additional experiments to address each point raised. We are eagerly looking forward to hearing your valuable feedback!

We would also like to take this opportunity to highlight the contributions of our work:

• A New Direction for LLM Bias Mitigation: Unlike existing works based on external adjustments of LLM outputs, we mitigate LLM bias through manipulation of LLM internal structure. This novel perspective could potentially stimulate future research on LLM bias mitigation from an internal perspective, offering a new direction for the field.
• Addressing an Important Issue of Prompt Brittleness: Our experimental results show that our method effectively mitigates prompt brittleness, a critical issue in ICL.
• Unveiling Internal Mechanisms of LLM Bias: We conduct a thorough exploration of the internal mechanisms underlying biases in LLMs, providing deep insights that go beyond superficial observations, revealing the inner mechanisms of these biases.
• Extensive Evaluation: UniBias is evaluated across 12 datasets and 4 LLMs, consistently demonstrating its effectiveness.

Thank you again for your time and efforts dedicated to reviewing our paper. We welcome any further questions and are very happy to discuss them.

With Gratitude,

Authors of Paper 12087