Step by Step, Easy Transformer Calculation by Hand

1. Words become lists of numbers

A language model cannot read words as they are. A sentence written in natural language needs to be converted into numbers before the model can process it. The text is first split into small pieces of text called tokens, and every token must come from a fixed list of text pieces called the vocabulary. For each entry in this list, the model stores a list of numbers that it can use in its calculations. Splitting a sentence into these pieces lets the model look up the numbers for each piece. In real models, a token is often a piece of a word (a subword), and the vocabulary is large: Llama 3 has 128,256 tokens. For an easy example, our vocabulary has six tokens, each one a whole word, and each token gets an ID from 0 to 5:

The ID is a row number that tells the model where to look up a token's numbers. And the row with id=3 holds the four numbers for cats, [1, 1, −1, −1]. An ordered list of numbers like this is called a vector, and this vector is what we call the token's embedding. A table of numbers arranged in rows and columns is called a matrix, so the whole six-row table in this example is the embedding matrix E. The model learns these numbers during training.

Any text a model reads has to be written with tokens from this vocabulary. Lets say we give the model an example input of a sentence I like cats too. A tokenizer, the program that splits texts into tokens, turns the sentence into four tokens in this example, I / like / cats / too, and replaces each token with its ID: [0, 2, 3, 5]. The tokens you and dogs are in the vocabulary but not in this sentence, so their rows are not used. A different input, such as you like dogs, would use different rows.

Looking up the rows for the four IDs and stacking them in sentence order gives a 4 × 4 matrix X, with one row per token and one column per number.

Each row of X is the hidden state of one token, which is the list of numbers the model uses to represent that token while it processes the sentence. The model never shows these numbers as output, which is why they are called "hidden".

At the start, a token's hidden state is simply its embedding. For cats, that is [1, 1, −1, −1].

The hidden state changes as the token passes through the model's layers. A layer is one step of computation applied to the hidden states, and a model is a stack of them. A Transformer uses two kinds of layer, attention and the feed-forward network, and sections 2 to 10 build both. One attention layer followed by one feed-forward layer makes a block, and a model stacks many blocks. Each layer reads the hidden states, computes an update, and adds that update back, so the same four positions hold different values after each layer. After attention, for example, the hidden state of cats also contains information taken from I and like:

A hidden state therefore starts as a description of the word on its own and becomes a description of the word in its sentence. At the very end, the model turns the last hidden states into a prediction of the next word.

The length of a hidden state never changes: 4 numbers here, 4,096 in Llama 3 8B, a model with 8 billion learned numbers. Because of that, every layer can read the same shape and add its update to it. From here on, we follow the hidden state of cats, which is row 2 of X (rows count from 0). The other three rows go through exactly the same steps.

2. Each hidden state is turned into three smaller vectors

Right now, each hidden state describes only its own token. The hidden state of cats says nothing about the words around it, yet understanding cats in this sentence requires knowing that it comes after like. Attention is the layer that adds this context: it lets each token take in information from the other tokens in the sentence.

Each token first decides how much to read from every token, including itself, and then reads that much from each one and adds the results together. Deciding and reading need different information, so every hidden state is turned into three new vectors:

• The query q is used when a token is the one reading. It describes what this token looks for in the others.
• The key k is used when a token is the one being read. Comparing a reader's query with another token's key gives a number that says how well the two match. A better match means more reading. Section 4 computes these comparisons.
• The value v is the content a token gives out when it is read. Each reader collects a mix of values, taking more from the tokens that matched well. Section 6 computes this mix.

Each token both reads and is read. When cats reads, it compares its query with the keys of I, like, and cats. (It skips too, which comes later; section 5 explains why.) When too reads, it compares its own query with the key of cats, among others, and the value of cats is what it receives. Keys and values are kept separate because the reason a token is found and the information it should pass on can differ. A word such as like might be found because it is a verb, while what it should pass on is what it means.

Each of the three comes from multiplying the hidden state by its own weight matrix. A weight is a number the model learns during training, like the numbers in the embedding matrix, and a weight matrix is a table of them. W_Q, W_K, and W_V each have shape 4 × 2 (4 rows and 2 columns), so a hidden state of four numbers becomes a query, a key, and a value of two numbers each.

Multiplying a vector by a matrix works one column at a time. For each column, multiply the numbers in matching positions of the vector and the column, then add the products. This multiply-and-add is called the dot product, and each column's dot product gives one number of the result.

For cats, the hidden state is [1, 1, −1, −1]. Column 0 of W_Q is [0, 0, 0, −1] and column 1 is [−0.5, 0, 0.5, 0], so the two numbers of its query are:

So the query of cats is [1, −1]. The same calculation with W_K and W_V gives its key [0, −1.5] and its value [1.5, −1.5]. Doing this for all four rows gives three 4 × 2 matrices, and each of their rows still belongs to one token. Written for all four rows at once, with X from section 1:

Click any number in Q, K, or V to see its calculation.

Real models use the same operation with much wider vectors. Llama 3 8B, for example, has 4,096 numbers per hidden state and 128 numbers per query, key, and value.

3. Marking where each token sits

The query, key, and value of cats depend only on the word cats, not on where it sits. The same holds for every token. If attention used them as they are, swapping I and like would leave the result for cats unchanged, and the model could not tell I like cats from like I cats.

Llama and most recent models fix this with rotary position embedding (RoPE). A vector of two numbers can be drawn as an arrow. RoPE turns the arrow of each query and key by an angle that grows with the token's position: position p turns it by p radians, and 1 radian is about 57°. Turning keeps the arrow's length. For an angle a:

cats is at position 2 (positions count from 0), and cos 2 = −0.416, sin 2 = 0.909. Its query [1, −1] becomes:

Its key [0, −1.5] becomes [1.364, 0.624] the same way. I sits at position 0, so its query and key do not turn at all. We write the turned matrices as Q′ and K′. Values are not turned.

Click any number in Q′ or K′ to see its rotation.

Turning two arrows changes their dot product only through the angle between them. The query and key of cats both turn by 2 radians, so their dot product stays 1×0 + (−1)×(−1.5) = 1.5, the same as before. A query at position 2 and a key at position 0 end up 2 radians further apart than they started, wherever the pair sits in the text. A score therefore depends on how far apart two tokens are.

Real queries and keys hold many pairs: Llama 3 8B has 128 numbers, so 64 pairs. Each pair turns at its own speed. The first turns 1 radian per position, and later pairs turn more slowly, so they can tell apart positions that are thousands of tokens away. The original Transformer took a different route and added fixed sine and cosine patterns to the embeddings from section 1.

4. Comparing one query with every key

With a rotated query and key for every token, cats can now compare its query [0.493, 1.325] with each token's key. The comparison is again a dot product, and its result is called an attention score. For now we compute a score with every key, including too; section 5 removes that one.

In general, the score of a query q and a key k with two numbers each is:

A product is positive when both numbers have the same sign, so a query and a key match well when they have large numbers with the same signs in the same positions. For cats, the largest score is 1.5, with its own key: [0.493, 1.325] and [1.364, 0.624] are positive in both positions.

Every token's query is compared with every key, which gives 4 × 4 = 16 scores. They all come from one matrix product, S = Q′K′^T. Multiplying two matrices repeats the rule from section 2 for each row of the first matrix. The ^T means K is transposed: its rows and columns are swapped, so each key becomes a column. Row i of S then holds the scores of token i's query against every key, and row 2 is the four numbers above.

5. Turning scores into weights that add up to 1

The scores of cats, [1.16, −1.53, 1.5, −2.07], can be negative and do not add up to anything in particular. Attention needs attention weights instead: numbers that are zero or positive and add up to 1. They are not the learned weights inside a weight matrix.

Scale

First, every score is divided by √d_k, where d_k is the length of a key, the count of its numbers. Ours is 2, so we divide by √2 ≈ 1.41. A dot product adds one product for every number in the vectors, so the longer the vectors, the larger the scores tend to be. The typical size of a score grows with √d_k, so dividing by √d_k brings the scores back to a steady size. Without it, a model with 128-number keys would get very large scores, and the steps below would put almost all the weight on a single token.

Mask

Second, each token is stopped from reading the tokens that come after it. The reason is how a language model produces text: one token at a time, each predicted from the tokens before it. Given I like, the model predicts the next token, and only then, given I like cats, the one after that. When it predicts what follows like, the token cats does not exist yet, so like must not be able to read it. During training the whole sentence is available at once, so the same rule is applied there too. Otherwise the model could simply read the next token instead of learning to predict it.

The rule is enforced by the causal mask. A mask is a pattern that marks which scores may be used. It replaces every blocked score with −∞ (negative infinity), and the next step turns −∞ into a weight of exactly 0, so a blocked token contributes nothing. "Causal" means a token may use only itself and the tokens before it, the way a cause always comes before its effect. In the table below, each row is a token that reads, and each column is a token it might read:

For cats, only the score for too is blocked, so −1.466 becomes −∞.

Softmax

Third, softmax turns the row into weights. For a row of scores s₀, s₁, s₂, s₃, the weight of token j is:

Here e is a fixed number, about 2.718, and e^x means e raised to the power x. It is always positive: it grows quickly as x grows, equals 1 at x = 0, and gets closer and closer to 0 as x becomes more negative. So the top of the fraction turns any score into a positive number and keeps a larger score larger. The bottom adds up the same thing for every score in the row, so dividing by it makes the weights add up to 1. Since e^−∞ = 0, the blocked token gets a weight of exactly 0.

So cats gives 52.6% of its attention to itself, 41.2% to I, 6.2% to like, and none to too. Each of the other rows is computed the same way.

Row 0 has only one allowed key, its own, so I puts all of its weight on itself.

Dot-product scores, scaled and turned into weights, are the core of what Vaswani et al. (2017) named scaled dot-product attention; section 6 completes it by mixing the values. Written for all rows at once, the mask is a matrix that holds 0 where reading is allowed and −∞ where it is blocked, so adding it leaves the allowed scores unchanged:

6. The weights mix the values

The attention weights say how much to read from each token. What gets read is each token's value. The attention output of a token is the weighted sum of all the values: multiply each value by its weight, then add the results.

For cats, the weight of too is 0, so three terms remain:

Because cats put 52.6% of its weight on itself, its own value [1.5, −1.5] dominates the result [0.61, −0.49].

Every token builds its output this way. Because the queries, keys, and values all come from the same sentence, this is called self-attention; translation models also use a kind that reads a second sentence. In matrix form, the outputs of all four tokens are Z = AV.

Z has two numbers per token, but a hidden state has four, and the next step adds the attention output to the hidden state. The two must have the same length for that. A weight matrix that changes the length of a vector is often called a projection. The output projection W_O, a 2 × 4 weight matrix, maps the two numbers back to four:

This four-number vector is the full result of attention for cats.

7. Several copies of attention, sharing keys and values

So far, attention has used one set of W_Q, W_K, and W_V. One such set, together with the attention it computes, is called a head. Real models run several heads side by side, each with its own weights, so each head can look for something different. One head might learn to attend to the previous word, and another to the subject of the sentence. We now use four heads. The head we computed becomes head 0, and heads 1 to 3 get new weights.

Each head produces its own two-number output. The four outputs are placed side by side, which gives eight numbers, and an output projection maps them back to four, as W_O did in section 6.

Every head needs its own query, because the query decides what the head looks for. Keys and values, by contrast, can be shared by several heads. The query part of a head is called a query head, and its key and value part a KV head.

A model generates text one token at a time, and the causal mask from section 5 means earlier tokens never read later ones, so their keys and values never change. The model therefore stores them instead of recomputing them at every step. This store is the KV cache, and every KV head adds to it for every token, so the number of KV heads decides how much memory generation needs. Models use one of three arrangements:

• MHA (multi-head attention): every query head has its own key and value head. 4 query heads, 4 KV heads.
• GQA (grouped-query attention): query heads are split into groups, and each group shares one KV head. 4 query heads, 2 KV heads.
• MQA (multi-query attention): all query heads share a single KV head. 4 query heads, 1 KV head.

In MHA, heads 0 and 1 have separate keys. For cats, head 1 spreads its attention weights as 0.17, 0.67, and 0.16 over I, like, and cats.

In GQA, heads 0 and 1 share one KV head. A common way to build the shared head is to average the weight matrices of the heads it replaces. Ainslie et al. (2023) convert a trained MHA model this way and then train it for a short time. With averaging, the key of I becomes the mean of head 0's [1, 0.5] and head 1's [−0.5, −1], which is [0.25, −0.25].

Each query head keeps its own query, but now compares it with the shared keys. Head 1's attention weights for cats change from 0.17 / 0.67 / 0.16 to 0.25 / 0.54 / 0.21. Head 0's change too, from 0.41 / 0.06 / 0.53 to 0.19 / 0.18 / 0.63, and with the shared values its output moves from [0.61, −0.49] to [0.63, −0.13]. MQA goes one step further and averages all four heads into one.

Per token, MHA caches 4 key heads and 4 value heads of 2 numbers each, which is 16 numbers. GQA caches 8, and MQA caches 4. Llama 3 8B has 32 layers, 32 query heads, and 8 KV heads of 128 numbers each. Computer memory is counted in bytes, and each number here takes 2 bytes, so the model caches 128 KB (about 128,000 bytes) per token, or 1 GB (about a billion bytes) for 8,192 tokens. With one KV head for each of the 32 query heads (MHA), the same 8,192 tokens would take 4 GB.

MQA saves the most memory, but one shared key and value head for every query gives the heads less to work with. Ainslie et al. report that MQA loses quality and can make training less stable. GQA sits between the two: it keeps most of MHA's quality and most of MQA's memory saving, which is why most recent open models use it.

8. Attention's result is added to the hidden state

The rest of the article goes back to the single head from sections 2 to 6. Its attention result for cats, [0.61, −0.49, −0.25, 0.31], is added to the hidden state of cats:

Adding a layer's output to the layer's own input is called a residual connection. The name says what the layer has to compute: only the residual, the amount by which the hidden state should change, while the input passes around the layer unchanged. Writing h for the hidden state:

Because of this addition, the original embedding is still inside the new hidden state, and attention only contributes an adjustment. If attention learned nothing useful, the token would simply keep its old vector.

Every later layer is connected the same way, so the hidden state becomes a running total that flows through the entire model. This running total is called the residual stream, and section 12 follows it through two full blocks.

9. Rescaling each layer's input to a fixed size

The addition changed the numbers of cats, and it also made them larger overall. One way to measure the overall size of a vector is the root mean square (RMS): square each number, take the mean of the squares, then take the square root.

Here n is how many numbers the vector has, 4 for cats. Before and after the addition, the RMS of cats is:

Every residual addition can grow the hidden state like this, and a model with dozens of layers adds dozens of times. Growing inputs are a problem because a layer's output grows with its input. The scores in section 4 would grow too, and section 5 showed that large scores make softmax put almost all the weight on one token. Rescaling a vector to a standard size is called normalization, and RMSNorm (RMS normalization) does it by dividing the vector by its RMS. It is applied to the copy that goes into a layer, before every layer; the hidden state itself keeps the unnormalized sum. For a vector x:

The new vector has an RMS of 1 again. In practice, RMS also adds a tiny constant ε (0.000001 here) under the square root, so a vector of all zeros does not cause a division by zero. γ is a learned scale for each of the four positions. Here every γ is 1, so the division is the only change.

Since RMSNorm also runs before attention, sections 2 to 6 should have normalized X first. They could skip it because every embedding in E is made of 1s and −1s, so its RMS is exactly 1, as the "before" calculation above shows, and dividing by 1 changes nothing.

An older method, LayerNorm, also subtracts the mean of the vector before dividing by its standard deviation, a measure of how spread out the numbers are. RMSNorm skips the subtraction, which makes it cheaper to compute, and Zhang and Sennrich (2019) found it trains to comparable quality. Llama, Mistral, and Gemma all use RMSNorm.

10. A second layer that updates each token on its own

The normalized vector [1.46, 0.46, −1.13, −0.63] is the input to the feed-forward network (FFN), the second kind of layer in a block. Unlike attention, the FFN looks at one token at a time: it updates cats without reading any other row. For a vector x, it computes:

Up projection

First, the matrix W_up (4 × 8) expands the four numbers into eight. Column 0 of W_up is [0, 1, 0, −1], so the first of the eight numbers is 0.46 − (−0.63) = 1.09. All eight are:

With eight numbers, each one can respond to a different pattern in the four inputs. Real models expand by three or four times; Llama 3 8B goes from 4,096 numbers to 14,336. This example expands by two.

Activation

Second, the same rule is applied to each of the eight numbers separately. A rule that turns one number into another is called a function, and this one is the activation function. We use ReLU, which keeps positive numbers and replaces negative ones with 0: ReLU(x) = max(0, x). Four of our eight numbers are negative, so four become 0.

The activation is the only part of the FFN that is not a matrix multiplication. It is needed because two matrix multiplications in a row can always be merged into one: without ReLU, W_up followed by W_down would equal a single 4 × 4 matrix, and the FFN could only compute what that one matrix can. ReLU is non-linear: it cannot be written as multiplying and adding, because its output depends on whether each number was positive. Stacking many such layers lets the network represent far more complicated rules than one matrix can. Most recent models replace ReLU with GELU, a smoother version, or with SwiGLU, which also adds a third weight matrix to the FFN.

Down projection

Third, W_down (8 × 4) maps the eight numbers back to four. Each row of our W_down has a single entry of 0.5 or −0.5, so each output is half of a sum of a few activations. For example, output 1 reads activations 0 and 5: 0.5 × 1.09 + 0.5 × 0 = 0.54.

Like the attention result, the FFN result is added to the stream:

Attention and the FFN together make one block, and each of them reads a normalized copy of the hidden state and adds its result back. The hidden state of cats entered the block as [1, 1, −1, −1] and leaves it as [1.61, 1.05, −1.81, −1.51].

11. Normalizing before each layer or after each addition

In our block, RMSNorm runs before each layer, on the copy that goes into it, while the stream itself stays unnormalized. This arrangement is called pre-norm. The original Transformer (Vaswani et al., 2017) normalized after each addition instead, which is called post-norm.

The two figures show the difference in the numbers for cats. In the first block, both layouts give the same attention sum, [1.61, 0.51, −1.25, −0.69], and the FFN receives the same normalized input in both. What differs is what stays in the stream. Pre-norm keeps the unnormalized sum and adds the FFN result to it. Post-norm replaces the stream with the normalized version and adds the FFN result to that.

Most current models, including Llama and Mistral, use pre-norm.

Models moved to pre-norm because it is easier to train. Training adjusts every weight a little at a time, and the change for each weight comes from its gradient: a number that says in which direction, and how strongly, changing that weight would reduce the model's error, meaning how far its predictions are from the correct next tokens. Gradients are computed at the output first and then passed backward through the layers to the early ones.

The learning rate sets how big each weight change is. In post-norm, the backward path passes through a normalization at every block, and at the start of training the gradients near the output are large. Large gradients mean large weight changes, which can push training off course before it has really started. Xiong et al. (2020) showed that this is why post-norm models need a learning-rate warm-up, which keeps the learning rate very small at the start and then raises it slowly.

Pre-norm leaves a path made only of additions from the last layer to the first. In an addition, changing an input by some amount changes the output by the same amount, so the gradient passes back through each addition unchanged and reaches the early layers directly. Training is then stable without the warm-up.

In pre-norm, nothing normalizes the stream itself, so it keeps growing from block to block, as the next section shows. Pre-norm models therefore apply one last RMSNorm after the final block. Post-norm can still match or beat pre-norm when it does train well, and some models use both: Gemma 2 normalizes the input and the output of each layer.

12. The residual stream through two blocks

A real model stacks many blocks, and the residual stream passes through all of them. Here we run two: attention, FFN, attention, FFN. In the diagram below, the vertical line is the stream of one token. Each dashed branch is a layer that reads from the stream and writes its update back, and each dot shows the stream after an addition.

With pre-norm, the stream of cats grows at every addition. Its RMS goes from 1 to 1.11, 1.52, 1.87, and 2.26. Each layer still receives an input with an RMS of 1, because it reads a normalized copy. With post-norm, the RMS returns to 1 after every addition.

Each step only added to the pre-norm stream, so its final vector is exactly the embedding plus the four updates:

Because every update stays in the sum, Elhage et al. (2021) describe the stream as a shared channel. Each layer adds to it, and every earlier update is still present when a later layer reads it.

13. Reading the prompt, then adding one token at a time

After the last block, the model turns the hidden state of the last token into a score for every token in the vocabulary, and the next token is picked from those scores, usually one with a high score. The article Why Use Weight Tying? works through those scores. The new token is appended to the text, and the model runs again to find the token after it.

Say the prompt is I like cats. The first run is called prefill: all three prompt tokens go through every block at once, the same way the four tokens did in sections 1 to 12. Each attention layer stores the rotated keys and the values of the three tokens in the KV cache. Prefill ends with the first new token, say too.

Every run after that is called decode. It processes only the newest token. For too, the first attention layer computes one query, one key, and one value [−1.5, 1.5]. too sits at position 3, so its query and key are turned by 3 radians, to [0.849, −1.131] and [−0.212, −1.485]. The key and value join the cache, and the query is compared with all four cached keys:

These four scores are the last row of the score matrix from section 4. The other rows stay as prefill left them, because the causal mask stops I, like, and cats from reading too. cats still gives 52.6% of its attention to itself, as it did in section 5, so decode computes only the new row.

The two phases are slow for different reasons. Prefill multiplies each weight matrix by one row per prompt token, so each weight read from memory is used once per token, and with a long prompt the arithmetic sets the speed. Decode has a single row, so every weight is read for one row's worth of work. Llama 3 8B reads all 8 billion weights, 16 GB at 2 bytes each, plus the cache, to produce each new token, and the speed of reading memory sets the pace.

14. Computing attention in small pieces to save memory

The 4 × 4 score matrix from section 4 has to be stored somewhere while it is computed. Models run on a GPU, a chip built for many multiplications at once. A GPU has a large main memory, which is slow to read and write, and a small, fast memory on the chip itself. Standard attention computes the whole score matrix, writes it to main memory, reads it back to apply softmax, writes the attention weights, and reads them again to multiply by V. With 4 tokens, that is 16 scores. With 32,000 tokens, it is about a billion scores per head, and moving them in and out of memory takes longer than the arithmetic itself.

FlashAttention (Dao et al., 2022) is a way of computing attention that never builds the full matrix. It splits Q, K, and V into groups of rows small enough for the fast on-chip memory and computes one tile of scores at a time, the scores of one group of queries against one group of keys. It keeps a running result for each query row. We use groups of 2 tokens, so the 4 × 4 score matrix becomes four 2 × 2 tiles.

Softmax divides each weight by the total of its whole row, but after the first tile only half of the row has been seen. FlashAttention solves this by keeping three running values for each query row:

• m: the largest score seen so far,
• ℓ: the sum of e^{score − m} so far,
• o: the sum of e^{score − m} × v so far.

Subtracting m does not change the softmax, because the same factor e^−m appears in every term and in the total. It does keep the exponentials from overflowing, that is, from growing too large for the computer to store. When a new tile brings a larger score, m increases, and the old ℓ and o are multiplied by α = e^{m_old − m_new} so that they use the new m. After the last tile, o / ℓ is the attention output.

For cats, the first tile covers the keys of I and like. Its scaled scores are 0.817 and −1.084, so m = 0.817, ℓ = e⁰ + e^−1.901 = 1 + 0.149 = 1.149, and o = 1 × [−0.5, 0.5] + 0.149 × [0.5, 1.5] = [−0.425, 0.724].

The second tile covers cats and too. The score for too is masked (−∞), so it adds e^−∞ = 0. The score for cats is 1.061, which is larger than the old m, so m becomes 1.061 and the old sums shrink by α = e^{0.817 − 1.061} = 0.784:

The result is the same [0.61, −0.49] that section 6 computed with the full matrix. The tile above the diagonal, where I and like would read cats and too, is entirely masked, so FlashAttention skips it without computing anything. In causal attention, that skip saves almost half of the tiles.

FlashAttention does the same arithmetic as standard attention, plus a little extra for the rescaling. It is faster because it reads and writes far less memory. Doubling the number of tokens doubles the memory it needs, where standard attention would need four times as much.