Source-linked AI summary
Scaling Attention Head Analysis via Gradient-Based Attribution in Context-Aware Machine Translation
Paweł Mąka, Yusuf Can Semerci, Jan Scholtes, Gerasimos Spanakis
TL;DR
Analyzing attention heads in LLMs at scale is difficult because exhaustive intervention studies require many evaluations. The paper uses gradient-based attribution to rank head-relation pairs for intervention, then finds multifunctional heads and functional redundancies across its analyses.
Problem
Exhaustive causal interventions across the many head-relation pairs in modern LLMs are computationally expensive, limiting analysis across larger models, language directions, and phenomena.
Source: §1 · p. 1Method
The method backpropagates a Token-level Max-Margin loss into attention maps to select promising head-relation pairs for Modifying Heads evaluation.
Source: §4 · p. 5Results
The analysis identifies general-purpose heads that improve performance when attending to different relations, while average attention assigned to a relation does not necessarily track performance.
Source: §1 · p. 2, §8 · p. 12Takeaways & Limitations
The findings suggest that models develop redundancies in head functions, and overlap in apparent functions is linked to transfer across language directions and phenomena.
Source: §8 · p. 12The Modifying Heads technique requires well-defined token relations and evaluation metrics that are challenging for the models; results for Formality and Inflection at C = 0.99 were not statistically significant.
Source: §6.1 · p. 10
Abstract
from arXiv · showhide
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.
1 Introduction
Analyzing LLM attention heads is costly because exhaustive causal interventions require thousands of head-relation evaluations. The paper proposes gradient attribution to scale this analysis and reports multifunctional heads, functional redundancies, and a reduction in computational requirements.
1 Introduction
- Exhaustive causal interventions require thousands of evaluations on modern LLMs, limiting analysis across larger models, language directions, and linguistic phenomena.The paper gives EuroLLM 1.7B as an example of a model requiring thousands of evaluation runs.Source: §1 · p. 1
- The framework backpropagates a Token-level Max-Margin loss into attention maps to rank promising head-relation pairs without exhaustively evaluating every pair.The authors describe it as an efficient way to approximate the effects of Modifying Heads.Source: §1 · pp. 1–2
- Focusing on top-k heads identified via gradients preserves analysis accuracy while reducing computational requirements by more than 96%.Source: §1 · p. 2
- The analysis identifies attention heads that attend to linguistically relevant relations and improve performance when attention to those relations is increased.Source: §1 · p. 2
- Some heads attend to multiple specific relations, and the authors link this overlap to model generalizability and possible transfer.Source: §1 · p. 2
- Average attention scores do not necessarily track model performance, suggesting that LLMs develop functional redundancies among heads.Source: §1 · p. 2
2 Related Work
Context-aware MT uses surrounding sentences to resolve ambiguity, but general-purpose translation metrics can miss context-dependent errors. Prior work studies contextual evaluation and attention-head functions, while intervention-based methods examine component behavior at substantial computational cost.
2 Related Work
- Decoder-only LLMs have achieved state-of-the-art results in Context-aware MT, but model scale makes their internal mechanisms challenging to analyze.Source: §2 · p. 2
- BLEU and COMET may miss context-use errors because only a few tokens often require context and standard metrics are insensitive to such mistakes.Contrastive and generative evaluation datasets have been introduced to address this issue.Source: §2 · p. 3
- Interpretability research examines attention's influence on behavior, with some studies linking specific heads to human-interpretable, linguistically plausible functions.Source: §2 · p. 3
- Gradient-based explanation work often identifies influential input parts, whereas this paper focuses on attention heads as parts of the model influencing predictions.The distinction is described as local versus global explanation.Source: §2 · p. 3
- Attention Patching replaces a head's attention map with one saved from the correct input, whereas Modifying Heads shifts attention toward a specified relation while preserving the distribution over other tokens.Modifying Heads can measure the potential positive effect of increased attention and enable targeted fine-tuning.Source: §2 · p. 3
3 Background
Section 3 introduces the attention mechanism, the Modifying Heads intervention, and contrastive evaluation as foundations for the paper’s analysis.
3 Background
- Section 3 sets out the Transformer attention, Modifying Heads intervention, and Context-aware MT contrastive evaluation used by the paper.These are presented as the concepts on which the work builds.Source: §3 · p. 3
3.1 Attention Mechanism
- Transformer attention projects inputs into Q, K, and V, uses scaled Q–K dot products and softmax scores, then weights V representations to produce each head’s output.The heads’ outputs are concatenated and linearly projected to form the attention module’s output.Source: §3.1 · p. 3
3.2 Modifying Heads
- Modifying Heads sets the total attention from tokens in X to tokens in Y to a chosen C while leaving other pre-softmax scores unchanged.The paper examines whether increasing attention to a relation (C = 0.99) benefits the task, alongside average attention and performance drops at C = 0.01 as other notions of head function.Source: §3.2 · p. 4
3.3 Contrastive Evaluation
- Contrastive evaluation tests whether a model ranks the context-dependent correct target sentence above alternatives differing by one word, with accuracy as the dataset-level score.The study uses ContraPro and LCPT for pronoun disambiguation and adapts ctxPro to contrastive evaluation across other linguistic phenomena.Source: §3.3 · p. 4
4 Approach
The approach estimates head-relation effects through gradients on attention scores, then evaluates only the highest-gradient pairs with Modifying Heads.
4 Approach
- For each dataset example, the method backpropagates a loss into each head’s attention scores and aggregates gradients for the specified token-to-token relation.For multi-token phrases and cues, it aggregates the maximum gradient over token pairs, corresponding to a head attending to any part of the cue from any part of the phrase.Source: §4 · p. 4
- Token-level Max-Margin loss uses contrastive-token probabilities to guide attention scores toward a larger relative probability difference when the correct token is wrong or the margin is below µ.For multi-token contrastive words or phrases, the loss uses the lowest-probability token.Source: §4 · p. 5
- Top-k gradient-selected head-relation pairs are tested with Modifying Heads at C = 0.99, reducing computation while retaining robust measurement despite potentially non-monotonic gradients.The gradients indicate accuracy differences but are not assumed to map linearly or monotonically to them.Source: §4 · p. 5
5 Experiments
Token-level Max-Margin loss retrieves accuracy-improving head-relation pairs better than other tested losses in most scenarios. Its aggregated gradients are monotonic with accuracy increases, while individual-example gradients can change sign.
5 Experiments
- The evaluation spans EuroLLM, Qwen, and Gemma models and multiple datasets; EuroLLM 9B uses only five selected relations because a full sweep is prohibitively costly.The five relations were selected based on accuracy increases measured with EuroLLM 1.7B.Source: §5 · p. 5
5.1 Gradient-Modifying Alignment
- Token-level Max-Margin loss achieves the best average precision in most cases, with exceptions on EuroLLM 1.7B at context size 1 and Qwen Auxiliary.Table 1 reports average precision for head-relation pairs whose modification improves accuracy beyond the unmodified model’s 99% confidence interval.Source: Table 1 · p. 6, §5.1 · p. 7
- Token-level Max-Margin loss has the highest Spearman correlation for positive accuracy differences in most scenarios, except Qwen on Auxiliary and LCPT.Source: §5.1 · pp. 6–7
- Its NDCG and head-relation ranking results are best in most scenarios; it still retrieves relevant pairs within the top 1% when surpassed.NDCG exceptions are Qwen LCPT and Auxiliary; ranking exceptions include EuroLLM 1.7B at context size 1 and Qwen LCPT and Auxiliary.Source: §5.1 · p. 7, Figure 2 · p. 7
5.2 Why Gradients Work?
- Aggregated gradients were monotonic with accuracy increases, but individual-example gradients were non-monotonic and sometimes changed sign.The analysis tested three selected heads across modification values from 0.01 to 0.99.Source: §5.2 · p. 7
6 Analysis
The analysis spans 50 scenarios across four models and four language directions while reducing computational cost by more than 96%. It examines multifunctional heads, unresponsive heads, cross-task overlap, and head tuning.
6 Analysis
- More than 96% lower computational cost enabled analysis of 50 scenarios across four models and four language directions.The authors focused evaluation on promising head-relation pairs; EuroLLM 9B was evaluated on a limited set of relations.Source: §6 · pp. 8–9
6.1 Results
- Heads can attend to or benefit from modifying multiple relations, including across language directions and phenomena.Examples include heads linked to the same relation across directions and general-purpose heads spanning different relations, directions, and phenomena.Source: §6.1 · pp. 9–10
6.2 Unresponsive Heads
- Heads with high average attention did not necessarily improve accuracy when modified, and tested increases above their average attention were not statistically significant.The authors suggest these heads may not be crucial for disambiguation because other heads attend to the same relation.Source: §6.2 · p. 10
6.3 Link to Generalizability
- Accuracy-based head overlap was high for Auxiliary across language directions, while Gender overlap was substantially lower than attention-based overlap.The authors argue that accuracy-based overlap better explains transfer patterns reported in prior work; Formality and Inflection accuracy results were not statistically significant.Source: §6.3 · pp. 10–11
6.4 Tuning Attention Heads
- Head tuning left BLEU and COMET stable, but ContraPro accuracy gains varied by head-relation pair.Tuning 19-4 TP →TC+1 almost replicated the modification improvement, while tuning 21-15 TP →SP+1 improved accuracy by 1 percentage point versus more than 3.5 when modified.Source: §6.4 · pp. 11–12, Table 2 · p. 11
7 Indirect Object Identification Task
On Indirect Object Identification, the authors test gradient attribution using a log-probability contrastive loss. The top-ranked pairs capture all changes above 0.60 while reducing computation by almost 97%.
7 Indirect Object Identification Task
- More than 57% improvement in the logit-difference metric came from modifying head-relation pair 20-14 TA → SB+1.The task had near-perfect accuracy, so the authors measured the logit difference between the correct and incorrect name instead.Source: §7 · p. 12, Figure 8 · p. 12
- Head-relation pairs with high average attention did not generally increase the metric when modified, and some heads improved it across different relations.Head 20-12 responded to five relations, including TA →SB+1 and TA →SA.Source: §7 · p. 12
- Pearson and Spearman correlations were 0.6700 and 0.6862, respectively, and average precision was 0.6505 for the log-probability contrastive loss.The top 3% of head-relation pairs included all results with a logit-difference change greater than 0.60, while reducing computational cost by almost 97%.Source: §7 · p. 12
8 Conclusions
The paper concludes that gradient-based attribution scales head analysis and identifies heads that improve performance across relations. Its appendix describes contrastive dataset construction for context-dependent phenomena.
8 Conclusions
- The gradient-based method correlates with Modifying Heads and enabled analysis at more than 96% lower computational cost across 50 phenomena, four models, and four language directions.The authors also apply the method to Indirect Object Identification as evidence of its potential generalizability.Source: §8 · p. 12
- Some heads improve performance across different relations, with overlaps linked by the authors to transfer effects and training redundancies.The authors plan to extend the analysis to other contextual tasks and larger models.Source: §8 · p. 12
A Generating Contrastive Dataset
- The contrastive ctxPro datasets replace target context-dependent phrases with phrases sampled from other examples, using four targets for Gender, five for Formality and Inflection, and eight for Auxiliary.ContraPro and LCPT were used instead of generated ctxPro contrastive versions for Gender in English-to-German and English-to-French, respectively.Source: A Generating Contrastive Dataset · p. 20
B Evaluation Details
The evaluation details specify model input preparation, hardware and inference settings, and dataset-size reporting.
B Evaluation Details
- Evaluations used batch size 2 and gradient backpropagation batch size 1, with five previous sentences as context except in the EuroLLM 1.7B Context 1 experiment.EuroLLM models used bf16 precision, while Qwen 2.5 and Gemma 3 used 8-bit precision.Source: B Evaluation Details · p. 20
- For model inputs, researchers concatenated context and current sentences on both source and target sides, tokenized the prompt with each model’s tokenizer, and used system prompts from the models’ HuggingFace model cards.Source: B Evaluation Details · p. 20
- Table 3 reports dataset sizes by example count and average number of contrastive targets, including the correct target, for each dataset and language direction.Source: Table 3 · p. 21
C Time and Memory Requirements
The appendix quantifies the evaluation cost of head-relation interventions and gradient calculations, then gives Head Tuning and evaluation settings. Full sweeps require many evaluations, while the reported top-pair method uses a small fraction of full-sweep time.
C Time and Memory Requirements
- EuroLLM 9B Instruct requires 133400 separate evaluations for the Modifying Heads technique, assuming 10 token-to-token relations.Table 4 reports model layers, heads, and evaluation counts under this assumption.Source: C Time and Memory Requirements · p. 21, Table 4 · p. 22
- Gradient calculation uses more GPU memory than evaluation despite a lower batch size, while its execution-time increase is small relative to the full head-relation sweep.The high number of evaluations makes the sweep a substantial part of total computation time.Source: C Time and Memory Requirements · p. 21, Table 5 · p. 22, Table 6 · p. 22
- 175 GPU hours versus 5824 GPU hours are required for EuroLLM 9B on ContraPro when evaluating the top 3% of head-relation pairs versus a full sweep.The top-3% method takes 3% of full-sweep time.Source: C Time and Memory Requirements · pp. 21–22
- 21 hours versus 537 GPU hours are required for Gemma 3 1B on ContraPro with the method versus a full sweep, respectively.The method takes 4% of full-sweep time.Source: C Time and Memory Requirements · p. 22
D Details of Head Tuning
- Head Tuning uses 13640 ctxPro English-to-German Gender examples, freezing all parameters except the selected layer’s Q and K projections and minimizing attention-score MSE.Training used 10 epochs, batch size 4, and 8 gradient-accumulation steps, and took less than one hour on an NVIDIA GeForce RTX 3090 24GB.Source: D Details of Head Tuning · p. 23
- The Head Tuning evaluation uses ContraPro and 10000 sampled OpenSubtitles 2018 English-to-German examples, greedy decoding, context size 1, and gold target-side context.The authors report BLEU and COMET as evaluation metrics.Source: D Details of Head Tuning · p. 23
E Details of Indirect Object Identification Experiment
The Indirect Object Identification experiment tests whether EuroLLM 1.7B Instruct prefers the correct name in two sequence templates. Its 14 targeted relations yield a large full-intervention evaluation burden compared with evaluating the top 3% of pairs.
E Details of Indirect Object Identification Experiment
- The task asks the model to prefer name [A] over [B] in BABA and ABBA templates, using 10000 generated examples.Multi-token names were removed because the original single-token constraint did not hold for EuroLLM 1.7B Instruct; near-perfect accuracy motivated measuring the names’ logit difference.Source: E Details of Indirect Object… · p. 23
- The experiment marks names in the first sequence part as SA and SB and in the second part as TA and TB, then evaluates 14 relations among these positions.A single evaluation takes 8 minutes and uses 3.1 GB peak GPU memory; gradient calculation takes 19 minutes and uses 6.8 GB.Source: E Details of Indirect Object… · p. 23
- Evaluating the top 3% of head-relation pairs takes 216 GPU hours versus 7168 GPU hours for all 53760 interventions at C = 0.99.The top-3% evaluation is 3.01% of the full-sweep time.Source: E Details of Indirect Object… · p. 23
F Extended Results
The extended results provide baseline contrastive accuracy and examine how attention and accumulated-gradient measures relate to accuracy changes after interventions. The figures and table specify model, relation, and evaluation scopes, but these passages do not report numerical outcomes.
F Extended Results
- Table 7 reports unmodified-model accuracy on contrastive datasets across language directions and phenomena, with 99% confidence intervals.The caption identifies the table’s scope but does not provide cell values here.Source: Table 7 · p. 24
- For EuroLLM 1.7B Instruct on ContraPro with context size 5, Figure 9 relates accumulated Token-level Max-Margin Loss gradients to accuracy change after setting C = 0.99.It shows all head-relation pairs and a separate view limited to pairs with ∆A > 0.Source: Figure 9 · p. 24
- Figures 10–12 plot average attention scores against accuracy increase ∆A for all tested phenomena and language directions in EuroLLM 9B, Qwen 2.5 1.5B, and Gemma 3 1B Instruct.Pairs with large negative ∆A are omitted for clarity.Source: Figure 10 · p. 24, Figure 11 · p. 26, Figure 12 · p. 26
- Table 8 reports Pearson r and Spearman ρ correlations between accumulated gradients and ∆A across all pairs, positive-∆A pairs, and nonpositive-∆A pairs, also comparing average attention scores.The table covers all models and settings; EuroLLM 9B uses a limited set of token-to-token relations.Source: Table 8 · p. 25