MechLight: Mechanistic Highlighting in LMs
- MechLight is a mechanistic interpretability–inspired method that assigns token-level importance scores based on attention head logit contrasts.
- It differentiates between context and memory usage by identifying the head that steers the model toward a specific answer, leading to precise highlight explanations.
- Compared to FA, IG, and ATTN, MechLight offers a computationally efficient and structurally grounded approach to diagnosing context attribution in transformer models.
Searching arXiv for the target paper to ground the article in the primary source. arXiv search query: (Sun et al., 3 Oct 2025) MechLight is a mechanistic-interpretability–inspired highlight explanation method for context utilisation in LLMs. In the formulation introduced in "Evaluation Framework for Highlight Explanations of Context Utilisation in LLMs" (Sun et al., 3 Oct 2025), it is not a generic mechanistic interpretability analysis pipeline; it is a concrete procedure that starts from head-level mechanistic attribution, identifies attention heads that “steer” the model towards context-based or memory-based answers, and converts that head-level attribution into token-level highlight scores that can be evaluated alongside Feature Ablation (FA), Integrated Gradients (IG), and Attention (ATTN). Its central aim is to reveal which context tokens, via which attention head, actually influenced the chosen answer, rather than only measuring generic saliency.
1. Conceptual role and problem setting
MechLight is introduced as “a mechanistic interpretability MI–inspired method where we propose to convert the MI insights (e.g., the attention head most important for context utilisation) to HEs.” In this setting, context utilisation denotes the ability of LLMs (LMs) to incorporate relevant information from the provided context when generating responses, while highlight explanations (HEs) are token-level explanations intended to indicate the exact context pieces and tokens that influenced model outputs.
The method is designed for a task with three closely related objectives: determining whether the model used context or parametric memory, determining which context passage was used when more than one is present, and identifying the exact answer tokens within the relevant passage. MechLight therefore functions as a full HE method: it produces one scalar importance per input token, so it can rank context versus question tokens, tokens belonging to versus , and individual answer spans.
A common misconception is that MechLight is a general-purpose mechanistic interpretability framework. The paper explicitly positions it more narrowly. It is an adaptation and operationalisation of existing mechanistic interpretability ideas, following Yu et al. (2023), under the assumption that some attention heads “steer” the model towards context-based or memory-based answers. This suggests that its contribution lies less in discovering new internal mechanisms than in translating head-level mechanistic signals into human-readable, evaluable token highlights.
2. Mechanism and formal definition
For one instance, MechLight takes as input a causal LLM with standard Transformer architecture, an input sequence , the model’s generated answer token , and a pair of competing candidate answers. In single-context regimes, and the candidate pair is ; in dual-context regimes, and the candidate pair is . Its output is a token-level explanation vector , with one score per input token.
The procedure begins with per-head logit contributions at the generation step. For each layer 0 and head 1, the paper defines
2
where 3 is the contribution of head 4 to the residual stream after the output projection. The contribution of that head to the logit of token 5 is then
6
MechLight does not use these head-level scores in isolation. Its crucial operation is a contrast between competing answers: 7 where 8. In single-context setups, 9 measures how much more the head supports the context answer than the memory answer. In dual-context setups, 0 and 1 quantify preference for one passage’s answer over the other’s.
Head selection is instance-adaptive. For instances where the model answered from context in the single-context case,
2
For memory-answer instances,
3
For dual-context setups where the answer comes from 4,
5
The mechanistic assumption is that there exists at least one head whose logit contribution strongly favours the used source; the attention pattern of that head is then treated as an informative explanation of context usage.
Once the head is chosen, MechLight converts head attention into token highlights: 6 The token score is therefore the attention weight from the answer position to token 7 under the selected head. In the paper’s own formulation, the output is directly interpretable as where the most context-steering head looked when producing the answer (Sun et al., 3 Oct 2025).
3. Relation to alternative highlight explanation methods
MechLight is evaluated against three established HE methods: Feature Ablation, Integrated Gradients, and Attention. The comparison is important because MechLight shares some surface similarities with attention-based explanations while differing in the criterion used to choose the relevant head.
Feature Ablation is perturbation-based. For each token 8,
9
where 0 is the logit of answer 1. It measures global impact on answer probability by masking tokens. The paper characterises it as very faithful but computationally expensive, especially on long contexts.
Integrated Gradients is gradient-based. For each token,
2
It aggregates the gradient along a path from a baseline of <pad> tokens to the input. The paper notes that it is classically strong on faithfulness benchmarks, but here often fails on context usage.
Attention selects a single head in the last layer based on its contribution to the answer logit: 3 Like MechLight, it ultimately uses attention weights as token scores, but it does not contrast context against memory or one context against another during head selection.
The key conceptual difference is therefore not that MechLight uses attention and the others do not, but that MechLight’s head selection is mechanistically grounded in a logit contrast 4. ATTN picks the head that generally maximises contribution to the chosen answer in the final layer; MechLight picks the head, from any layer, that most discriminates between competing answers. Compared to FA and IG, MechLight is described as cheaper than FA on long contexts and more structurally aligned with the model’s attention computation than IG. This suggests that its distinctive contribution is the explicit encoding of context utilisation at the feature-selection stage.
4. Evaluation framework and experimental regimes
The paper introduces the first gold standard HE evaluation framework for context attribution using controlled test cases with known ground-truth context usage (Sun et al., 3 Oct 2025). It evaluates four HE methods—FA, IG, ATTN, and MechLight—across four context scenarios, four datasets, and five LMs.
The four input regimes are as follows.
| Regime | Configuration | Ground truth |
|---|---|---|
| Conflicting (single) | One context passage 5 whose answer contradicts the model’s parametric knowledge | Whether the model answered from context or parametric memory |
| Irrelevant (single) | One context passage 6 with an irrelevant answer token | Correct answers must rely on parametric memory |
| Double-Conflicting (dual) | Two passages 7, both contradict PK but with different answers | Which passage contains the selected answer |
| Mixed (dual) | One irrelevant passage and one conflicting passage | Which passage’s answer was used |
For dual-context setups, “Swap” variants reverse the order of 8 to probe positional bias. The datasets are WorldCapital, CounterFact, Fakepedia, and ConflictQA. Their descriptions in the paper are precise: WorldCapital contains short geographical contexts with a single capital-of relation; CounterFact contains entity biography questions with edited versus irrelevant contexts; Fakepedia contains long synthetic encyclopaedic contexts of approximately 9–0 tokens; ConflictQA contains long multi-domain questions of approximately 1–2 tokens. The evaluated open LMs are GPT2-XL (1.5B), Pythia-2.8B, Pythia-6.9B, Qwen2.5-3B, and Qwen2.5-7B.
The evaluation metrics are tailored to three research questions. For document-level attribution, the framework uses 3 and the cross-group rank margin
4
with positive values meaning the segment is ranked higher in the relevant group. For dual-context setups it also defines a per-instance margin
5
For simulatability, it defines two metrics from the top-6 importance scores: Normalised Mutual Information,
7
and MDL-Bits,
8
For token-level attribution, it uses reciprocal rank and mean reciprocal rank: 9 All methods are evaluated with identical metrics, typically on top-0 tokens, with robustness checks at 1.
5. Empirical performance
The paper’s overall summary is that MechLight performs best across all context scenarios (Sun et al., 3 Oct 2025). That statement is then unpacked by research question.
For RQ1, whether the model used context or memory, MechLight has the most consistently positive rank margins 2 across datasets and models, often best or second-best. FA also performs reasonably well but is more variable and sensitive to context length, with strong results on short contexts and degraded results on long ones. IG and ATTN generally fail to distinguish context from memory, with margins near zero. In simulatability for the Conflicting setup using top-5 tokens, FA is typically strongest, reducing about 3 of answer label uncertainty, while MechLight is second-best with about 4 uncertainty reduction. In the Irrelevant setup, all methods improve, and MechLight is better than FA overall on both MDL and NMutInf.
For RQ2, which context document was used, MechLight shows positive margins across most model–dataset pairs in Double-Conflicting and Mixed settings, meaning the answer passage’s tokens are ranked higher in the group that actually uses that passage. FA is second-best but shows large negative margins on long-context datasets, particularly when the answer lies in the first piece. IG and ATTN give margins close to zero. On per-instance margins, no method achieves consistently positive margins in all conditions, but MechLight is strongest overall, best in one regime and second-best in the other, with mostly positive margins. In simulatability, MechLight achieves the best result, removing about 5 of label uncertainty, while FA is second with 6 uncertainty reduction.
For RQ3, pinpointing the exact answer tokens, MechLight has the highest MRR among the four methods in all setups, although its performance drops on the longest contexts. In the single-context Conflicting regime, median MRRs across model/dataset pairs are approximately 7 for MechLight, 8 for IG, 9 for FA, and 0 for ATTN. The paper interprets this as MechLight and IG typically placing the answer in roughly the top-3 tokens, FA in about the top-6, and ATTN substantially worse. At the same time, performance degrades for all methods with longer contexts such as Fakepedia and ConflictQA, where answer tokens often fall outside the top-10.
These results indicate that mechanistic head selection based on logit contrasts improves explanation quality relative to standard post-hoc methods, particularly for identifying which passage was used and for token-level localisation. They do not imply that the problem is solved; the same evaluation shows that absolute explanation quality remains limited in long and positionally confounded settings.
6. Positional bias, long-context failure modes, and interpretability significance
Although MechLight is best overall, the paper emphasises several limitations. The most prominent is positional bias. In dual-context settings, MechLight and ATTN tend to favour early context tokens. When the answer is in the second passage, their margins often become negative; after swapping passage order, the bias follows the position rather than the content. FA and IG exhibit the opposite tendency, often favouring later passages. The paper suggests that the attention heads selected by MechLight often have an inherent recency or locality pattern and are not purely content-driven.
Long-context degradation is the second major limitation. In Fakepedia and ConflictQA, where context lengths are approximately 1–2 tokens, MechLight’s token-level MRR drops substantially. Answer tokens are frequently not in the top-10 highlighted tokens. Document-level margins remain somewhat positive, but fine-grained localisation is poor. Appendix case studies note that MechLight sometimes highlights generic tokens such as “the” and “capital of,” entity descriptors, or question words instead of the answer token itself, especially in long or noisy contexts.
The paper also identifies structural limitations in the method. MechLight relies on the assumption that single heads, identified by per-head logit contrasts, capture context-versus-memory decisions. In reality, multiple heads and layers might form circuits, so focusing on one head may oversimplify. It only uses attention, not MLP neurons or deeper patching techniques, and may therefore miss important non-attentional mechanisms for context usage. Even where it performs relatively well, its best-case margins can remain modest, indicating fundamental challenges in providing faithful, stable context attribution when contexts are long and contain multiple distractors.
From a mechanistic interpretability perspective, MechLight occupies an intermediate position. It is grounded in the idea that specific attention heads are responsible for steering the model towards particular factual associations, and it uses per-head logit attribution to quantify that support. However, it does not perform full circuit-level interpretability: there is no activation patching across layers and no causal scrubbing. Instead, it treats each head as a quasi-atomic unit, identifies “the most context-steering head” per instance, and uses that head’s attention map as a proxy for the relevant causal path. A plausible implication is that MechLight is best understood as a bridge between mechanistic attribution and user-facing explanations, rather than as a full causal account of context utilisation.
7. Practical use and broader implications
The paper frames MechLight as useful for diagnosing whether a model is using provided context or parametric memory, checking which retrieved document actually influenced the answer, and inspecting which tokens serve as evidence. In Conflicting and Irrelevant settings, MechLight’s top-3 highlights can be fed into simple probes to predict context-knowledge versus parametric-knowledge behaviour with substantial accuracy, as reflected by NMutInf and MDL. In dual-document settings, it is the most reliable of the tested methods for telling which passage’s answer was used. At the token level, it usually ranks the true answer span near the top in short contexts.
Its implementation profile is also explicit. MechLight requires access to internal attention weights and head outputs, rather than only API-level logits, so it is best suited to open-weight models. It is computationally lighter than FA for long contexts, because FA requires one ablation forward pass per token, whereas MechLight needs a forward pass plus per-head logit calculations; the cost scales mostly with the number of heads, not the number of tokens. The method still requires engineering to access per-head outputs and compute per-head logit contributions, but the paper describes this as straightforward in most Transformer libraries.
The broader implication drawn in the paper is two-sided. On one side, the poor performance of vanilla ATTN relative to MechLight shows that simply visualising attention without head selection is not enough; mechanistic selection via per-head logit contrasts matters. On the other side, MechLight’s own limitations show that even mechanistic-interpretability–inspired HEs struggle with long and complex contexts. This suggests that future methods may need multi-head or circuit-level explanations, explicit treatment of positional bias, and deeper causal tools, while still translating internal model behaviour into user-friendly token highlights. In that sense, MechLight provides a concrete template for mechanistically grounded highlight explanations without resolving the broader problem of reliable context attribution at scale.