JustujuAI, layer by layer.Open in the interactive map

Transformers & attention

The architecture behind modern language models: stacked layers in which every token can look back at earlier tokens and pull in what's relevant.

Layer 5 · Foundation Models · Build and run time · reviewed 2026-10-30

Read first: Tokens & tokenization.

D1 · StoryFor young children and total beginners

Read the sentence slowly. Which word does 'it' point back to?

Read this: “The cat sat on the mat because it was tired.” You know “it” means the cat, not the mat.

Attention is how the AI works that out: for each word it looks back at the earlier words and decides which ones matter most.

Like: Highlighting the words that matter.

D2 · PictureFor curious beginners

A transformer is a stack of identical layers. Each layer has two parts: attention, which mixes in information from other positions, and a feed-forward network, which processes each position on its own.

In attention, each position asks a question (a query), compares it with a label on every earlier position (its key), and takes a weighted blend of what those positions carry (their values). Many attention heads do this in parallel, looking for different patterns.

In models that generate text, a causal mask stops any position from looking ahead at tokens that haven't been written yet.

Like: Highlighting the words that matter.
Where it breaks: You highlight a few words once. Attention gives every earlier token some weight, dozens of heads do it at once in every layer, and the result is spread across many layers rather than one clear highlight.

D3 · PlayFor motivated learners

Each word builds a query (what it is looking for) and a key (what it offers). The query is compared with the keys of earlier words using a dot product, and softmax turns those scores into attention weights that add up to 100%.

A transformer runs many heads side by side. Each head has its own queries and keys, so each looks for something different. Choose a word, then switch heads and watch where its attention goes.

Try the interactive version in the map.

Things to try

D4 · MechanismFor practitioners

Stack every token's query into a matrix Q, its key into K and its value into V. Attention is softmax(QKᵀ ÷ √d) · V, where d is the length of each query and key.

Row i of QKᵀ holds the dot product of token i's query with every key. Dividing by √d stops the scores growing just because the vectors are longer. The causal mask sets scores for later tokens to −∞ so softmax gives them zero weight. Each output row is a weighted blend of the value vectors.

Edit the vectors below; every table is computed live from the formula.

Try the interactive version in the map.

Things to try

def attention(Q, K, V, causal=True):
    d = Q.shape[-1]                      # length of each query and key
    scores = Q @ K.T / sqrt(d)           # one score for every (query, key) pair
    if causal:
        scores[j > i] = -inf             # hide tokens that come later
    weights = softmax(scores, axis=-1)   # each row adds up to 1
    return weights @ V                   # blend the values

# multi-head: run several of these with different learned Q, K, V
# projections, then join the outputs and mix them with one more matrix

Trade-offs

D5 · FrontierFor experts

D4's formula is exact, but its cost grows with the square of the context, and the picture of heads looking for patterns is only partly understood.

Exact but faster. Kernels that avoid memory traffic made exact attention much faster without changing its quadratic arithmetic (Dao et al., 2022).

Replacing attention. State-space layers run in linear time with about five times the throughput (Gu and Dao, 2023), yet efficient alternatives lag attention at recalling earlier context, and hybrids close most of the gap (Arora et al., 2023).

Using long context. Models use information at the start or end of a long input better than in the middle (Liu et al., 2023).

What heads do. Heads that copy earlier patterns may underlie learning from examples in the prompt, but the causal evidence is strong only for small attention-only models (Olsson et al., 2022).

Examples as of 2026-10-30

Sources

  1. Attention Is All You Need (2017)
  2. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022)
  3. Mamba: Linear-Time Sequence Modeling with Selective State Spaces (2023)
  4. Gemma 2: Improving Open Language Models at a Practical Size (2024-06-27)
  5. Zoology: Measuring and Improving Recall in Efficient Language Models (Arora et al.) (2023-12-08)
  6. Lost in the Middle: How Language Models Use Long Contexts (Liu et al.) (2023-07-06)
  7. In-context Learning and Induction Heads (Olsson et al.) (2022-09-24)

How it connects

More in Layer 5 · Foundation Models