2.1 — Transformer Architecture and Attention

Learning objectives

After this lesson, you should be able to:

• Explain what a token is and what attention does with it.
• Write scaled dot-product attention in NumPy.
• Explain why attention is divided by the square root of the vector size.
• Describe a transformer block and why position information is needed.

The one-sentence idea

A transformer lets every token look at every other token and mix in what is useful.

Attention is a weighted average whose weights are learned.

The idea in plain words

A transformer reads a list of tokens (pieces of text, or patches of an image) and lets every token look at every other token to decide what matters for it. The look-and-mix step is called attention. The architecture was introduced in “Attention is all you need” as a design based entirely on attention, without recurrence or convolutions.^[1] Almost every VLA you will meet is built on it, so it is worth knowing the parts.

Think of a sentence like “put the red block in the box”. To understand “block” the model should pay more attention to “red” than to “the”. Attention is a way to learn that kind of weighting.

Attention in six lines of NumPy

Each token is a vector. The model makes three versions of it: a query (what am I looking for?), a key (what do I contain?) and a value (what do I hand over if someone attends to me?). Scaled dot-product attention scores every query against every key, turns the scores into weights with softmax, and returns the weighted mix of the values: softmax(QKᵀ/√d)V.^[1]

import numpy as np

def softmax(x):
    e = np.exp(x - x.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

def attention(Q, K, V):
    scores = Q @ K.T / np.sqrt(K.shape[-1])      # how well each query matches each key
    weights = softmax(scores)                    # each row sums to 1
    return weights @ V, weights                  # a weighted mix of the values

tokens = ["red", "block", "box"]
X = np.array([[1.0, 0.0, 0.0, 1.0],              # hand-made 4-number vectors
              [1.0, 1.0, 0.0, 0.0],
              [0.0, 0.0, 1.0, 0.0]])
out, w = attention(X, X, X)                      # self-attention: queries, keys and values all come from X
print("attention weights (rows = who is looking, columns = who is looked at):")
print(np.round(w, 2))
print("each row sums to:", np.round(w.sum(axis=1), 2).tolist())
print("output shape:", out.shape)

attention weights (rows = who is looking, columns = who is looked at):
[[0.51 0.31 0.19]
 [0.31 0.51 0.19]
 [0.27 0.27 0.45]]
each row sums to: [1.0, 1.0, 1.0]
output shape: (3, 4)

Toy model. The vectors are made up and there are no learned weights. In a real model, Q, K and V come from learned linear maps of X.

Read the weights row by row. The first row says how much “red” listens to each of the three tokens; the numbers add up to 1. The output for each token is a new vector that mixes in information from the tokens it attended to.

Why divide by the square root?

The division by √d is not decoration. The paper’s reason is that dot products grow with the vector length, which would push softmax into regions where almost all the weight lands on one token.^[1] You can see it directly:

import numpy as np
rng = np.random.default_rng(0)

def softmax(x):
    e = np.exp(x - x.max())
    return e / e.sum()

for d in (4, 64, 1024):
    q = rng.normal(size=d)
    keys = rng.normal(size=(5, d))
    raw = keys @ q
    print("d =", d, "| biggest weight without scaling:", round(float(softmax(raw).max()), 2),
          "| with scaling:", round(float(softmax(raw / np.sqrt(d)).max()), 2))

d = 4 | biggest weight without scaling: 0.43 | with scaling: 0.31
d = 64 | biggest weight without scaling: 1.0 | with scaling: 0.51
d = 1024 | biggest weight without scaling: 1.0 | with scaling: 0.41

As the vectors get longer, the unscaled version tends to put nearly all its weight on one key, which makes learning hard. The scaled version stays balanced.

Many heads, and a block around them

One attention pattern is limited, so transformers run several at once, called heads, each on a smaller slice of the vector, and concatenate the results. The base model in the paper used 512-dimensional vectors and 8 heads.^[1] Around attention sits a small feed-forward network, and each of the two parts is wrapped in a residual connection (the input is added back to the output) followed by layer normalisation.^[1] Stacking many such blocks gives the transformer.

import numpy as np
rng = np.random.default_rng(0)

n_tokens, d_model, n_heads = 6, 512, 8
x = rng.normal(size=(n_tokens, d_model))
d_head = d_model // n_heads
heads = x.reshape(n_tokens, n_heads, d_head).transpose(1, 0, 2)   # (heads, tokens, per-head size)
print("one slice per head:", heads.shape)
merged = heads.transpose(1, 0, 2).reshape(n_tokens, d_model)       # concatenate the heads again
print("merged back:", merged.shape, "| identical to the input:", bool(np.array_equal(merged, x)))

def layer_norm(v):
    return (v - v.mean(axis=-1, keepdims=True)) / (v.std(axis=-1, keepdims=True) + 1e-5)

block_out = layer_norm(x + 0.1 * rng.normal(size=x.shape))          # "residual + normalise", with a stand-in for the sub-layer
print("block output shape (same as input):", block_out.shape)

one slice per head: (8, 6, 64)
merged back: (6, 512) | identical to the input: True
block output shape (same as input): (6, 512)

Toy model. The “sub-layer” here is random noise standing in for a real attention or feed-forward layer. Only the shapes and the residual pattern are real.

The key point: a block takes a list of token vectors and returns a list of the same shape, so blocks can be stacked as deep as you like.

Attention does not know the order

Attention treats its input as a set. Shuffle the tokens and the outputs shuffle in the same way, nothing more. That is why the paper adds position information to the inputs, since the model has no recurrence or convolution to provide it.^[1] A quick check:

import numpy as np

def softmax(x):
    e = np.exp(x - x.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)

def attention(X):
    return softmax(X @ X.T / np.sqrt(X.shape[1])) @ X

rng = np.random.default_rng(0)
X = rng.normal(size=(5, 8))
perm = rng.permutation(5)

same = np.allclose(attention(X[perm]), attention(X)[perm])
print("shuffling the tokens just shuffles the outputs:", same)

shuffling the tokens just shuffles the outputs: True

For a VLA this matters twice: the words of an instruction have an order, and the patches of an image have positions. Both must be injected, which the next lessons come back to.

Common misconceptions

Engineering takeaways

1. Attention = softmax(QKᵀ/√d)V.
2. Rows of attention weights sum to 1.
3. Heads work in parallel on slices of the vector.
4. A block keeps the shape, so blocks stack.
5. Order must be added explicitly.

Attention is a learned, weighted look at everything else.

Knowledge check

Try each question first, then reveal the explanation.

What’s next

Next: the Vision Transformer, which turns an image into tokens so this machinery can read it.

Back to the VLA roadmap