Attention in 3 lines

Table of contents

Pemberton, BC

Attention in 3 lines

"They're made out of weights."

"Weights?"

"Weights. Floating-point numbers. We checked the whole thing through. It's nothing but weights."

"Weights doing what? Where do the words come from?"

Short excerpt from https://maxleiter.com/blog/weights

LLMs can feel mysterious because the finished systems are enormous. The basic transformer idea is smaller than that: turn text into numbers, let every token look at the tokens before it, and use the result to predict what comes next. A lot of their capabilities emerge from the scale of training they undergo; see why scale matters.

The important word in that sentence is look. Attention is the part of the transformer that decides which previous tokens are useful for understanding the current token i.e. which token to pay more Attention too

Text Becomes Tokens

A model receives numbers, not words. So the first step is tokenization: splitting text into pieces — a word, part of a word, punctuation, whitespace — and giving each piece an integer ID.

For example:

"the bird is red" -> ["the", "bird", "is", "red"]

Those token strings are then mapped to IDs:

["the", "bird", "is", "red"] -> [1820, 10214, 374, 2579]

The IDs are just labels. If bird is token 10214, that number does not contain bird-ness. It is no more meaningful than a row number in a spreadsheet.

Tokenization → Embedding

Words become vectors

Text to embedding Each word is mapped to an embedding vector. The animation cycles through the words of a sentence and shows their vector representations forming. the embed

Eight of the vector's dimensions. Real models use hundreds or thousands.

Tokens Become Vectors

IDs become useful only once we map them into vectors, called embeddings: lists of numbers learned during training. Instead of an isolated dictionary entry, each token becomes a point in a learned space, and tokens used in similar contexts end up near each other.

word2vec made this famous. Nobody taught it gender, grammar, or analogy. The training goal was much simpler:

Given The cat sat on the mat., seeing cat should raise the probability of the, sat, and mat over unrelated words. Do that across enough text and words from similar neighborhoods drift together.

Which is why this became so striking:

king - man + woman ~= queen

Embedding arithmetic

Vector directions can carry relationships

King minus man plus woman lands near queen A looping vector addition diagram showing king minus man plus woman tracing a path that ends close to, but not exactly on, the queen vector in embedding space. dim 1 dim 2 king −man +woman queen

Nobody told the model "king is to man as queen is to woman." That structure emerged because the vectors were optimized to be useful for prediction.

One caveat: individual dimensions rarely map cleanly to human concepts. Dimension 1 is not "gender" and dimension 2 is not "color". Meaning is spread across many dimensions at once.

Embedding space

Random IDs become useful neighborhoods

before training
Before and after embedding clusters A scatter plot where word vectors begin randomly scattered and move into semantic clusters as training progress increases.
With training, words of similiar concepts end up together

At this point each token has a vector, but each vector is still mostly context-free. The embedding for bird represents a broad idea of bird. It does not yet know whether this bird is red, flying, extinct, angry, or part of a company name.

That is where attention enters.

Attention Adds Context

Take the sentence:

the bird is red

After tokenization and embedding, the model has one vector for each token:

the   -> vector
bird  -> vector
is    -> vector
red   -> vector

But the token bird should not stay generic. In this sentence, it should absorb information from red. The representation we want is closer to:

bird = bird, but informed by "red"

A naive version would be:

bird_context =
  some_amount_of("the") +
  some_amount_of("bird") +
  some_amount_of("is") +
  some_amount_of("red")

That is Attention in plain English. Each token builds a weighted mixture of other token vectors. High weight means "this token matters to me right now." Low weight means "ignore most of this."

Attention

Context is added, not looked up

Adding a red vector to a bird vector A bird vector, a red context vector added to it, and the resulting red bird vector. The three stages draw in as you scroll. start bird bird vector + add context bird + red bird + red direction = updated meaning red bird new context-aware vector

Scroll to draw. bird keeps its own direction and picks up a push from red.

The result is a new vector for every token. Not just bird, but also the, is, and red. Every position gets rewritten as a context-aware representation.

The 3 Lines

Here is the core attention operation in three lines:

x contains the embedded vectors from previous step

scores = (x @ x.T) / math.sqrt(d)
weights = masked_softmax(scores)
context = weights @ x

That is the heart of it.

The dot product x @ x.T measures similarity between every pair of tokens. If two token vectors point in similar directions, their dot product is larger. Larger score means the model should pay more attention.

Attention weights

Softmaxed weights with causal masking

Attention weight matrix A 4x4 matrix showing softmaxed attention weights between the tokens the, bird, is, and red. Cells above the diagonal are masked (causal). Darker cells mean higher weight. the bird is red the bird is red 1.00 – – – 0.35 0.65 – – 0.20 0.25 0.55 – 0.15 0.20 0.10 0.55 token being looked at token doing the looking
low
high softmaxed attention weights (– = causal-masked)

The division by sqrt(d) keeps the scores from becoming too large as vector size grows. Without that scaling, softmax can become too sharp too early, where one token gets almost all the weight and the others vanish.

masked_softmax does two things.

First, the mask prevents cheating. During text generation, token 3 is allowed to look at tokens 1, 2, and 3, but not token 4. Future tokens do not exist yet.

Second, softmax converts raw scores into weights that add up to 1:

raw scores -> attention weights

Then the final line computes the weighted sum:

context = weights @ x

Each row of weights says how much that token borrows from every other token.

Context vectors

Each row produces a weighted sum of embeddings

Context vector formation A looping animation cycling through each row of the attention weight matrix, showing the weighted sum equation and the resulting context vector bars. weights the bird is red the bird is red 1.00 – – – 0.35 0.65 – – 0.20 0.25 0.55 – 0.15 0.20 0.10 0.55 weighted sum the' = 1.00·the context vector

This is why attention is often described as "weighted lookup" or "soft lookup." It does not choose exactly one previous token. It blends information from many tokens in different amounts.

What The MLP Does

Attention only moves information between tokens. The MLP does the heavy lifting.

Ask a model:

What day was Michael Jackson born?

Attention cannot answer that. Your prompt does not contain the date, so there is nothing to copy from. That fact lives in the MLP weights.

The two parts split the work:

A transformer block usually repeats this pattern many times:

attention -> MLP -> attention -> MLP -> ...

Early layers catch local patterns; later layers build more abstract ones. The exact interpretation is messy, but the shape is useful: attention moves information between positions, MLPs process it at each position.

Eventually the model takes the final vector at the current position and turns it into a score for every token in the vocabulary. These scores are called logits.

Feed-forward → projection

Context vectors through the MLP and into logits

MLP and logit projection Context vectors flow through a 2-layer MLP with a SiLU activation, then a linear layer projects the result to a logit distribution over the vocabulary. context vectors the' cat' is' here' (8-dim each) MLP input 8 hidden 16 output 8 linear → logits Linear logits (vocab) the cat is here sat on mat ... argmax

After softmax, logits become probabilities:

P(".")    = 0.27
P("and")  = 0.17
P("bird") = 0.14
...

The model is not directly writing English. It is repeatedly producing a probability distribution over the next token.

Sampling The Next Token

Once the model has probabilities, we need to choose a token.

The simplest choice is greedy decoding: pick the highest-probability token every time. That is deterministic, but it can become dull or get stuck in repetitive patterns.

Other sampling strategies add controlled randomness:

Sampling

Pick a token from the logit distribution

Sampling from logits A logit distribution over the vocabulary. The animation highlights each candidate token, then selects one as the output. logit distribution
sampled token → ?

After a token is sampled, it is appended to the input. Then the model runs again to predict the next token. Then again. Then again.

That is autoregressive generation:

input -> predict one token -> append it -> predict one token -> append it

End to end

The three lines, actually evaluated

The full generation loop with real values Four tokens with four-dimensional embeddings. The scores, weights and context matrices are computed from those embeddings by the three attention lines, and the numbers shown are the real results. the three lines append the sampled token, run it all again

Every number is computed, not decorative. Only the embeddings and the unembedding are invented; everything downstream is the real arithmetic. Two things worth checking by eye: the first row of weights is exactly 1.00, so context's first row comes out identical to x's — the has nothing but itself to look at. And scores is symmetric, which is precisely what separate query and key projections would break.

Two Missing Details

There are two details worth adding before the picture feels complete.

First, attention by itself does not know word order. If you only compare token vectors, the set ["the", "bird", "is", "red"] looks too much like ["red", "is", "bird", "the"]. Transformers add positional information to the token embeddings so the model can tell where each token appears.

Second, real transformers do this attention operation many times in parallel. This is called multi-head attention. Each head has its own learned query, key, and value projections, so different heads can specialize in different relationships. One head might track nearby syntax. Another might connect names to pronouns. Another might focus on punctuation or formatting.

The key idea does not change:

compare tokens -> make weights -> mix values

Why Scale Matters

The recipe is simple, but scale changes what it can learn. A tiny transformer picks up local grammar. A large one picks up facts, style, code patterns, reasoning traces, translation.

GPT-2 had 1.5 billion parameters. GPT-3 had 175 billion. Language Models are Few-Shot Learners showed that at sufficient size, a model could perform many tasks straight from the prompt, with no task-specific fine-tuning.

Nobody taught it a general reasoning algorithm. Next-token prediction, at enough scale, forced it to learn representations useful for many behaviors.

As Max Leiter puts it in Weights:

They are made out of weights? Yes.

That is the unsettling and beautiful part. The model is "just" matrices, vectors, nonlinearities, and probabilities. But those weights encode a huge amount of structure about language.

Ending Thought

Attention is not the whole transformer, and transformers are not the whole story. There are tokenizers, positional encodings, layer norms, residual connections, MLPs, optimizers, datasets, sampling tricks, caches, and a lot of engineering.

But the center is surprisingly compact:

The three lines are not the whole model, but they are the part that makes the model context-aware:

scores = (x @ x.T) / math.sqrt(d)
weights = masked_softmax(scores)
context = weights @ x

That is attention: every token asking, "given where I am, which other tokens matter?"

Complete Code along with comments: gitlab

Tags: #attention #transformer #LLM