Mind Your Manners: Detoxifying Language Models via Attention Head Intervention

Authors: jordan  pettyjohn
nathaniel  hudson
mansi  sakavardia
aswathy  ajith
kyle  chard
Venue: EMNLP_2024_Workshop_BlackBoxNLP
Published: 2024-08-20 00:59:07

Abstract: Transformer-based Language Models have advanced natural language processing with their ability to generate fluent text. However, these models exhibit and amplify toxicity and bias learned from training data---posing new ethical challenges. This work builds upon the AttentionLens framework to allow for scalable decoding of attention mechanism information. We then use this decoded information to implement a pipeline to localize and remove toxic memories from pre-trained language models in a way that is both human interpretable and effective while retaining model performance.

Main paper: 
 introduction: 
 As Language Models (LMs) become widespread, their potential for both helpful and harmful applications increases. Recent studies have exposed the inherit biases in LMs (Bender et al., 2021) . For example, LM-based chatbots perpetuate racial biases based on the names present in a prompt and their association with certain racial groups (Schulz, 2024) . These concerns motivate our research on better understanding how LMs interpret prompts and ultimately generate toxic text.We focus on understanding how individual attention heads contribute to the generation of toxic text. Our work is driven by the observation that attention heads can have highly specialized roles in text completion tasks (Sakarvadia et al., 2023; Hanna et al., 2024; Wang et al., 2022; Nanda et al., 2023) . While many prior works have attempted to interpret attention head functions, we notice few are focused on interpreting where and how biases are encoded within these models and their harmful effects.To address this challenge, we propose a two-step toxicity analysis and intervention pipeline:(i) Degenerate Attention Response Tracking (DART) and (ii) TOXicity INtervention (TOXIN).DART uses the AttentionLens framework (Sakarvadia et al., 2023) to train custom lenses for interpreting the contribution of each attention head of the open-source GPT-2 Small model to the next predicted token. These decoded attention heads can then be used to better identify which attention heads are most prone to contribute toxic tokens while completing the prompt, and what those toxic tokens are, over a sample of 500 prompts. Our TOXIN step can then form a memory (Sakarvadia et al., 2024) and remove it from the "toxic" attention heads during inference thus reducing toxicity in LM-generated text. 
methodology and implementation: 
 DART, seen in Figure 2 , efficiently identifies toxic attention heads by comparing the top-k output logits of a trained AttentionLens with a predefined toxic dictionary (Patch et al., 2020; Grad and Orthrus-Lexicon, 2021) , outputs these problematic tokens, and visualizes the distribution of toxicity across layers and attention heads as shown in TOXIN then uses these tokens that DART identifies from the most toxic attention heads and forms a toxic memory by taking k toxic tokens from the most toxic attention head, reshaping the tokens to form a R k×|V | matrix, and reducing it to a onehot vector B ∈ R |V | . Our toxic memory is then m = BW T U where W T U is the unembedding matrix. This memory m is scaled by the injection strength τ and subtracted from attention head h ℓ,j , where ℓ denotes the layer and j denotes head position: h ℓ,j = h ℓ,j -τ m. 
experiments and results: 
 To enable our DART+TOXIN pipeline, we must first train AttentionLens for GPT-2 Small; we use a subset of the BookCorpus (Zhu et al., 2015) dataset for this training. For DART analysis, we form a dictionary using Orthrus Toxic Dictionary (Grad and Orthrus-Lexicon, 2021) and LDNOOBW (Patch et al., 2020) . For TOXIN analysis we use the Wiki Toxic Dataset (Sorensen et al., 2017) . For experiments measuring perplexity we use WikiText Corpus (Merity et al., 2016) .To quantitatively measure the toxicity of LM text generation before and after interventions, we use Toxic BERT (Hanu and Unitary team, 2020). In Figure 3 we see the results of targeting the most toxic attention heads-specifically L8-H2 (Layer 8, Head 2), L9-H3, L10-H4 and L11-H3-by removing their respective toxic memory at varying injection strengths. In many of our experiments we find that a well-tuned injection strength can both decrease toxicity while only modestly affecting perplexity such as the injection on L10-H4 with a strength of 0.25 resulting in a substantial (∼29.1%) reduction in measured toxicity with only a modest (5.2%) increase in model perplexity.These promising results demonstrate tremendous potential for future work which will focus on expanding our framework to larger models (e.g., Llama2 (Touvron et al., 2023) , Llama3 (Dubey et al., 2024) ), expanding our toxicity dataset, and considering toxicity mitigation interventions on multiple attention heads at the same time. Eventually, we will deploy our toxicity mitigation workflow in an user study to assess its real world potential for safer LM-based text generation. 
conclusion: 
 We have enhanced the Attention Lens framework to identify and address toxicity and bias in transformer-based language models. We trained large lenses and developed new pipelines to identify degenerate attention heads, generate and remove toxic memories for specific heads, and measure the impact of this excision on toxicity reduction and language modeling capabilities. By removing identified toxic memories, we achieved a targeted reduction in model toxicity with only modest reduction in model performance. These advancements significantly improve the AttentionLens framework's applicability and effectiveness in mitigating harmful biases in language models. 
 
 Appendix:  


Decision: This ongoing work focuses on understanding how individual attention heads contribute to the generation of toxic text; an appropriate topic for the workshop.