PaperlayerView exact revision on arXiv

Source-linked AI summary

Emergent Structure in the Marginal Attention Space of Language Models

Valentino Maiorca, Walter Nelson, Francesco Locatello

arXiv:2610.03109v1cs.CL

TL;DR

This paper addresses the limited characterization of attention mechanics across language models. It studies query-position-marginalized attention and connects its token profile to Jacobian statistics, finding text-shared token profiles and model-specific, document-stable head profiles. The head profile supports offline KV-cache budgeting, with competitive eviction performance using a per-token score.

  • Problem

    How attention mechanics behave across models is less characterized than cross-model representation similarity.

    Source: §1 · p. 1

  • Method

    The paper marginalizes attention over query positions, studies token and head reductions, and connects per-token marginal attention to Jacobian statistics and their cross-model identifiability.

    Source: §1 · pp. 1–2

  • Results

    Per-token profiles are shared across models, per-head profiles are stable across documents and model-specific, and the resulting eviction method is competitive with methods that recompute budgets or train per target.

    Source: §6 · pp. 9–10

  • Takeaways & Limitations

    A per-head cache budget can be computed offline on pretraining text, leaving inference with a per-token score that can be supplied by a much smaller model.

    Source: §1 · p. 2

    The theory covers only per-token marginal influence, its link to marginal attention is empirical, and the census is mostly English with documents capped at 2048 tokens.

    Source: §6 · p. 10

Abstract

from arXiv · show

While representation similarity across independently trained language models is well-documented, how internal mechanics such as attention behave across models remains far less characterized. Inspired by this gap, we examine the structure of post-softmax attention weights by marginalizing over query positions, mapping them into a joint token-head "marginal attention space". Evaluating across 60+ diverse LLMs, we find that different properties emerge when reducing this space along its token and head axes. When reduced token-wise, marginal attention yields a text-intrinsic signal robustly conserved across models. To explain this property, we empirically connect marginal attention to the input-output Jacobian of the network, and prove theoretically that under a smoothness assumption, models with similar next-token distributions are guaranteed to have similar input-output Jacobian statistics. When reduced head-wise, it forms a model-private signature conserved across documents. Practically, this provides a natural way to estimate a per-head budget for key-value (KV) cache eviction, effectively decoupling model-specific budget allocation from text-intrinsic token scoring. On standard eviction benchmarks, a per-head budget precomputed offline on pretraining text, combined with a training-free token score, shows competitive performance with methods that recompute the budget on every document or train it per target. Code available at https://github.com/Flegyas/marginal-attention

1 INTRODUCTION

The paper shifts attention research from cross-model representation similarity to the less-characterized mechanics that produce representations. It studies two reductions of marginal attention, finding a text-shared token signal and a model-specific head signal with practical use for KV-cache budgeting.

1 INTRODUCTION

  • Marginal attention approximates input-output Jacobian statistics, and smooth models with similar next-token distributions must have similar Jacobians.For models sharing a tokenizer, closer predictions correspond to greater agreement in marginal influence and per-token attention.Source: §1 · pp. 1–2
  • Per-token marginal attention agrees across 67 models on the same document (¯ρ = 0.87), while per-head profiles match a model’s other documents (¯ρ = 0.95).These reductions distinguish a text-intrinsic token signal from a model-specific head signal.Source: §1 · pp. 1–2, Figure 1 · p. 2
  • The per-head profile is stable across documents and private to each model, supporting a one-time, offline KV-cache head budget computed on pretraining text.At inference, token scoring can use the cross-model per-token signal, including a score supplied by a much smaller model.Source: §1 · p. 2

2 RELATED WORK

Prior work has studied similarities in representations across independently trained models, while attention mechanics and their cross-model behavior remain less characterized. This paper positions marginal attention as a way to study those mechanics and as a signal for KV-cache eviction.

2 RELATED WORK

3 MARGINAL ATTENTION

Marginal attention reduces a head-by-position representation along its head and token axes to yield distinct per-token and per-head profiles. Across models and documents, the token profile tracks the text, while the head profile tracks the model.

3 MARGINAL ATTENTION

  • A head’s marginal attention averages its weighted attention to each input position over output positions, using a weight that rescales attention by the norm of the value sent through the head.Both factors used in the weighting can be computed in one forward pass.Source: §3 · pp. 3–4
  • Summing marginal attention across heads gives a normalized per-token profile, while summing over positions gives each head’s share of total marginal attention.The per-token profile is compared across models by document position; the per-head profile is compared across documents by model head.Source: §3 · p. 4

3.1 THE TWO REDUCTIONS ACROSS MODELS AND ACROSS DOCUMENTS

  • The evaluation compares 67 open models across 38 families and 28 architectures on 850 natural documents, with per-head comparisons using 170 documents.Tokenizers differ, so token values are mapped to covered characters; head profiles are binned by relative depth for cross-model comparison.Source: §3.1 · pp. 4–5
  • Per-token profiles agree across models on the same document (0.87) but not across different documents (0.02); per-head profiles agree across document halves within a model (0.99) and less across different architectures (0.37).The findings support the per-token profile as document-specific and the per-head profile as model-specific.Source: Figure 2 · p. 4, §3.1 · p. 5, Figure 3 · p. 5

4 IDENTIFIABILITY IN THE MARGINAL ATTENTION

The section relates the cross-model stability of per-token marginal attention to the input-output Jacobian, then develops a theory of when models’ marginal influences must be similar. Its empirical tests connect stronger next-token-logit agreement with stronger agreement in both marginal influence and marginal attention.

4.1 MARGINAL ATTENTION APPROXIMATES THE JACOBIAN

  • Across 85 documents, Qwen2.5-0.5B’s per-token marginal attention and marginal influence correlate at average Spearman ρ = 0.635; uniform weights reduce agreement to 0.534.Agreement is lowest on EuroParl, while agreement between distinct documents is nearly zero (0.022).Source: §4.1 · p. 6, Figure 4 · p. 6
  • Across 16 models, influence maps agree with the other-model consensus at 0.772, versus 0.861 for unweighted per-token marginal attention.Marginal influence from attention-free Mamba and RWKV-4 models also agrees with the consensus similarly to a transformer model.Source: §4.1 · p. 6
  • Marginal attention approximates marginal influence because it sums per-head Jacobian contributions, when those contributions do not systematically cancel and the input direction is generic with roughly constant norm.The head-summed block norm stays within 5% of its orthogonal-head value during Pythia-410M training and is at most 17% below it on four finished models.Source: §4.1 · p. 7

4.2 JACOBIAN IDENTIFIABILITY

  • Under three assumptions, models whose next-token logits agree under small single-position input perturbations have per-token marginal influences that agree, with an error bound independent of sequence length and embedding dimension.The bound tends to zero as the smoothness constant M and logit gap δ tend to zero.Source: §4.2 · p. 7
  • Within tokenizer-sharing model pairs, larger next-token-logit gaps accompany lower agreement in both per-token marginal influence and unweighted per-token marginal attention.Over training, cross-family agreement rises as loss falls and is near zero until it beats a bigram baseline.Source: §4.2 · pp. 7–8, Figure 5 · p. 7

5 APPLICATION: KV CACHE EVICTION

The application uses a static, model-specific per-head budget derived from Head Mass alongside a per-entry score for KV-cache eviction. The resulting method is competitive with approaches that recompute budgets per document or train them per target.

5 APPLICATION: KV CACHE EVICTION

  • Eviction ranks each entry using its head’s budget multiplied by a per-entry score, paired here with Expected Attention.Each head keeps its highest-scoring entries within its budget.Source: §5 · p. 8

5.1 HEAD BUDGET FROM THE PER-HEAD MARGINAL ATTENTION

  • Head Mass uses reconstruction queries, KVzip+ scoring, and KVzip’s global cut; its per-head profile is averaged across 200 pretraining documents and computed once per model.The profile is clipped to a floor of 0.01 and a cap of 1, then rescaled so heads spend the total budget on average.Source: §5.1 · p. 8

5.2 RESULTS

  • HM + EA improves over its same-scorer baseline on every target and setting, closing most of the gap to KVzip without a second target forward.The table reports RULER 4096 scores over 13 tasks and documents at 4×, 8×, and 16× compression.Source: Table 1 · p. 9, §5.2 · p. 9
  • HM + EA is ahead of AdaKV + KVzap in all but one cell, while its budget is an offline fixed table rather than a trained surrogate.The KVzap surrogate trains for hours and requires light online computation; the HM budget is built offline in minutes and needs no representation access.Source: §5.2 · p. 9
  • A proportional budget based on reading queries consistently loses to AdaKV + EA with no fixed budget, showing that cross-document stability alone does not ensure an eviction-useful profile.The reconstruction-query profile is the one used for Head Mass.Source: §5.1 · p. 8, §5.2 · p. 9
  • Tilting EA with the guide’s smaller-model map raises HM + EA in all but two columns, probing whether smaller-model attention carries part of the eviction signal.The guide is not intended as a practical method.Source: §5.2 · p. 9

6 CONCLUSIONS

The conclusions distinguish a text-intrinsic per-token profile from a model-specific per-head profile, connect the former to Jacobian statistics, and apply the latter to KV-cache budgets. The paper also notes limits on its theory, census, and proposed guide method.

6 CONCLUSIONS

  • Per-token marginal attention is shared across models, while per-head marginal attention is model-specific and stable across documents.The shared token profile separates documents without strongly aligning with classical syntactic structure.Source: §6 · p. 9
  • Marginal attention empirically approximates blockwise input-output Jacobian norms, and smooth models with close next-token distributions have close Jacobians under the stated smoothness assumption.The theory assumes discrete conditional distributions admit a sufficiently good smooth approximation.Source: §6 · p. 9
  • A per-head profile computed offline on pretraining text can serve as a KV-cache budget, with a shared per-token score supplied by a smaller model; the method is competitive with per-document or per-target alternatives.HM + Guide requires an extra small-model forward pass per document and is presented as a probe, not a practical method.Source: §6 · pp. 9–10
  • The census is mostly English and reads at most 2048 tokens per document, while the theory covers only per-token marginal influence and its link to marginal attention remains empirical.The 2048-token limit was retained to compare models with differing capabilities and context limits.Source: §6 · p. 10

AI USE STATEMENT

The authors describe AI-tool use in experiments, writing, theory exposition, and model-citation collection, while retaining responsibility for and review of the work.

AI USE STATEMENT

  • The authors used Claude and Claude Code, with authors setting the questions, claims, and experiments and reviewing retained contributions.The authors state that they take responsibility for the final content.Source: AI USE STATEMENT · p. 10
  • AI tools helped write experiment scripts and manage experiments; an author reviewed the results and code.The authors wrote the core library by hand, with autocompletion.Source: AI USE STATEMENT · p. 10
  • AI tools integrated raw results, made bulk edits and layout changes, and edited the paper; Figure 1 was drawn from an author sketch.Source: AI USE STATEMENT · p. 10
  • The authors proposed the theory’s definitions, hypotheses, theorems, and proofs; AI tools helped find gaps and improve exposition, and authors checked every item.Source: AI USE STATEMENT · p. 10
  • AI tools scraped model citations from corresponding HuggingFace repositories, and an author checked them.Source: AI USE STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

The reproducibility materials include the estimator library, experiment notebooks, and resources for constructing the corpus and running the KV-eviction experiments.

REPRODUCIBILITY STATEMENT

  • The library implements the marginal-attention estimator and its ranking, residualization, projection, and decomposition steps; each experiment is self-contained.Included notebooks reproduce figures and tables end to end.Source: REPRODUCIBILITY STATEMENT · p. 10
  • All 67 census models are public checkpoints, and the code release includes a corpus-building script.Source: REPRODUCIBILITY STATEMENT · p. 10
  • KV-eviction experiments use the kvpress library and its RULER and LongBench harnesses unchanged, with the authors’ budget and scorer added as a press.Appendices report controls, design sensitivity, assumption diagnostics, and theorem statements and proofs.Source: REPRODUCIBILITY STATEMENT · p. 10

A.1 DATA

The study compares per-token marginal attention across diverse text and model sets, asking whether agreement reflects a shared text signal rather than a generic profile. Agreement with other models is strongest against a multi-model consensus, varies by model relationship, and is near zero across different documents.

A.1 DATA

  • The mix1000 corpus samples 50 documents from each of 20 domains, while mix10k uses ten times as many documents per domain.The domains include natural text and three generated-text sets; the main text uses natural documents.Source: §A.1 · p. 19, Table 2 · p. 19

A.2 MODELS

  • Of the 69 models used, 67 form both the per-token census and per-head analyses; the remaining two are Qwen3-32B and Mistral-7B-Instruct-v0.3, used only as eviction targets.Source: Table 4 · p. 22

B.1 AGREEMENT BY PAIR CLASS

  • 0.8673 mean agreement with the other 66 models’ consensus exceeds the 0.7607 mean between-model agreement across all census pairs.The consensus comparison is averaged over 850 natural documents; the between-model value is reported with interval [0.7556, 0.7658].Source: Table 3 · p. 21, §B.1 · p. 23

B.2 AGREEMENT SATURATES AT ABOUT TWENTY MODELS

  • 0.8432 agreement with five-model consensus rises to 0.8567 with ten models and 0.8630 with twenty, then reaches 0.8673 with the full roster.The paper reports that roster size is not what produces the headline.Source: §B.2 · p. 23

B.3 A DIFFERENT DOCUMENT GIVES NEAR-ZERO AGREEMENT

  • 0.0155 agreement with other models across a different document is near zero, unlike agreement on the document at hand.The between-model different-document value is 0.0139.Source: §B.3 · p. 23

B.4 ROBUSTNESS: PUNCTUATION, WHITESPACE AND HEAD WEIGHTING

The section tests how per-token agreement responds to character removal, text scrambling, and head weighting, then examines per-head stability and the Jacobian connection. The evidence also reports heuristic conditions and measurement limits for interpreting that connection.

B.6 AGREEMENT ON SCRAMBLED TEXT

  • 0.8277 agreement on real text exceeds 0.7201 on shuffled words and 0.5974 on random token ids in the Pythia group.The comparison is not length-matched because a 2048-token cap spans fewer characters on scrambled text; sizes below shuffled words are provisional, but real text remains above shuffled words.Source: §B.6 · pp. 23–24

B.4 ROBUSTNESS: PUNCTUATION, WHITESPACE AND HEAD WEIGHTING

  • 0.8705 agreement after deleting punctuation and 0.8734 after deleting punctuation and whitespace exceed the undeleted 0.8673 value on 850 natural documents.Whitespace alone is the low row at 0.7336; it covers 2.3% of the axis. The unweighted variant scores 0.8605 on the same documents.Source: §B.4 · p. 23, Table 5 · p. 24

C PER-HEAD MARGINAL ATTENTION: ADDITIONAL EVIDENCE

  • 0.9529 leave-one-document-out agreement within a model contrasts with 0.3728 agreement between different-architecture models and −0.0003 for permuted heads.Across model pairs, base-versus-derived models agree at 0.9900; pair classes are distinct from those used for the per-token profile.Source: C PER-HEAD MARGINAL ATTENTION:… · p. 25, Table 6 · p. 25

D.1 A SINGLE ATTENTION LAYER

  • Per-token marginal attention and marginal influence align when head contributions do not systematically cancel, such as when heads write into linearly independent subspaces.Attention-only weighting drops each head’s value-channel gain; the Jacobian comparison also depends on the input direction and value-channel geometry.Source: §D.1 · pp. 25–26, Figure 9 · p. 26, Figure 10 · p. 26

B.5 THE SEED CEILING

  • The unweighted per-token profile and per-token marginal influence differ because one sums per-head norms while the other takes the norm of their sum.This distinction helps explain why attention and Jacobian-based statistics need not match exactly.Source: §D.1 · p. 25, §E.1 · p. 28

D.2 SEVERAL LAYERS

  • The Jacobian expansion treats direct attention edges as first-order paths and products of layer Jacobians as longer token-space paths.Under an even-distribution assumption, the two-hop gain is smaller than the general bound by about d_head/d_model, but trained composing heads are not randomly oriented, so this is not a bound.Source: §D.2 · pp. 26–27

E THEORY: MEASUREMENTS

  • The table compares per-token attention and three row-gradient estimates of marginal influence across 16 models; marginal attention is more shared across models.The comparison includes each model against a consensus excluding its family, as well as between-model and different-document references.Source: E THEORY: MEASUREMENTS · p. 27, Table 8 · p. 28

E.1 ESTIMATORS OF THE MARGINAL INFLUENCE

  • The block norm is a Frobenius norm of Jacobian blocks summed over output positions, but it is computed only for Qwen2.5-0.5B.Its agreement with per-token marginal attention is 0.64, versus 0.53 for uniform weights; the row gradient is a proxy used when block norms are too costly to compute across models.Source: §E.1 · p. 28

E.2 EMPIRICAL DIAGNOSTICS OF THE ASSUMPTIONS

The appendices test theoretical assumptions, track agreement during training, examine models without attention, and assess text-property prediction and head-budget stability. These results include empirical scope limits and support precomputing head budgets from pretraining documents.

E.2 EMPIRICAL DIAGNOSTICS OF THE ASSUMPTIONS

  • The diagnostics do not establish the theorem’s required uniform bounds or evaluate its quantitative bound.They examine sampled Jacobian blocks and same-document cases; measured lag dependence decays as a power law with exponent about one half to two thirds, with one fifth to one third more than 512 tokens back.Source: §E.2 · pp. 28–29

E.3 EMERGENCE DURING TRAINING

  • 0.84 agreement at the final Pythia-410M checkpoint contrasts with near-zero agreement at initialization, with most of the rise occurring by step 8,000.Agreement rises only after loss falls below the bigram reference of 2.51 bits per character; the evidence covers one model, one seed, and 60 documents.Source: §E.3 · p. 29, Figure 11 · p. 29

E.4 MODELS WITHOUT ATTENTION

  • 0.42 to 0.47 agreement for four Mamba models approaches the transformer reference’s 0.50, so the shared profile does not require attention.The evaluation uses one transformer reference and includes generated text; it does not test the theorem because logit gaps are unavailable for the attention-free models.Source: §E.4 · p. 30, Figure 12 · p. 30

E.5 WHAT THE SHARED PER-TOKEN PROFILE ENCODES

  • The five word-property groups predict the consensus profile at 0.52, below a model’s own profile agreement of 0.88.Grammatical role predicts best alone, followed by surface form, local predictability, and word frequency; position predicts nothing.Source: §E.5 · p. 30

F.1 HEAD BUDGET DETAILS

  • Head Mass tables from disjoint pretraining-document sets agree at median Spearman 0.98 or higher across target key-value heads.Figure 13 reports agreement for several document splits, including 1 versus 199, 5 versus 195, 20 versus 180, and 100 versus 100.Source: Table 9 · p. 31, Figure 13 · p. 31

F.2 HEAD BUDGET MAPS, ONE PER TARGET

  • 0.91 to 0.93 Spearman correlation links Head Mass tables to KVzip+’s realized held-out head keep rates, while prefill-query budgets agree far less.The table compares each target’s layer-by-key-value-head grid at 8× compression; Head Mass averages 200 pretraining documents.Source: §F.2 · pp. 31–32, Table 10 · p. 32, Figure 14 · p. 32

F.3 PROPORTIONAL BUDGET ABLATION

The appendix compares proportional and guided per-head budgets, explains per-task score plots, and states assumptions behind the Jacobian results. The proportional budget’s performance depends on the target and compression ratio.

F.4 THE GUIDED SCORE

  • The guided score tilts EA using a per-target ridge regression fitted on 48 documents, while leaving the Head Mass budget unchanged.It uses guide-model attention profiles from reading and reconstruction passes to predict KVzip+ log scores for target heads.Source: §F.4 · pp. 33–34

F.3 PROPORTIONAL BUDGET ABLATION

  • On Qwen3-32B at 16×, the proportional budget trails HM + EA by more than 11 points; on Qwen3-8B it stays within half a point at every setting.Elsewhere, it trails by up to 3.3 points. The paired budgets use KVzip’s reconstruction pass on the same pretraining documents.Source: §F.3 · p. 32, Table 12 · p. 35

F.4 THE GUIDED SCORE

  • At context 8192, HM + EA scores above AdaKV + EA in all 10 RULER columns; on LongBench, LU-KV + EA and HM + EA differ by at most 2.5 points over eight columns.The RULER result is the mean stringmatch score over 13 tasks at 6500 documents, with one run per cell. The LongBench mean is unweighted over 16 English subsets at their full official split (3750 documents); phi-4 was not run.Source: Table 14 · p. 36, Table 15 · p. 36

F.5 PER-TASK BREAKDOWN

  • Per-task figures group bars by task or LongBench subset, compare methods against an uncompressed-model line, and show where aggregate scores conceal differences.The aggregate scores are means of the per-task scores.Source: §F.5 · p. 37

G.1 ASSUMPTIONS

  • The Jacobian argument assumes smoothness and a lower bound on dependence on the immediately preceding token, alongside lag decay for earlier-token dependence.Lag decay is observed to be slow and close to a power law; whether a milder assumption suffices remains open. Centering logits handles their insensitivity to a fixed additive constant.Source: §G.1 · p. 46, G PROOFS · p. 46

G.2 PROOFS

The proof bounds Jacobian block differences from next-token logit similarity, then uses smoothness and lag decay to control per-token marginal-influence differences. The resulting bound does not depend on sequence length or embedding dimension.

G.2 PROOFS

  • The Jacobian-block lemma’s explicit estimates depend on smoothness and logit mismatch parameters, but not on sequence length or embedding dimension.The lemma concerns two models with smooth Jacobians and similar next-token logits.Source: §G.2 · pp. 46–47
  • Under smoothness, preceding-token dependence, and lag-decay assumptions, similar next-token logits imply per-token marginal influences differ by at most ε.The theorem’s ε tends to zero as the smoothness parameter M and logit mismatch δ tend to zero, and is independent of sequence length n and dimension d.Source: §G.2 · pp. 47–49
  • The proof first bounds differences between Jacobian blocks, then separates recent token-position pairs from distant pairs controlled by lag decay.It converts the resulting bound to normalized marginal influences, whose L1 distance is at most 2 unconditionally.Source: §G.2 · pp. 48–49

Related papers

Loading 2610.03109v1…