Scaling head attribution reveals general-purpose attention roles in translation
Scaling Attention Head Analysis via Gradient-Based Attribution in Context-Aware Machine Translation
Computation and LanguageArtificial Intelligence
Summary
Understanding how different parts of AI language models help translate words correctly can be hard. The authors developed a new way to trace which pieces of the model (called attention heads) influence decisions by using gradients, a technique that shows what matters most during learning. They tested this on models translating between several languages and found that some attention heads play a general helper role, improving accuracy across many tasks. Interestingly, just looking at which heads pay the most attention isn’t enough to know their importance, hinting that models create extra backup roles during training.
What this means in practice
- •For machine translation developers: Identify important attention heads that improve translation accuracy across contexts to optimize and debug translation models.
- •For natural language processing engineers: Analyze and enhance large language models by pinpointing attention mechanisms that contribute importantly to various language tasks beyond translation.
Authors
Paweł Mąka, Yusuf Can Semerci, Jan Scholtes, Gerasimos Spanakis
Abstract
In this paper, we introduce a gradient-based head attribution strategy where the Token-level Max-Margin loss is backpropagated to the attention maps. This framework enables a large-scale causal analysis of attention heads, making it suitable for LLMs. We evaluate our method on the task of disambiguation in Context-aware Machine Translation, where we analyze 50 phenomena across 4 models and 4 language directions. We empirically show the alignment of our method with the effects of increasing the attention scores of token-to-token relations on three models and two language directions, ensuring the robustness of our method. Our analysis reveals the presence of the "general-purpose" attention heads that improve the model's performance when attending to different relations. We find that the average attention a head assigns to a relation does not necessarily relate to the model's performance, which suggests that the models developed redundancies during training in terms of the head functions.