Identifying Crucial Attention Heads for Multilingual Language Models: Retrieval and Retrieval-Transition Heads
TL;DR
A cross-lingual needle-in-the-haystack task is designed, to identify RTHs that retrieve key information and further map it to target-language output and advances the understanding of multilingual LMs by isolating the attention heads responsible for mapping to target languages.
Abstract
Retrieval heads, a subset of attention heads in Transformers, were studied in English, showing its crucial role in retrieving information from the context. We expand the study of retrieval heads to multilingual context and find that while nearly half of all retrieval heads are often shared across multiple languages, language-specific retrieval heads also emerge. We further design a cross-lingual needle-in-the-haystack task, to identify $\textit{Retrieval Transition Heads (RTH)}$ that retrieve key information and further map it to target-language output. Our experiments reveal that RTHs do not always overlap with RH, and can be more vital for Chain-of-Thought reasoning in multilingual LLMs. Across four multilingual benchmarks (MMLU-ProX, MGSM, MLQA, and XQuAD) and two model families (Qwen2.5-7B Instruct and Llama3.1-8B Instruct), we demonstrate that masking RTH induces bigger performance drop than masking Retrieval Heads (RH). Specifically, for Llama3.1-8B Instruct, masking the top-25 RTHs results in an average 37.3-point drop (against 19.5-point drop from RH masking) in reasoning accuracy (MMLU-ProX, MGSM) and 9.0 F1-score reduction (against 5.1 drop for RH) in extractive QA (MLQA, XQuAD). Our work advances the understanding of multilingual LMs by isolating the attention heads responsible for mapping to target languages.
