Advanced Attention

Attention is powerful, but it has a price: every token looks at every other token. Double the text, and the work becomes four times bigger. In this session, let us see how real models handle this, look at attention on real sentences, and understand what goes wrong when the context becomes very long.

01The n² problem

For n tokens, attention builds an n × n table of scores, in every head of every layer. So the work grows with n²:

02Efficient attention variants

Researchers have found many ways to make attention cheaper. Some approximate it (look at fewer pairs), and some compute it exactly but smarter.

MHA vs GQA vs MQA

During generation, the model stores the keys and values of all previous tokens (the KV cache). With many heads, this cache becomes huge. GQA and MQA let several query heads share the same keys and values:

03Attention on real sentences

Change one word at the end of the sentence and see how the attention of “it” moves:

You can see real attention weights from any Hugging Face model by asking for them:

from transformers import AutoTokenizer, AutoModel

tok = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModel.from_pretrained("bert-base-uncased", output_attentions=True)

inputs = tok("The animal didn't cross the street because it was too tired", return_tensors="pt")
out = model(**inputs)

print(len(out.attentions))        # one entry per layer
print(out.attentions[0].shape)    # (batch, heads, tokens, tokens)

For a visual, interactive view of every head, try the open-source BertViz library (pip install bertviz) in a notebook.

04What breaks as context grows

Modern models advertise 128K or even 1M-token windows. That does not mean they use every token equally well. Here is what goes wrong:

• Memory and cost: the KV cache grows with every token. In our calculator, 128K tokens need about 16 GB just for the cache of an 8B model, more than the model weights themselves.
• Lost in the middle: models find facts at the start and end of a long prompt more reliably than facts buried in the middle.
• Attention gets spread thin: with thousands of tokens competing, important ones can get a small share of attention.
• Position limits: a model trained on 8K-token texts may behave badly at 32K unless its positional encoding is stretched (methods like RoPE scaling and YaRN).

Practical tips for long prompts

• Put the instructions and the most important facts at the start or the end.
• Prefer RAG: retrieve the 5 most relevant chunks instead of pasting 300 pages.
• Summarise old conversation turns instead of keeping every message.
• Test with “needle in a haystack” style checks on your own documents before trusting a long-context model.

05Summary

• Attention cost grows with n²; the KV cache grows with n.
• Sliding-window and sparse attention skip pairs; FlashAttention stays exact but computes smarter.
• GQA / MQA share keys and values across heads to shrink the KV cache.
• Attention maps are a useful view, but not a full explanation.
• Long context brings memory cost, lost-in-the-middle and position limits, so design prompts and RAG with care.

06Quick quiz

07References

• Dao et al., FlashAttention (2022)
• Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models (2023)
• Beltagy et al., Longformer (2020)
• Liu et al., Lost in the Middle: How Language Models Use Long Contexts (2023)