Next Word Prediction Using Attention and Transformer

Direct access and the end of recurrence

Author

Gian C

Published

July 27, 2026

Overview

It is definitely possible to build a capable next-word predictor with RNN and LSTM and GRU. However, Vaswani et al. (2017) proposed a new architecture for processing sequential data called the transformer. It removed recurrence entirely, eliminating the sequential computation bottleneck and making it possible to process a sequence in parallel during training, which results in more efficient utilization of GPUs.

Trained Data Walkthrough

Unlike the previous articles in this series, from my experience, I think the transformer is best explained through a trained example. Without that context, all the formulas can feel hand-wavy, making it difficult to develop an intuitive understanding.

There are a few variations in transformer architecture, but the most basic workflow looks like this:

Embed (+ Position) → Project → Score (+ Scale) → Mask → Normalize → Mix → Output

The core of this process is called self-attention. Through the demonstration, I hope the reasoning behind the nomenclature should become clear.

Goal

We’ve seen the rise of the machine in a previous example, but the sentence is too short to show the intricacies of the workflow. It’s been lengthened so we can demonstrate how the model picks up on the cues from the previous context:

the rise of the machine was slow but the fall of the machine was fast

The model will use “rise” and “fall” to produce the adjectives “slow” and “fast.”

Notice that “was”” appears twice. The first must be followed by “slow”, the second by “fast”.

Training Strategy

A transformer model trained on just this one sentence would overfit. The model would simply memorize the sequence, and the result would not be eloquent enough to showcase its capabilities. Therefore, the model is trained on a variation of sentences:

  • the rise was slow but the fall was fast
  • the rise was slow but the fall of the machine was fast
  • the rise of the machine was slow but the fall was fast
  • the rise of the machine was slow but the fall of the machine was fast
  • the fall was fast but the rise was slow
  • the fall was fast but the rise of the machine was slow
  • the fall of the machine was fast but the rise was slow
  • the fall of the machine was fast but the rise of the machine was slow

Swapping the clause order matters as much as varying the length, because we need to make sure that model reliably produces adjective based on the previous noun regardless of the order.

Prediction

In RNNs and LSTMs, data is processed one token (a word, or part of a word) at a time, so the recurrence implicitly captures the position and order of each word. In a transformer, positions are processed in parallel, so nothing about the architecture tells the model where a token sits in the sequence. We need additional information to distinguish the position of each word. One approach is to use positional embeddings.

Embedding & Position Matrices

One row per vocabulary token

\[ E = \begin{aligned} \text{the} \\ \text{rise} \\ \text{of} \\ \text{machine} \\ \text{was} \\ \text{slow} \\ \text{but} \\ \text{fall} \\ \text{fast} \\ \end{aligned} \begin{bmatrix} -0.24 & -1.53 & 0.51 & 1.22 & 0.30 & 0.08 \\[-0.6pt] -0.19 & -0.44 & 1.47 & -1.43 & 0.78 & 0.59 \\[-0.6pt] -0.26 & 0.58 & 0.32 & -0.16 & -1.55 & -2.06 \\[-0.6pt] 0.49 & 0.82 & 4.62 & -1.45 & -3.26 & 1.93 \\[-0.6pt] 1.34 & 0.16 & -1.17 & -0.47 & 2.12 & -1.04 \\[-0.6pt] 1.16 & -1.52 & -2.11 & -1.29 & -1.56 & 0.13 \\[-0.6pt] 0.85 & 1.55 & -0.67 & -0.60 & -0.28 & 0.35 \\[-0.6pt] -1.29 & -0.12 & 0.72 & -0.98 & 1.65 & -0.61 \\[-0.6pt] -2.10 & 0.30 & -1.52 & 0.17 & -1.00 & 0.39 \\[-0.6pt] \end{bmatrix} \]

One row per position

\[ P = \begin{aligned} \text{the} \\ \text{rise} \\ \text{of} \\ \text{the} \\ \text{machine} \\ \text{was} \\ \text{slow} \\ \text{but} \\ \text{the} \\ \text{fall} \\ \text{of} \\ \text{the} \\ \text{machine} \\ \text{was} \\ \end{aligned} \begin{bmatrix} 0.18 & 0.18 & 0.16 & -0.66 & 0.38 & -0.08 \\[-0.6pt] -0.26 & -0.69 & 0.07 & -0.89 & -0.02 & -0.19 \\[-0.6pt] -0.33 & -0.36 & -0.49 & 1.20 & 1.01 & 0.57 \\[-0.6pt] -0.30 & -0.54 & -0.01 & -0.51 & 0.90 & 0.56 \\[-0.6pt] 0.38 & 0.99 & 0.38 & -0.37 & -0.04 & -0.24 \\[-0.6pt] 0.71 & -0.35 & -0.07 & 1.24 & 0.53 & -0.18 \\[-0.6pt] -0.61 & -0.55 & 0.43 & -1.23 & 0.37 & 0.23 \\[-0.6pt] 0.12 & 0.24 & -0.53 & 0.43 & -0.05 & -1.67 \\[-0.6pt] 0.00 & -0.72 & 0.48 & 0.94 & -0.15 & -0.21 \\[-0.6pt] -0.04 & 0.22 & 0.26 & -1.00 & 0.44 & 0.70 \\[-0.6pt] 0.37 & 0.44 & -0.49 & -0.08 & 0.34 & -1.10 \\[-0.6pt] -0.31 & 1.34 & -0.17 & -0.32 & 0.42 & 1.14 \\[-0.6pt] 1.58 & 0.74 & 1.98 & -1.12 & -1.79 & 0.58 \\[-0.6pt] 0.76 & 0.64 & -0.54 & -0.23 & 0.79 & 0.04 \\[-0.6pt] \end{bmatrix} \]

\(P\) is the parameter to learn the position of the words.

“fast” is not in positional parameter, because that’s the target and not the input.

To turn the input into a representation carrying both token and positional information, we add the two vectors. The token at position 6 is “slow,” so its input vector is \(E_{\text{slow}} + P_6\). In general:

\[ x_i = E(\text{token}_i)+P_i \]

Make transformer input

3Blue1Brown put the logic quite eloquently. A positional embedding can have any number of dimensions. However, imagine a 3-dimensional space where a vector represents the meaning of a word. Similar words would tend to be closer together, while words with very different meanings would tend to be farther apart. This representation is useful because it allows the model to encode semantic relationships as patterns in the embedding space.

The result will occupy 6-dimensional space, so \(X\) can be expressed as:

\(_\text{token}\backslash^\text{coordinate}\) 0 1 2 3 4 5
0 the -0.06 -1.35 0.67 0.56 0.68 0.
1 rise -0.45 -1.13 1.54 -2.32 0.76 0.4
2 of -0.59 0.22 -0.17 1.04 -0.54 -1.49
3 the -0.54 -2.07 0.5 0.71 1.2 0.64
4 machine 0.87 1.81 5. -1.82 -3.3 1.69
5 was 2.05 -0.19 -1.24 0.77 2.65 -1.22
6 slow 0.55 -2.07 -1.68 -2.52 -1.19 0.36
7 but 0.97 1.79 -1.2 -0.17 -0.33 -1.32
8 the -0.24 -2.25 0.99 2.16 0.15 -0.13
9 fall -1.33 0.1 0.98 -1.98 2.09 0.09
10 of 0.11 1.02 -0.17 -0.24 -1.21 -3.16
11 the -0.55 -0.19 0.34 0.9 0.72 1.22
12 machine 2.07 1.56 6.6 -2.57 -5.05 2.51
13 was 2.1 0.8 -1.71 -0.7 2.91 -1.

We see that each token is represented by vector of 6 numbers, therefore: \[ d_{\text{model}} = 6 \]

Another way to inpterpret it is that each token has 6 features.

For context, the largest GPT-3 model uses \(d_{\text{model}} = 12{,}288\); the smallest variant uses 768.

Projection

You might be familiar with dictionary in the context of programming. For example, in Python a dictionary is just next_word_map={"the": "rise", "rise": "of" ...}, one of the keys being the and value being rise. To query the value, we use next_word_map["the"].

In a transformer, the concept is similar, but with mathematical operations in between to match a query with a key and deliver value by approximation rather than static encoding. During this process, the previous tokens create context for next prediction. Which means even though “was” needs to produce both “slow” and “fast”, it’s still possible based on context cue “rise” and “fall”.

The same input \(X\) is projected three times, through three separate learned matrices, into three different roles (request - \(Q\), an advertisement - \(K\), and a payload - \(V\)).

Query

A query is the vector that represents what a token is currently trying to retrieve from the sequence.

\[ q_i = x_i W_Q, \quad Q=XW_Q \]

Trained query parameter:

\[ W_Q = \begin{bmatrix} 0.83 & 0.88 & -2.05 & -1.02 & 1.94 & -0.86 \\ 0.27 & -0.20 & -0.66 & -0.37 & 0.49 & -0.04 \\ 0.27 & 0.59 & -0.05 & 0.24 & -0.43 & 0.97 \\ -1.65 & -1.29 & 0.42 & 1.13 & -1.36 & -1.41 \\ 0.05 & 0.22 & 0.46 & -0.07 & 0.31 & 0.05 \\ -0.31 & -0.10 & -0.08 & -0.60 & -0.08 & -0.32 \end{bmatrix} \]

Since this is just simple matrix multiplication, we’ll fast-track to final result \(Q\):

\(_\text{token}\backslash^\text{coordinate}\) 0 1 2 3 4 5
0 the -1.12 0.04 1.53 1.31 -1.62 -0.
1 rise 3.48 3.86 0.93 -1.67 1.27 5.11
2 of -1.76 -1.97 1.38 2.59 -2.43 -0.68
3 the -2.18 -0.48 3.25 1.77 -2.92 -0.11
4 machine 4.87 4.81 -5.65 -3.2 1.74 5.89
5 was 0.56 0.82 -2.38 -0.9 4.29 -3.52
6 slow 3.43 2.86 -1.31 -3.18 3.8 1.36
7 but 1.64 0.07 -3.23 -1.32 3.51 -1.42
8 the -4.06 -1.92 2.91 3.82 -4.87 -1.74
9 fall 2.53 2.39 2.73 -0.88 0.38 4.96
10 of 1.64 0.15 -1.29 1.18 0.99 0.99
11 the -2.24 -1.37 1.85 0.95 -2.4 -0.81
12 machine 7.13 7.36 -9.21 -5.16 3.67 7.13
13 was 3.11 2.32 -3.62 -3.24 7.14 -2.04

By skimming over the matrix, “was” at two different position has different vector, which means the result they retrieve will also be different.

Key

A key is the vector that represents how a token can be matched against queries.

\[ k_i = x_i W_K, \quad K=XW_K \]

Trained key parameter:

\[ W_K = \begin{bmatrix} -0.03 & -0.56 & -0.27 & 0.60 & -1.54 & 1.02 \\ -0.66 & -0.77 & 0.39 & 0.72 & 0.24 & -1.38 \\ 0.07 & 0.36 & 0.07 & 0.00 & 0.63 & 0.24 \\ -0.21 & -1.08 & 1.03 & 1.32 & -1.24 & 1.40 \\ 0.61 & 0.19 & 0.47 & -0.30 & 1.56 & 0.45 \\ 0.66 & 0.07 & 0.28 & -0.54 & 1.04 & 0.46 \end{bmatrix} \]

\(_\text{token}\backslash^\text{coordinate}\) 0 1 2 3 4 5
0 the 1.24 0.84 0.43 -0.47 0.56 3.05
1 rise 2.08 4.35 -2.13 -4.59 5.87 -1.25
2 of -1.67 -1.23 0.63 2.14 -2.83 -0.42
3 the 2.42 1.58 0.85 -1.58 2.31 4.25
4 machine -1.39 1.38 -2.13 -0.5 1.11 -3.67
5 was 0.63 -1.86 0.98 1.97 -2.07 3.76
6 slow 1.27 3.2 -4.13 -4.32 -0.76 -0.88
7 but -2.33 -2.33 -0.35 2.46 -3.5 -2.76
8 the 1.11 -0.09 1.52 1.11 -2.13 6.13
9 fall 1.79 3.56 -0.57 -4.02 8.5 -3.05
10 of -3.46 -1.1 -1.34 2.55 -4.91 -3.67
11 the 1.22 -0.17 1.71 -0.15 2.29 1.93
12 machine -1.51 2.01 -3.81 -0.87 -0.74 -3.17
13 was 0.55 -1.17 -0.01 0.58 0.25 0.5

Notice how “rise” and “fall” in coordinate 4 is holding the largest value. It’s fair to assume that this column is doing majority of the work to discern the context cue for the following adjectives.

The length of each vector representing each token in \(Q\) and \(K\) is known as \(d_k\). Intuitively, it can be viewed as “features used to determine the relevance between query and key”

Here \(d_k = 6\), the same as \(d_\text{model}\), but they don’t have to be the same size in practice.

Score and Scale

Every query is compared against every key. Similar to next_word_map["the"], the representation here is \(q_i \cdot k_j\), with \(j\) referring to the column index.

\[ s_{ij} = \frac{q_i \cdot k_j}{\sqrt{d_k}}, \quad S = \frac{QK^\top}{\sqrt{d_k}} \]

The division by \(\sqrt{d_k}\) is used to prevent the score inflation, so later softmax does not push any number into its saturated region, where one weight sits at nearly 1, the rest at nearly 0, and the gradient is almost flat. Dividing by \(\sqrt{d_k}\) cancels the dimensional growth so the scores land at a workable scale.

To make the computation process perfectly clear, here’s how the result for the number in position (0, 0) is calculated.

\[ \underbrace{ \begin{bmatrix} -1.12 & 0.04 & 1.53 & 1.31 & -1.62 & -0. \end{bmatrix} } _{q_0} \underbrace{ \begin{bmatrix} 1.24 \\ 0.84 \\ 0.43 \\ -0.47 \\ 0.56 \\3.05 \end{bmatrix} } _{k_0} = -0.2202 , \quad \frac{-0.2202}{\sqrt{6}} = -0.9 \]

\(_i\backslash^j\) the rise of the mac. was slow but the fall of the mac. was
0 the -0.9 -8.54 4.15 -2.92 -1.67 2.72 -4.91 4.43 2.43 -8.87 5.33 -1.09 -1.63 -0.13
1 rise 10.22 12.56 -7.87 17.4 -7.34 3.76 5.99 -16.35 12.94 8.72 -19.1 7.42 -6.85 -0.29
2 of -3.22 -16.52 7.73 -7.68 -1.92 4.69 -9.39 10.18 1.7 -16.28 11.19 -2.75 -1.98 0.77
3 the -1.84 -15.8 7.51 -5.44 -3.38 4.84 -9.42 8.15 4.09 -15.95 9.39 -1.73 -3.69 -0.17
4 mac. 11.46 24.76 -13.01 19.91 -2.52 0.33 21.31 -20.74 10.32 15.81 -21.6 4.63 2.68 -0.56
5 was -3.1 17.77 -6.54 -1.23 9.61 -11.2 6.9 -4.03 -14.19 22.89 -4.11 -0.14 7.61 -0.75
6 slow 5.65 23.52 -11.52 12.78 1.15 -5.51 11.67 -15.95 -0.7 23.69 -18.39 5.42 0.48 -0.68
7 but -0.44 15.92 -6.94 2.23 5.91 -7.15 8.13 -5.89 -8.46 18.15 -6.84 0.81 5.31 0.1
8 the -6.2 -27.35 13.75 -14.32 -1.71 6.12 -14.14 18.02 1.64 -27.44 21.35 -6.03 -1.23 0.05
9 fall 9.02 4.06 -4.28 14.54 -9.53 6.51 -0.51 -12.09 14.43 1.3 -15.26 7.31 -10.08 0.26
10 of 1.88 2.45 -1.81 3.16 -0.98 1.42 0.49 -2.86 2.09 1.99 -3.91 1.53 -0.87 0.88
11 the -3.02 -13.07 6.44 -6.75 -1.18 2.76 -6.7 8.47 0.66 -12.94 9.8 -2.68 -1.17 -0.04
12 mac. 15.21 41.98 -20.91 27.79 0.16 -3.77 34.24 -30.92 9.58 30.37 -31.74 5.99 7.42 -1.28
13 was 1.42 34.15 -14.95 8.59 9.66 -14.22 15.01 -15.78 -13.7 39.11 -18.06 4.13 7.26 -0.85

Notice in row 13 how “rise” and “fall” have much larger dot products than anything else in that row. The projections have learned that the representation of “was”” at position 13 strongly matches the feature configuration of those two nouns.

One issue is that the model is looking ahead. “rise” at row 1 shouldn’t be able to see “of”” or anything after it. To address this, we apply a causal mask.

The Causal Mask

Causal masking is the lower-triangular attention mask that stops each position attending to positions to its right. In other words, a position may not look ahead.

\[ S^{\text{masked}}_{ij} = \begin{cases} S_{ij} & j \le i \\ -\infty & j > i \end{cases} , \quad M = \underbrace{ \begin{bmatrix} 0 & -\infty & -\infty & \cdots & -\infty & -\infty \\ 0 & 0 & -\infty & \cdots & -\infty & -\infty \\ \vdots & \vdots & \vdots & \ddots & \vdots & \vdots \\ 0 & 0 & 0 & \cdots & 0 & -\infty \\ 0 & 0 & 0 & \cdots & 0 & 0 \\ \end{bmatrix} }_\text{14x14 causal mask matrix} \]

In a mathematical sense, a triangular matrix can be constructed by setting everything above the diagonal to \(-\infty\), and everything on or below it to zero. \(-\infty\) because \(e^{-100{\dots}}\) is so small that it’s practically zero. It is then applied by adding M to the score matrix:

NOTE: some models use self-attention that looks in both directions, with no causal mask. That’s called encoder-style self-attention, and it’s what BERT uses. Our example is decoder-style, the kind used in GPT.

\(_i\backslash^j\) the rise of the mac. was slow but the fall of the mac. was
0 the -0.9 \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\)
1 rise 10.22 12.56 \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\)
2 of -3.22 -16.52 7.73 \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\)
3 the -1.84 -15.8 7.51 -5.44 \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\)
4 mac. 11.46 24.76 -13.01 19.91 -2.52 \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\)
5 was -3.1 17.77 -6.54 -1.23 9.61 -11.2 \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\)
6 slow 5.65 23.52 -11.52 12.78 1.15 -5.51 11.67 \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\)
7 but -0.44 15.92 -6.94 2.23 5.91 -7.15 8.13 -5.89 \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\)
8 the -6.2 -27.35 13.75 -14.32 -1.71 6.12 -14.14 18.02 1.64 \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\)
9 fall 9.02 4.06 -4.28 14.54 -9.53 6.51 -0.51 -12.09 14.43 1.3 \(-\infty\) \(-\infty\) \(-\infty\) \(-\infty\)
10 of 1.88 2.45 -1.81 3.16 -0.98 1.42 0.49 -2.86 2.09 1.99 -3.91 \(-\infty\) \(-\infty\) \(-\infty\)
11 the -3.02 -13.07 6.44 -6.75 -1.18 2.76 -6.7 8.47 0.66 -12.94 9.8 -2.68 \(-\infty\) \(-\infty\)
12 mac. 15.21 41.98 -20.91 27.79 0.16 -3.77 34.24 -30.92 9.58 30.37 -31.74 5.99 7.42 \(-\infty\)
13 was 1.42 34.15 -14.95 8.59 9.66 -14.22 15.01 -15.78 -13.7 39.11 -18.06 4.13 7.26 -0.85

Normalize

Applying softmax across each row gives a probability distribution over positions, we’ve seen this in the previous articles:

\[ a_{ij} = \frac{e^{s_{ij}}}{\sum_{l} e^{s_{il}}}, \quad A = \text{softmax}_{\text{row}}\!\left(\frac{QK^\top}{\sqrt{d_k}} + M\right) \]

Plugging in the raw numbers produces absurdly large values, such as \(e^{39.11} \approx 10^{17}\). Softmax is therefore evaluated in practice by subtracting the row maximum first, which cancels from the ratio and puts the largest exponent at \(e^0 = 1\):

\[ a_{ij} = \frac{e^{\,s_{ij} - \max_l s_{il}}}{\sum_{l} e^{\,s_{il} - \max_l s_{il}}} \]

\(_i\backslash^j\) the rise of the mac. was slow but the fall of the mac. was
0 the 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
1 rise 0.09 0.91 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
2 of 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
3 the 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
4 mac. 0. 0.99 0. 0.01 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
5 was 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
6 slow 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
7 but 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
8 the 0. 0. 0.01 0. 0. 0. 0. 0.99 0. 0. 0. 0. 0. 0.
9 fall 0. 0. 0. 0.53 0. 0. 0. 0. 0.47 0. 0. 0. 0. 0.
10 of 0.1 0.18 0. 0.37 0.01 0.07 0.03 0. 0.13 0.12 0. 0. 0. 0.
11 the 0. 0. 0.03 0. 0. 0. 0. 0.2 0. 0. 0.77 0. 0. 0.
12 mac. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
13 was 0. 0.01 0. 0. 0. 0. 0. 0. 0. 0.99 0. 0. 0. 0.

Row 5 puts \(0.99971\) on “rise”; row 13 puts \(0.99304\) on “fall”. The positional embedding led two identical words point at different cues.

In row 13, note how decisive the weights are. fall led rise by only 5 units in score space. Softmax converted small gaps into near-total dominance.

Values

A value is the information that gets actually aggregated and passed forward once the attention weights have been computed.

\[ v_i = x_i W_V, \quad V = XW_V \]

\[ W_V = \begin{bmatrix} -0.18 & 0.49 & -0.45 & 0.26 & -0.22 & -0.54 \\ 0.07 & 0.03 & 0.58 & 0.33 & -0.12 & 0.11 \\ 0.07 & -0.23 & 0.49 & 0.25 & -0.81 & 0.06 \\ -0.21 & -0.24 & 0.50 & -0.20 & 0.09 & 0.16 \\ -0.95 & 0.66 & 0.21 & -0.09 & -0.08 & 0.19 \\ 0.03 & 0.16 & -0.58 & 0.25 & 0.09 & 0.19 \end{bmatrix} \]

Applying that to the input we get \(V\):

\(_\text{token}\backslash^\text{coordinate}\) 0 1 2 3 4 5
0 the -0.8 0.09 -0. -0.47 -0.37 0.14
1 rise -0.11 0.51 -0.93 0.39 -1.25 0.06
2 of 0.36 -1.09 1.58 -0.66 0.24 0.11
3 the -1.28 0.28 -0.48 -0.79 -0.01 0.56
4 machine 3.89 -2.14 0.53 3.16 -4.21 -0.57
5 was -3.18 2.65 0.01 -0.54 0.32 -0.81
6 slow 1.31 0.47 -3.99 -0.26 1.39 -1.19
7 but 0.18 0.42 0.62 0.28 0.44 -0.74
8 the -0.65 -0.85 0.48 -1.04 -0.31 0.29
9 fall -1.25 0.99 0.53 0.16 -0.85 0.89
10 of 1.14 -1.12 1.92 -0.31 -0.22 -0.83
11 the -0.73 0.1 0.2 -0.06 0. 0.81
12 machine 5.61 -2.77 -0.59 4.3 -5.59 -1.44
13 was -3.09 3.37 -0.48 0.01 0.44 -0.9

In value matrix, coordinate 2 signals the cue word since they have the largest difference betwen “rise” and “fall”.

Key coordinate 4 is large for both rise (\(5.87\)) and fall (\(8.50\)), enough to identify the cue but not enough to tell them apart. Value coordinate 2 separates them, \(-0.93\) against \(0.53\).

Mix

Mixing combines information across positions, so the model can capture dependencies regardless of distance.

\[ z_i = \sum_{j} a_{ij}\, v_j, \quad Z=AV \]

\(_\text{token}\backslash^\text{coordinate}\) 0 1 2 3 4 5
0 the -0.8 0.09 -0. -0.47 -0.37 0.14
1 rise -0.17 0.48 -0.85 0.32 -1.17 0.07
2 of 0.36 -1.09 1.58 -0.66 0.24 0.11
3 the 0.36 -1.09 1.58 -0.66 0.24 0.11
4 machine -0.12 0.51 -0.93 0.38 -1.24 0.06
5 was -0.11 0.51 -0.93 0.39 -1.25 0.06
6 slow -0.11 0.51 -0.93 0.39 -1.25 0.06
7 but -0.11 0.51 -0.93 0.39 -1.25 0.06
8 the 0.18 0.4 0.64 0.26 0.43 -0.73
9 fall -0.98 -0.25 -0.03 -0.9 -0.15 0.43
10 of -0.96 0.38 -0.32 -0.41 -0.37 0.28
11 the 0.92 -0.8 1.64 -0.2 -0.07 -0.78
12 machine -0.11 0.51 -0.93 0.39 -1.25 0.06
13 was -1.24 0.99 0.52 0.16 -0.85 0.88

Position 13 now carries a dependency on “fall”’s value. This is the core benefit of attention: a position can pick up information from tokens arbitrarily far back.

\[ \begin{aligned} z_{13} &= (-1.2441,\; 0.9916,\; 0.5232,\; 0.1642,\; -0.8533,\; 0.8797) \\ v_{9} &= (-1.2520,\; 0.9949,\; 0.5334,\; 0.1626,\; -0.8505,\; 0.8854) \end{aligned} \]

The attention function can be written as:

\[ \text{Attention}(Q, K, V) = \text{softmax}_{\text{row}}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right) V \]

We’ve seen that QKV are generated directly from the input \(X\), hence the term self-attention. A functioning model typically has multiple attention projections, known as heads, and this is typically referred to as multi-head attention.

Output

\(z_{13}\) essentially holds “fall” value. But the vector needs to carry both the cue and the token “was” at position 13. To do this we form the hidden state with a residual connection:

\[ \text{hidden}_i = x_i + z_i W_O \]

\(W_O\) maps the head’s output back into the model’s coordinate system before the addition.

NOTE: \(W_O\) is especially useful in multi-head attention. Each head produces its own \(z\) of width \(d_k\), and \(W_O\) transforms the concatenated result back to \(d_{model}\).

\[ W_O = \begin{bmatrix} 0.58 & 0.05 & 0.28 & -0.75 & 1.00 & -0.54 \\ 0.11 & 0.30 & 0.50 & -0.95 & 0.18 & -0.69 \\ 0.62 & 0.46 & -0.90 & 0.81 & 0.83 & -0.96 \\ -0.01 & -0.32 & 0.15 & 0.06 & 0.07 & 0.62 \\ -0.83 & 0.41 & -0.12 & 1.01 & 0.15 & 0.88 \\ -0.39 & 0.81 & -0.78 & -0.02 & 0.84 & -0.36 \end{bmatrix} \]

The resulting \(\text{hidden}\) matrix:

\(_\text{token}\backslash^\text{coordinate}\) 0 1 2 3 4 5
0 the -0.26 -1.25 0.36 0.66 -0.08 -0.29
1 rise -0.08 -1.91 2.63 -4.49 -0.13 0.12
2 of 0.24 1.04 -2.25 3.29 1.02 -2.68
3 the 0.29 -1.25 -1.58 2.96 2.76 -0.55
4 machine 1.28 0.95 6.21 -4.19 -4.2 1.42
5 was 2.47 -1.06 -0.02 -1.62 1.75 -1.5
6 slow 0.97 -2.94 -0.46 -4.91 -2.09 0.08
7 but 1.39 0.92 0.03 -2.56 -1.23 -1.59
8 the 0.22 -2.32 1.22 2.63 0.4 -0.31
9 fall -1.97 0.54 0.15 -1.24 1.32 -0.03
10 of -0.4 1.15 -0.2 -0.55 -2.21 -3.29
11 the 1.28 -0.23 -0.69 2.23 2.18 -0.2
12 machine 2.49 0.69 7.83 -4.96 -5.95 2.24
13 was 2.18 1.59 -2.59 -1.15 2.9 -2.48

Feedforward Network

What we’ve seen so far is the attention mechanism. Functional attention-only transformers do exist. However, every production transformer since Vaswani et al. (2017) seems to pairs each attention block with a feed-forward network, the same densely-connected structure we built in the perceptron section.

If you ask an LLM, it’ll tell you the FFN introduces “non-linearity”. It’s a nonsense wrapped in some truth, because softmax also applies non-linearity. Here are some actual theories from various research:

  • Memory capacity: Models store roughly 2 bits of knowledge per parameter, meaning a 7B model can hold more than English Wikipedia and textbooks combined. Most of those parameters sit in the feed-forward layers.
  • Key-value memory: The first layer detects patterns in the input, the second maps them to distributions over the vocabulary.
  • Factual storage: Specific facts appear to live in mid-layer feed-forward weights, and can be edited there with rank-one updates.
  • Preventing collapse: Without skip connections or MLPs, deep attention converges to a rank-1 matrix, making every token identical.
  • Sparse feature computation: In trained models only around 3% of feed-forward units fire for a given token, and larger models are sparser. The layer appears to represent more features than it has neurons, with the activation function filtering the interference that creates.

Our model includes one for learning purposes, to gain familiarity with the standard architecture. A typical feed-forward formula looks like this:

\[ y_i = \text{hidden}_i + \text{FFN}(\text{hidden}_i) ,\quad \text{FFN}(\text{hidden}) = \text{ReLU}(\text{hidden}W_1 + b_1)W_2 + b_2 \]

Trained parameters:

\[ \begin{aligned} W_1 &= \begin{bmatrix} -0.51 & -0.57 & -0.12 & 0.34 & -0.13 & -1.60 & 0.66 & -0.86 & 0.87 & -1.27 & -0.61 & 0.61 \\ 0.65 & 0.14 & -1.11 & -1.05 & 0.22 & 0.72 & -0.73 & 1.38 & 0.46 & -0.45 & -0.94 & -0.56 \\ 0.81 & 0.87 & -0.54 & 0.12 & -0.23 & 0.79 & 0.19 & 0.17 & -0.14 & 0.10 & -0.09 & 0.88 \\ 0.01 & -0.87 & -0.98 & -0.24 & -0.42 & 0.25 & 0.98 & -1.10 & -0.48 & -0.29 & 0.77 & -0.09 \\ -0.39 & 0.79 & -1.79 & 0.68 & -2.11 & -1.53 & 0.19 & -0.88 & 1.34 & -0.48 & -0.11 & -0.27 \\ -1.50 & -0.12 & 0.23 & 0.01 & 0.39 & 0.65 & 0.57 & -0.29 & -1.89 & -0.25 & -0.57 & -0.50 \end{bmatrix} \\ W_2 &= \begin{bmatrix} 1.36 & 0.38 & -0.06 & -0.25 & -1.05 & 0.15 \\ 1.52 & -0.34 & 0.41 & -0.91 & -0.15 & 0.50 \\ -0.26 & -0.36 & -0.63 & 1.30 & 0.12 & -0.35 \\ -0.25 & 0.33 & 0.65 & -0.56 & 0.73 & 0.85 \\ -0.21 & 0.43 & -0.84 & -0.03 & 0.30 & -0.89 \\ -0.13 & -0.61 & -0.65 & -1.11 & 0.20 & 0.44 \\ -0.21 & 0.48 & -0.59 & -0.57 & -0.16 & 1.13 \\ 1.44 & -0.40 & -0.54 & 1.29 & -0.21 & -0.24 \\ 0.81 & 1.07 & -0.05 & -0.68 & 0.75 & -0.81 \\ 0.06 & -1.15 & -0.33 & 0.70 & -0.44 & 0.14 \\ -0.82 & -0.34 & -0.09 & 1.66 & -0.35 & 0.76 \\ -1.35 & 0.13 & -0.14 & -0.92 & -1.45 & 0.14 \end{bmatrix} \\ b_1 &= \begin{bmatrix} 0.76 & 0.53 & 0.14 & 0.52 & -0.08 & -0.60 & 0.18 & 0.75 & -0.12 & 0.37 & 0.14 & 0.91 \end{bmatrix} \\ b_2 &= \begin{bmatrix} 0.01 & 0.07 & 0.14 & 0.21 & -0.44 & 0.65 \end{bmatrix} \end{aligned} \]

After matrix operation, \(y\) becomes:

\(_\text{token}\backslash^\text{coordinate}\) 0 1 2 3 4 5
0 the -3.89 -1.88 -0.65 2.44 -4.43 5.36
1 rise 8.53 -7.59 -0.24 -3.17 -7.87 4.44
2 of 6.37 7.87 -3.49 3.47 0.31 -2.63
3 the -1.7 3.85 -2.44 3.03 4.86 7.39
4 machine 15.93 -8.13 -17.64 -3.82 -18.94 -2.7
5 was 6.24 8.51 2.39 -14.03 4.19 -1.47
6 slow -2.09 -8.12 -12.77 9.65 -4.22 -4.46
7 but 14. 4.63 -7.46 1.24 -6.1 -7.77
8 the -8.67 -0.33 -0.59 2.77 -5.54 11.89
9 fall 12. -3.61 -1.01 -0.62 -2.23 2.47
10 of 16.03 2.19 -11.33 6.12 -11.58 -8.16
11 the 0.92 5.32 -1.65 -1.35 3.9 4.61
12 machine 13.11 -8.9 -23.49 -3.83 -22.25 -5.47
13 was 14.49 15.11 -1.83 -10.93 11.2 -9.68

The final vector is projected into a probability distribution:

\[ \text{logits} = y_i W_U, \quad \hat{y} = \text{softmax}(\text{logits}) \]

Trained parameters:

\[ W_U = \begin{bmatrix} 0.63 & -0.88 & 0.78 & 0.20 & 0.77 & 0.19 & -0.31 & -0.66 & 0.16 \\ 1.17 & -0.48 & -0.87 & 0.44 & -0.68 & 0.02 & -0.64 & -0.24 & 0.21 \\ -0.27 & 0.43 & 0.43 & -0.85 & -0.11 & 0.47 & -0.69 & 0.54 & -0.05 \\ 1.58 & 0.50 & -0.26 & 0.42 & -0.36 & -1.33 & 0.99 & 0.59 & 0.30 \\ -0.51 & -0.84 & -0.45 & 0.98 & -0.64 & -0.21 & 0.11 & -1.01 & 1.10 \\ -0.94 & 0.70 & 0.10 & 0.74 & 0.01 & -0.63 & -0.83 & 0.77 & -1.87 \\ \text{"the"} & \text{"rise"} & \text{"of"} & \text{"machine"} & \text{"was"} & \text{"slow"} & \text{"but"} & \text{"fall"} & \text{"fast"} \end{bmatrix} \]

The distribution for final prediction \(\hat{y}\) is:

\(_\text{position}\backslash^\text{token}\) the rise of mac. was slow but fall fast
0 the 0. 0.51 0. 0. 0. 0. 0. 0.49 0.
1 rise 0. 0. 0.5 0. 0.5 0. 0. 0. 0.
2 of 1. 0. 0. 0. 0. 0. 0. 0. 0.
3 the 0. 0. 0. 1. 0. 0. 0. 0. 0.
4 machine 0. 0. 0. 0. 1. 0. 0. 0. 0.
5 was 0. 0. 0. 0. 0. 1. 0. 0. 0.
6 slow 0. 0. 0. 0. 0. 0. 1. 0. 0.
7 but 1. 0. 0. 0. 0. 0. 0. 0. 0.
8 the 0. 0.51 0. 0. 0. 0. 0. 0.49 0.
9 fall 0. 0. 0.5 0. 0.5 0. 0. 0. 0.
10 of 1. 0. 0. 0. 0. 0. 0. 0. 0.
11 the 0. 0. 0. 1. 0. 0. 0. 0. 0.
12 machine 0. 0. 0. 0. 1. 0. 0. 0. 0.
13 was 0. 0. 0. 0. 0. 0. 0. 0. 1.

The main takeaway is that although some positions split their prediction near 50/50, the model is decisive where it counts: “machine” → “was” → “slow” at position 5, and “machine” → “was” → “fast” at position 13.

Wrap Up

This particular architecture is commonly known as a decoder-only transformer. It is the building block of the GPT family. We’ll explore how to train a Transformer model using real data in the next article. Along the way, we’ll cover essential topics that were not showcased in this walkthrough, such as multi-head attention, layer normalization, and rotary positional embeddings (RoPE). In addition, we’ll explore some of the nuances of PyTorch.

Additional Resources