DEV Community

jidonglab
jidonglab

Posted on

Attention Entropy: Why Softmax Blurs Retrieval at 128k Context

Your needle-in-a-haystack eval is green at 200k. Your production RAG pipeline, stuffing 80 retrieved chunks into the same model, answers from the wrong chunk about one time in six. Same model, same context length, same temperature. The difference isn't "context rot" as a vibe — it's arithmetic. Attention entropy grows with the number of competing keys, and softmax has a hard structural limit on how much probability mass one key can hold when the others are numerous and not that much worse.

This is a property of softmax, not a defect of any particular checkpoint. Understanding the bound tells you exactly which knob actually helps.

TL;DR

  • Softmax attention over n keys can only concentrate on one key if that key's logit beats the rest by a margin that grows like ln(n). Fixed margin + more tokens = dispersed attention.
  • With n-1 equal distractors, holding 90% of the mass on one key needs a logit gap of ln(9(n-1)). Going 4k → 128k costs an extra ~3.5 nats of gap the model was never trained to produce.
  • The competing background isn't just count. For roughly Gaussian distractor logits it's ln n + μ + σ²/2 — so logit variance across your context raises the floor too, and near-duplicate chunks stack multiplicatively.
  • Measure exp(H) per attention head (the "effective number of attended tokens"), not raw entropy. A few retrieval heads stay sharp; most don't.
  • You can't lower attention temperature through an API. You can shrink n: rerank harder, dedupe, drop boilerplate. Cutting 80 chunks to 20 buys ln 4 ≈ 1.4 nats of margin for free.

Why does attention entropy grow with context length?

Because softmax normalizes over everything you put in the window, and the denominator grows with n whether or not the added tokens are relevant.

Take one head, one query, n keys. Logits are z_i = q·k_i / √d_head. Say the correct key ("the needle") has logit z* = Δ and all n-1 distractors sit at 0. Then:

p_needle = e^Δ / (e^Δ + (n - 1))
Enter fullscreen mode Exit fullscreen mode

Solve for the gap you need to hold 90% of the mass:

Δ = ln(9) + ln(n - 1) ≈ 2.20 + ln(n - 1)
Enter fullscreen mode Exit fullscreen mode
  • n = 4,096 → Δ ≈ 10.5
  • n = 32,768 → Δ ≈ 12.6
  • n = 131,072 → Δ ≈ 14.0

Every 2x in context costs ln 2 ≈ 0.69 nats of required margin, forever. The absolute numbers matter less than the shape: sharpness is not free, and its price is logarithmic in context length. A model that produces a 10.5-nat gap — plenty at 4k — puts about 12% of its mass on the needle at 128k. The other 88% is spread across content that is individually irrelevant and collectively louder.

That's what "the model ignored the chunk it retrieved" looks like from inside the head.

Why doesn't a 1M-token context window fix this?

Because the window size sets n, and nothing in the architecture automatically grows the logit gap to match.

The gap has a budget. z_i = q·k_i/√d_head is bounded by ‖q‖‖k_i‖/√d_head, and both norms are shaped by RMSNorm plus learned projection scales that settled during training — mostly at sequence lengths far shorter than your serving length. The model learned to emit gaps that were sufficient for its training distribution. Extending the window with RoPE interpolation extends addressability; it doesn't extend the logit range.

Some architectures make the ceiling explicit. Gemma 2 soft-capped attention logits with a tanh (attn_logit_softcapping in the config) — a training-stability win that hard-bounds the maximum achievable gap, which is precisely the quantity long-context retrieval needs. Gemma 3 moved to QK-norm instead, which normalizes queries and keys and reintroduces the scale as a learned gain: better, but still a cap, and the gain has to be large enough to fund ln n.

This is also the honest reason a 1M-token context window doesn't obsolete retrieval. It obsoletes chunking as a memory constraint. It doesn't obsolete ranking, because ranking is what controls n.

What actually raises the competing background?

Count is only the first term. Model the distractor logits as roughly Gaussian with mean μ and standard deviation σ. The denominator is a log-sum-exp, and in the regime where σ is modest relative to √(2 ln n):

LSE ≈ ln n + μ + σ²/2
Enter fullscreen mode Exit fullscreen mode

Two things fall out that count-based intuition misses:

  1. Variance is a tax. A context with heterogeneous, spiky relevance (σ large) has a higher floor than a flat one at the same mean. Mixed-quality retrieval is worse than uniformly mediocre retrieval at equal n.
  2. Duplicates stack. Five near-identical boilerplate headers aren't one distractor; they're ln 5 ≈ 1.6 nats of background at the same logit level. Legal disclaimers, repeated nav chrome, and re-retrieved overlapping windows are the cheapest thing in your pipeline to delete and the most expensive to keep.

When σ gets large the sum becomes max-dominated and LSE ≈ μ + σ√(2 ln n), which is worse still — one unusually attractive distractor drags the whole denominator up on its own. That's the mechanism behind "one confusable chunk poisoned the answer."

Here's the calculator, which is more useful than it looks:

import numpy as np

def needle_mass(gap_nats, distractor_logits):
    """Attention mass on the needle given its logit gap over the distractor mean."""
    z = np.asarray(distractor_logits, dtype=np.float64)
    z = z - z.mean()                      # measure the gap from the distractor mean
    lse = np.logaddexp.reduce(z)          # log sum exp of the background
    return float(np.exp(gap_nats - np.logaddexp(gap_nats, lse)))

rng = np.random.default_rng(0)

for n in (2_000, 20_000, 128_000):
    flat  = rng.normal(0.0, 0.5, n)   # homogeneous context
    spiky = rng.normal(0.0, 2.0, n)   # heterogeneous, some strong distractors
    print(n,
          round(needle_mass(11.0, flat), 3),
          round(needle_mass(11.0, spiky), 3))

# 2000    0.855  0.383
# 20000   0.379  0.058
# 128000  0.088  0.010
Enter fullscreen mode Exit fullscreen mode

A fixed 11-nat gap — a genuinely strong, confident head — degrades from dominant to negligible purely from the denominator. Nothing about the needle changed.

How do I measure attention entropy in my own model?

Report exp(H), the perplexity of the attention distribution. It reads as an effective number of attended tokens: exp(H) = 1 means the head is a pointer, exp(H) = 500 means it's averaging 500 tokens' worth of values.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL = "Qwen/Qwen3-8B"  # any HF causal LM with eager attention
tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(
    MODEL, torch_dtype=torch.bfloat16, device_map="auto",
    attn_implementation="eager",         # required: SDPA/Flash never materialize the matrix
)

def effective_attended_tokens(prompt):
    ids = tok(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        out = model(**ids, output_attentions=True)
    rows = []
    for layer, attn in enumerate(out.attentions):     # [B, H, Q, K]
        p = attn[0, :, -1, :].float()                 # last query row, all heads
        H = -(p * torch.log(p.clamp_min(1e-12))).sum(-1)
        rows.append(torch.exp(H).cpu())               # effective attended tokens
    return torch.stack(rows)                          # [layers, heads]

short = effective_attended_tokens(build_prompt(n_distractors=4))
long  = effective_attended_tokens(build_prompt(n_distractors=64))

ratio = (long / short.clamp_min(1e-6))
print("heads that stayed sharp:", (ratio < 2.0).float().mean().item())
Enter fullscreen mode Exit fullscreen mode

Hold the needle and the question fixed; vary only the distractor count. Then look at the distribution, not the mean. What you'll typically see: most heads' exp(H) scales close to linearly with context — they're doing diffuse aggregation and always were. A small minority stay nearly flat. Those are your retrieval heads, and they are the ones carrying the answer. If the sharp fraction collapses between your 4-chunk and 64-chunk prompt, you've localized the failure to dispersion rather than to the retriever.

Two gotchas: output_attentions=True silently forces the eager path, so it's slow and memory-hungry — cap it at a few thousand tokens and extrapolate the trend. And with GQA the head axis is query heads, which is what you want here.

How do I fix attention entropy without retraining?

Attack n and σ. Those are the only two terms you own from an API.

Rerank to a hard budget, not a score threshold. A cross-encoder that lets 80 chunks through when the query is broad is handing the model an 80-way softmax. Fixing the budget at 15–20 costs a little recall and buys ln(80/20) ≈ 1.4 nats of margin on every head at once. In practice that trade is usually positive because the recall you lose is at ranks the model was going to lose to dispersion anyway.

Dedupe before you concatenate, not after. Overlapping sliding-window chunks are the single largest source of stacked-duplicate background in most RAG stacks. Near-duplicate collapse at the character-shingle level takes an hour to write and removes multiplicative background.

Strip repeated chrome. Same headers, same disclaimers, same "Source:" preambles across every chunk are pure denominator.

Two-stage instead of one-shot. If you genuinely need 100 candidates, run an extraction pass over batches of 10 and a synthesis pass over the extracts. Ten 10-way softmaxes concentrate far better than one 100-way. This is not prompt aesthetics — it's the log term.

If you own the serving stack, attention temperature is a real lever. YaRN's often-overlooked second half scales attention logits by a factor that grows with ln(scale), folded into the RoPE tables so it costs nothing at runtime:

{
  "rope_scaling": {
    "rope_type": "yarn",
    "factor": 4.0,
    "original_max_position_embeddings": 32768,
    "attention_factor": null
  }
}
Enter fullscreen mode Exit fullscreen mode

Leaving attention_factor null lets the implementation derive the logarithmic default rather than silently applying 1.0. Extending the window without it is the classic half-migration: the positions reach, the attention doesn't.

So why does softmax blur retrieval at 128k context?

Because softmax attention allocates probability by comparison, and every token you add joins the comparison. Concentrating on one key requires a logit gap that grows like ln n plus a variance term σ²/2 from the spread of the competing logits — and a trained model's gap budget is bounded by norms and projection scales fixed at training time, not by your serving window. So the needle's share decays as the haystack grows, even when the needle's own score never moves: an 11-nat gap that holds 85% of the attention mass at 2k holds under 10% at 128k. The window controls what the model can address; ranking controls what it can resolve. Since attention temperature isn't exposed through any API, the fixes that work are the ones that shrink the denominator — hard rerank budgets, deduplication, chrome stripping, and staged extraction — and the diagnostic that tells you it's happening is the effective-attended-tokens curve per head as you hold the needle fixed and add distractors.

Top comments (0)