The Emergence of Relevance Through Axiomatic Attention Patterns During LoRA Fine-Tuning

AI summary

Overview

Research area: Natural Language Processing, specifically neural information retrieval and the mechanistic interpretability of LoRA fine-tuning in LLM-based rerankers.

Technical level: Advanced. The paper assumes familiarity with transformer attention, Low-Rank Adaptation, IR ranking metrics, and axiomatic IR theory.

Scope: A localized ablation and attention-analysis study identifying which attention components of RankLLaMA-7B are necessary and sufficient for reranking gains from LoRA, and whether those regions coincide with learned axiomatic IR attention patterns.

What This Paper Is About

LoRA fine-tuning is a standard way to adapt large language models into rerankers, but it is unclear where inside the network task-specific relevance behavior is learned and what attention-level changes accompany that learning. The authors localize LoRA attention updates in RankLLaMA-7B through head-, layer-, and window-level ablations, then measure how attention mass shifts toward interpretable token-pair features tied to classical IR axioms, and finally correlate the two. The goal is an interpretable, correlational account of how relevance-oriented behavior emerges during LoRA fine-tuning.

Key Contributions

1. A hierarchical keep/omit ablation framework at three granularities (single head, single layer, and sliding windows of layers) that separately identifies attention components that are sufficient versus necessary for reranking performance, given LoRA-fine-tuned MLPs throughout the network.
2. Normalized Feature Attention, a metric that excludes attention-sink mass (empirically the first token in the sequence) to reliably measure how much attention a head allocates to interpretable token-pair features.
3. A head-level and layer-level measurement of how LoRA fine-tuning changes attention to lexical matching, rarity sensitivity, and document-query interaction, plus several compositional features that combine these conditions.
4. A joint correlation analysis (Spearman) showing that regions where LoRA attention updates most affect NDCG overlap with regions where fine-tuning increases attention to axiomatic IR features.

Main Findings

• Mid-network attention updates carry most of the gain. Given LoRA fine-tuned MLPs throughout the network, restricting LoRA attention updates to layers 10–18 is sufficient to recover more than half of the performance gained by applying LoRA to all attention layers, and omitting attention fine-tuning in this region hurts performance more than elsewhere.

• A compact critical region plus one outlier. Omit ablations consistently weakened performance most in layers 10–18, with layer 29 emerging as an isolated critical site. Omitting the 13–18 window yielded the largest NDCG drop, and windows in 10–18 improved NDCG by over 0.4 points when kept in isolation.

• Redundancy across the network. Keep ablations showed heads in layers 10–18 and 25–32 contribute most in isolation, while keeping fine-tuning produced roughly uniform per-layer gains other than notably higher gains at layers 14 and 29. Regions that are sufficient but not necessary suggest redundancy.

• Baseline performance. For layer- and window-wise ablations (50 queries × 100 documents), the base model scored 0.199 NDCG and −0.204 Mean Score Margin, the fine-tuned model scored 0.911 and 8.768, for gains of 0.712 and 8.972. For head-wise ablations (50 queries × 10 documents), the base model scored 0.437 NDCG and −0.106 Mean Score Margin, the fine-tuned model scored 0.970 and 8.866, for gains of 0.533 and 8.972.

• Feature attention shifts after fine-tuning. Attention to lexical matching generally decreases, particularly from layer 11 onward, except in the final layer, which the authors attribute to self-attention collapsing toward identity. Rarity sensitivity increases in layers 8–19 and decreases from layer 20 onward, with high per-head variation despite consistent layer-wise trends. Document-query interaction shows the clearest learned signal, increasing throughout layers 8–22.

• Attention changes correlate with performance. Over sliding windows of size 6, Spearman correlations between per-window ablation effect on NDCG and per-window change in normalized feature attention were 0.37 (keep) and −0.63 (omit) for lexical matching, 0.92 and −0.89 for rarity sensitivity, and 0.71 and −0.68 for document-query interaction.

• Compositional features are equally or more predictive. Rare document tokens attending rare lexical match query tokens scored 0.88 and −0.94; document tokens attending lexical match query tokens scored 0.92 and −0.92; document tokens attending rare tokens scored 0.92 and −0.88. Correlations held at window sizes 3 and 4, though they were slightly lower for smaller windows.

• The pattern replicates on a second objective. On the document variant of RankLLaMA-7B, critical LoRA attention updates concentrated in approximately layers 7–16, keeping only this region recovered over 70% of the fully LoRA fine-tuned model's NDCG gain over the base model, and document-query interaction again aligned with the performance-critical region.

Methodology in Plain English

The authors start from RankLLaMA-7B, a decoder-only reranker built from meta-llama/Llama-2-7b-hf with LoRA adapters (rank 32, alpha 64) injected into attention and MLP weights of all 32 layers, with the base model frozen. The MLP always keeps its LoRA updates so that any measured differences can be attributed to attention alone.

They run two kinds of ablation. Keep ablations turn on LoRA attention only for the component being tested and revert everything else to base model weights, showing what is sufficient. Omit ablations keep LoRA everywhere except the tested component, showing what is necessary. Reverting to base weights (rather than zeroing or averaging) isolates the value of fine-tuning rather than the damage of destroying a component. Ablations operate at the level of individual heads, individual layers, and sliding windows of layers.

Evaluation uses 50 queries randomly sampled from the MS MARCO Dev split. Per-head ablations use 10 candidate documents per query (one relevant, nine random non-relevant); layer and window ablations use 100 candidates (one relevant, 99 random non-relevant). Each variant reranks the full set, and performance is measured with NDCG and a Mean Score Margin (the mean score gap between the relevant document and each non-relevant document), which stays sensitive when rankings do not change.

For the attention side, the authors define token-pair features from axiomatic IR: lexical matching (identical words under lowercased, whitespace-trimmed comparison), rarity sensitivity (words whose IDF, computed over 500k MS MARCO training documents, exceeds a threshold near the IDF of the 180th most frequent word), and document-query interaction (document tokens attending query tokens, since input ordering places the query first). They also test compositional features that require several conditions at once. To control for attention sinks, they introduce Normalized Feature Attention, which excludes pairs attending the sink token and normalizes over remaining attention mass. They compute per-head averages over 1,000 random relevant query-document pairs and report the signed change between the fine-tuned and base models. Finally, they compute Spearman correlations between per-window ablation effects and per-window feature-attention changes.

Why This Matters

The work connects two literatures that had not been directly joined: mechanistic interpretability of neural rerankers and the practical question of which LoRA parameters actually matter. Prior interpretability work showed that rerankers encode IDF-like and term-frequency-like signals, and prior LoRA work compared coarse all-attention versus all-MLP conditions. This paper localizes attention updates to specific layers and shows those layers are also where axiomatic attention patterns emerge, offering a structured account of where relevance is installed during adaptation.

Real-world applications:

• Search and retrieval pipelines: targeted attention fine-tuning could reduce adaptation cost for LLM-based rerankers in production multi-stage retrieval.
• Retrieval-augmented generation: cheaper or more focused reranker adaptation improves the quality of retrieved context feeding downstream generators.
• Domain-specific search such as legal, medical, or enterprise document search, where fine-tuning budgets are constrained but ranking quality matters.
• Model debugging and auditing: attention to rarity and cross-segment matching gives engineers an interpretable signal for whether a reranker is learning genuine relevance behavior.

Industry relevance centers on compute and engineering cost. The paper's final experiment configuration required 7B-parameter models and approximately 25 hours of compute on a single NVIDIA A100 GPU (Appendix H reports an anonymized code release). If only a compact mid-network region needs attention updates, practitioners can shrink the trainable attention footprint.

Future Directions

• Move from correlation to causation: the authors state that attention trends alone do not constitute a causal relevance mechanism and call for causal intervention methods.
• Address semantic matching: lexical comparison cannot capture matches between non-identical tokens, and the paper notes that empirical review of attention heatmaps suggests dedicated semantic matching heads are learned.
• Generalize beyond one model and one dataset: experiments use RankLLaMA-7B on MS MARCO with small candidate sets and random negatives, so other architectures and retrieval settings remain untested.
• Explore harder evaluation and more features: the authors call for harder negatives, additional features, and multiple models to build a complete account of how relevance is learned during LoRA.

Target Audience

Researchers and graduate students in information retrieval and NLP interpretability, especially those studying neural reranking, LoRA adaptation, or attention head specialization. It also suits machine learning engineers who fine-tune LLM rerankers and want to know which attention layers are worth training, and mechanistically minded practitioners interested in the link between classical IR axioms and learned attention behavior. Some familiarity with attention mechanisms, LoRA, and NDCG is needed to follow the experiments closely.

Authors’ abstract

LoRA fine-tuning is standard for adapting LLMs to reranking, but it remains unclear where in the network task-specific relevance behavior is learned and what attention-level changes accompany that learning. Through ablation and attention experiments, we identify where LoRA attention updates to RankLLaMA improve performance and whether those gains coincide with interpretable relevance-oriented attention patterns such as lexical matching, rarity sensitivity, and query-document interaction. We find that given LoRA fine-tuned MLPs throughout the network, restricting LoRA attention updates to a compact mid-network region is sufficient for recovering over half of the performance gained by applying LoRA to all attention layers, and that omitting attention fine-tuning in this region hurts performance more than elsewhere in the network. Additionally, we show that regions where applying LoRA affects performance the most overlap with regions where fine-tuning increased attention to axiomatic IR features. Rarity sensitivity, document-query interaction, and several compositional features are highly correlated with gains in ranking performance. Our results support an interpretable, correlational account of how relevance-oriented behavior emerges during LoRA fine-tuning and point toward improved strategies for adapting rerankers.