Linear Predictability of Attention Heads in Large Language Models

Khalid Shaikh; Asmit Kumar Singh; Rebecca Christopher Dsouza; Shikhar Shiromani

cs.LGMar 4, 2026

Linear Predictability of Attention Heads in Large Language Models

Authors: Khalid Shaikh, Asmit Kumar Singh, Rebecca Christopher Dsouza, Shikhar Shiromani

Organizations: Georgia Institute of Technology · Santa Clara University

Abstract

Large language model (LLM) inference is increasingly bottlenecked by the Key-Value (KV) cache, yet the fine-grained structure of attention-head activations remains poorly understood. We show that pretrained Transformers exhibit a pervasive inter-head linear structure: for a given token, the Query, Key, and Value (QKV) vectors of an attention head can often be reconstructed as a linear combination of a small number of peer heads, typically within the same layer. Across Llama-3.1-8B, Falcon3-10B, OLMo-2-7B, and Qwen3-32B, just 2-5 reference heads recover many target heads with high fidelity (e.g., mean R^2 approx 0.76 for Keys on C4 with five references, and frequently R^2 > 0.85 on GSM8K). This predictability is learned rather than architectural: it is largely absent at random initialization, rises rapidly during pretraining as we track through OLMo-2 checkpoints, and is supported by a theoretical lower bound showing high mean-squared error for linear prediction at initialization. We further connect this emergence to increasing intra-layer alignment of Key projection subspaces. Finally, we exploit this redundancy for efficiency by caching only reference-head KV states and reconstructing the remaining heads on the fly via lightweight linear maps, achieving 2x KV-cache reduction with model-dependent accuracy trade-offs (4.5-5.5 percentage point average drop on Falcon3-10B and Qwen3-32B across five benchmarks, and larger drops on Llama-3.1-8B), and we find that reconstructing Keys is substantially less harmful than reconstructing Values.

Figures & tables

Panel of Figure 1 — Figure 1: Linear predictability among attention heads enables significant KV-cache compression. (Left) Schematic of approximating a target head’s (orange) Key activations using a learned linear projection from reference heads (green). (Middle) In a pretrained Llama-3.1 8B model, most heads’ Key activations are highly predictable (high $R^{2}$ ) from a few peers on the C4 dataset ( Raffel et al., 2023 ) , indicating a shared low-rank subspace. (Right) This predictability allows for substantial KV-cache compressionwith minimal performance degradation on benchmarks like TruthfulQA, MMLU-STEM, and Winogrande. The dashed line marks 2 $\times$ compression.

Figure 2: Linearity is pervasive—keys, queries, and values all become more predictable as the number of reference heads grows. Shown is the mean

R^{2}

when reconstructing key , query , and value activations on a shared 150-sequence slice of C4 for three pretrained models: (Left) LLaMA-3.1 8B (Middle) Falcon-3 10B (Right) OLMo-2 7B. Across all panels, every stream (K, Q, V) shows a monotonic rise in

R^{2}

as additional references are added, confirming that head-level representations in each model occupy a low-rank sub-space that can be captured with only a handful of heads.

Figure 3: Dominant predictors migrate from cross-layer to within-layer . For OLMo-2 7B we track, across checkpoints, where a head’s strongest linear predictors reside. (Left): percentage of heads whose single best (

R^{2}

-max) predictor lies in the same layer versus a different layer. (Right): the same percentages after aggregating each head’s five best predictors. At initialization almost all dominant links are inter-layer; by 50 k–100 k steps intra-layer links overtake and remain prevalent, indicating that pre-training progressively concentrates shared computation inside individual layers.

Figure 5: Pipeline for KV-cache compression via head-level redundancy. (1) We log key/value activations on a calibration set. (2) Pair-wise linear probes quantify how well one head predicts another. (3) A thresholded graph selects a minimal reference set that covers all targets. (4) Compact weight matrices $W_{r\!\to t}$ are trained for those links. (5) During inference, we cache only the reference heads and reconstruct the others on-the-fly, shrinking memory with minimal extra compute.

Model	Method	TruthfulQA	HellaSwag	MMLU Stem	GPQA	Winogrande	Average
Qwen3 32B	Baseline	39.16	82.70	82.39	40.62	73.16	63.61
Qwen3 32B	50% KV Pred	40.02	77.02	74.21	31.91	67.24	58.08
Falcon3 10B	Baseline	37.05	80.95	68.47	34.82	71.19	58.50
Falcon3 10B	50% KV Pred	35.00	74.23	61.90	27.67	71.42	54.04
Llama3.1 8B	Baseline	37.08	79.28	58.70	29.68	73.87	55.72
Llama3.1 8B	50% KV Pred	30.35	66.53	44.53	26.34	61.40	45.83

Table 1: Evaluation accuracy (%) across multiple benchmarks with and without 50% KV prediction.

Evaluation Domain	Task/Metric		Calibration Domain
Evaluation Domain	Task/Metric	Baseline	Maths	Coding	Psychology
			(GSM8K)	(Humaneval)	(MMLU Psyc)
MATH	GSM8k	$75.35$	$71.03$	$67.77$	$62.31$
	Math Intermediate Algebra	$10.74$	$6.97$	$7.41$	$7.64$
	Math Prealgebra	$41.90$	$36.28$	$32.60$	$34.32$
	Math Precalculus	$11.42$	$6.59$	$6.77$	$7.14$

Table 2: Model performance across different tasks for 50% key and value predictions on Llama3.1 8b instruct

Compression setting	Avg. accuracy	Drop vs. baseline
Baseline (no compression)	36.77%	—
Key-only (50% of keys predicted)	34.80%	$-$ 1.97 pp
Value-only (50% of values predicted)	32.47%	$-$ 4.30 pp
Joint K+V (50% of each predicted)	30.04%	$-$ 6.73 pp

Table 3: Average accuracy on the GSM8k under three compression settings (50% of heads reconstructed in each case). Keys are substantially more resilient to linear approximation than Values, consistent with the asymmetric subspace overlap growth observed during pretraining (§ 3 ).

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Evaluation Domain	Task/Metric		Calibration Domain
Evaluation Domain	Task/Metric	Baseline	Maths	Coding	Psychology
			(GSM8K)	(Humaneval)	(MMLU Psyc)
MATH	GSM8k	$75.35$	$74.60$	$74.60$	$74.22$
	Math Intermediate Algebra	$10.74$	$9.85$	$11.18$	$9.30$
	Math Prealgebra	$41.90$	$40.64$	$40.18$	$39.83$
	Math Precalculus	$11.42$	$9.70$	$10.99$	$10.07$

Appendix

Table 4: Model performance across different tasks for 50% key predictions on Llama3.1 8b instruct

Evaluation Domain	Task/Metric		Calibration Domain
Evaluation Domain	Task/Metric	Baseline	Maths	Coding	Psychology
			(GSM8K)	(Humaneval)	(MMLU Psyc)
MATH	Gsm8k	$75.35$	$72.78$	$69.21$	$68.61$
	Math Intermediate Algebra	$10.74$	$9.41$	$9.30$	$8.41$
	Math Prealgebra	$41.90$	$40.64$	$34.90$	$42.13$
	Math Precalculus	$11.42$	$8.24$	$7.69$	$7.87$

Appendix

Table 5: Model performance across different tasks for 50% value predictions on Llama3.1 8b instruct

Window	Proximity	% Conn.	Mean $R^{2}$
2-layer	Near ( $\leq 2$ )	53.6% $\uparrow$	0.2946 $\uparrow$
	Far ( $\geq 3$ )	46.40% $\downarrow$	0.2550 $\downarrow$
4-layer	Near ( $\leq 4$ )	53.59% $\uparrow$	0.2895 $\uparrow$
	Far ( $\geq 5$ )	46.41% $\downarrow$	0.2506 $\downarrow$

Appendix

Table 6: Local vs. distant dominant linear connections. Boldface and arrows highlight the higher density and strength of near-layer links.

Panel of Figure 6 — Figure 6: GSM8k – Distribution of linear predictability ( $R^{2}$ ) for LLaMA-3.1 8B. Top row: cumulative distribution functions (CDFs); bottom row: kernel-density estimates (KDEs). Ordering (left→right): Key, Query, Value.

Panel of Figure 7 — Figure 7: Olmo 2 7B – Distribution of linear predictability ( $R^{2}$ ) on its pre-training corpus. Top row: cumulative distribution functions (CDFs). Bottom row: kernel-density estimates (KDEs). Left→right in each row: Key, Query, Value.

Panel of Figure 8 — Figure 8: Heatmap of head-level linear predictability across transformer layers of Llama3.1 8B on C4. Each row corresponds to one attention head and each column to a model layer; the color intensity encodes the coefficient of determination $R^{2}$ when reconstructing that head’s activations from its strongest reference heads. The gradient from warm to cool colors illustrates that linear redundancy is highest in early layers and gradually diminishes in deeper layers for Llama3.1 8B. Left key states, right value states.

Explore similar work

Autonomy-of-Heads: Data-Free Sparse Attention from Frozen Query-Key Geometry
Aug 7, 2026Yehan Yang, Junyuan Shang, Yang Li +3Attention Head Analysis LLM Inference Acceleration
Institute of Computing Technology, Chinese Academy of Sciences · University of Chinese Academy of Sciences · Baidu Inc. +1
Sliding-window beats linear attention
Aug 28, 2026Alexia Jolicoeur-Martineau, Rhea Sanjay Sukthanker, Pashmina Cameron +1KV Caching LLM Inference Acceleration
Microsoft Applied Sciences Group (ASG) · Independent
Multi-Head Low-Rank Attention
Mar 2, 2026Songtao Liu, Hongwu Peng, Zhiwei Zhang +2Long-Context Language Model Inference LLM Inference Acceleration
The Pennsylvania State University · University of Connecticut · Carnegie Mellon University +1

Aug 7, 2026J/K move · Enter open · S save

Autonomy-of-Heads: Data-Free Sparse Attention from Frozen Query-Key Geometry

Yehan Yang, Junyuan Shang, Yang Li, Guanqun Zhao+2

Institute of Computing Technology, Chinese Academy of Sciences · University of Chinese Academy of Sciences · Baidu Inc.+1

Attention Head Analysis LLM Inference Acceleration Sparse Attention Attention Head Specialization

Open paper AskPDF arXiv

Long-context LLM inference is bottlenecked by quadratic attention computation and growing KV-cache costs. Existing sparse attention and KV-compression methods typically decide which tokens or heads to preserve from runtime attention scores, observation windows, calibration prompts, or learned gates, making head diagnosis input-dependent and costly to deploy. We propose Autonomy-of-Heads (AoH), a data-free method that identifies retrieval and streaming heads from the spectral geometry of query-key projections. AoH defines the kernel attention operator

M_h = W_K^{h\top}W_Q^h

and uses its effective-rank as a weight-space measure of head function: concentrated spectra indicate a small number of dominant query-key matching directions and are associated with retrieval heads, whereas diffuse spectra indicate the absence of a dominant global matching direction and are associated with streaming heads. We further derive an efficient

d_\text{head}

-dimensional computation that avoids constructing the full

d_\text{model}\times d_\text{model}

matrix. We conducted extensive experiments across models demonstrating that at 50% sparsity, AoH retains 96.5% of Full Attention performance on average while reducing prefill and decode latency by up to 41.4% and 66.0%, respectively, and KV-cache memory by 50.0% at 256K tokens.

Abstract

Explore similar work

Autonomy-of-Heads: Data-Free Sparse Attention from Frozen Query-Key Geometry

Sliding-window beats linear attention

Multi-Head Low-Rank Attention