Chapter 10.4 - Complete Inference Walkthrough
Let's do a complete, end-to-end math walkthrough of the GPT-2 inference process!
Our example sentence is: "Tomorrow I am flying to"
After tokenization, we get the following Token IDs:
[49488, 314, 716, 7348, 284]
1. Token Embeddings Lookup
Let's assume each token is represented by a 3-dimensional matrix. We look up each Token ID in the token embedding matrix Wₑ (which has a shape of 50257 × 3).
| Word | Token ID | Embedding Vector |
|---|---|---|
| Tomorrow | 49488 | [ 0.62, -0.15, 0.48] |
| I | 314 | [-0.31, 0.84, 0.12] |
| am | 716 | [ 0.17, 0.39, -0.56] |
| flying | 7348 | [ 0.91, -0.44, 0.27] |
| to | 284 | [-0.08, 0.71, 0.65] |
2. Positional Embeddings
Next, we look up the positional embeddings for each position in the sentence:
Position 0 → [ 0.10, 0.00, -0.10]
Position 1 → [ 0.20, -0.10, 0.05]
Position 2 → [ 0.30, 0.10, 0.00]
Position 3 → [ 0.40, -0.20, 0.10]
Position 4 → [ 0.50, 0.00, -0.05]
We add the Token Embeddings and Positional Embeddings together to create our Input Embeddings:
[ 0.62, -0.15, 0.48] + [ 0.10, 0.00, -0.10] = [ 0.72, -0.15, 0.38]
[-0.31, 0.84, 0.12] + [ 0.20, -0.10, 0.05] = [-0.11, 0.74, 0.17]
[ 0.17, 0.39, -0.56] + [ 0.30, 0.10, 0.00] = [ 0.47, 0.49, -0.56]
[ 0.91, -0.44, 0.27] + [ 0.40, -0.20, 0.10] = [ 1.31, -0.64, 0.37]
[-0.08, 0.71, 0.65] + [ 0.50, 0.00, -0.05] = [ 0.42, 0.71, 0.60]
Final Input Embeddings Matrix:
[
[ 0.72, -0.15, 0.38],
[-0.11, 0.74, 0.17],
[ 0.47, 0.49, -0.56],
[ 1.31, -0.64, 0.37],
[ 0.42, 0.71, 0.60]
]
3. Multi-Head Attention
The input matrix is passed through the Self-Attention mechanism. For a detailed breakdown of this step, revisit Chapter 7.3 - Transformer Story.
Assuming we run this through 3 Attention Heads, we get 3 context vectors:
# Head 1
[
[ 0.083, -0.260],
[ 0.026, 0.120],
[ 0.170, 0.212],
[ 0.134, 0.056],
[ 0.127, 0.013]
]
# Head 2
[
[-0.192, 0.381],
[ 0.114, 0.248],
[ 0.292, -0.071],
[ 0.056, 0.319],
[ 0.173, 0.201]
]
# Head 3
[
[ 0.214, 0.090],
[ 0.097, -0.161],
[ 0.333, 0.282],
[ 0.011, -0.084],
[ 0.205, 0.118]
]
We concatenate these side-by-side:
# [ Head 1 | Head 2 | Head 3 ]
[
[ 0.083, -0.260 , -0.192, 0.381 , 0.214, 0.090 ],
[ 0.026, 0.120 , 0.114, 0.248 , 0.097, -0.161 ],
[ 0.170, 0.212 , 0.292, -0.071 , 0.333, 0.282 ],
[ 0.134, 0.056 , 0.056, 0.319 , 0.011, -0.084 ],
[ 0.127, 0.013 , 0.173, 0.201 , 0.205, 0.118 ]
]
We multiply this by the Output Projection Matrix W_O (6×3) to get our final Attention Output:
[
[ 0.413, -0.149, 0.123 ],
[-0.066, 0.171, -0.027 ],
[ 0.296, 0.291, 0.042 ],
[ 0.291, -0.029, -0.002 ],
[ 0.243, 0.069, 0.048 ]
]
4. First Residual Connection & LayerNorm
We add the Attention Output back to the original Input Matrix to get the Residual Output:
[
[ 1.133, -0.299, 0.503],
[-0.176, 0.911, 0.143],
[ 0.766, 0.781, -0.518],
[ 1.601, -0.669, 0.368],
[ 0.663, 0.779, 0.648]
]
After normalizing this matrix (see Chapter 7.2), we get our 1st LayerNorm Output:
[
[ 1.40, -1.52, 0.12],
[-1.15, 1.29, -0.14],
[ 0.71, 0.70, -1.41],
[ 1.31, -1.14, -0.17],
[-0.16, 0.91, -0.75]
]
5. Feed Forward Network (FFN)
We pass the LayerNorm Output through the FFN (Linear → GELU → Linear) to get our FFN Output:
[
[ 1.228, -0.564, 0.583],
[-0.810, 1.420, -0.561],
[-0.210, 0.117, -0.015],
[ 0.836, -0.351, 0.390],
[-0.339, 0.563, -0.140]
]
6. Second Residual Connection & LayerNorm
We add the 1st LayerNorm Output and the FFN Output together:
# Residual Output 2
[
[ 2.628, -2.084, 0.703],
[-1.960, 2.710, -0.701],
[ 0.500, 0.817, -1.425],
[ 2.146, -1.491, 0.220],
[-0.499, 1.473, -0.890]
]
And finally, we apply our 2nd Layer Normalization:
[
[ 1.28, -1.16, -0.12],
[-1.02, 1.36, -0.34],
[ 0.48, 0.91, -1.39],
[ 1.30, -1.14, -0.16],
[-0.41, 1.38, -0.97]
]
If this is not the last layer, this matrix becomes the input for the next layer's attention mechanism, and the entire cycle repeats!
7. Next Word Generation (LM Head)
If this is the last block, we only care about the very last row in our matrix (the final token), which is [-0.41, 1.38, -0.97].
We multiply this vector by the LM Head Matrix. The LM Head has a shape of (Model Dimension) × (Vocabulary Size). In this example, let's pretend our vocabulary only has 5 words.
# Last Token Vector
[-0.41, 1.38, -0.97]
# Multiplied by LM Head (3x5)
[
[ 0.3, -0.2, 0.7, 0.1, -0.4],
[-0.5, 0.8, -0.1, 0.6, 0.2],
[ 0.4, 0.3, -0.6, 0.5, -0.7]
]
# Results in our Logits Vector (Length 5)
[ -1.201, 0.891, 0.131, 0.307, 1.078 ]
The Prediction
The maximum value in our Logits Vector is 1.078 at index 4.
If we look up index 4 in our vocabulary, we find the word "Delhi".
Our completed sentence is now: "Tomorrow I am flying to Delhi"!
💻 Code Implementation
Here is the exact PyTorch implementation for the concepts discussed above:
from d_transformer import TransformerBlock
from a_layerNorm import LayerNorm
import torch.nn as nn
import torch
class GPTModel(nn.Module):
def __init__(self, config):
super().__init__()
self.tok_emb = nn.Embedding(config["vocab_size"], config["emb_dim"]) #nn.Embedding(50257,768)
self.pos_emb = nn.Embedding(config["context_length"], config["emb_dim"]) #nn.Embedding(1024,768)
self.drop_emb = nn.Dropout(config["drop_rate"]) #nn.Dropout(0.1)
self.trf_blocks = nn.Sequential(
*[TransformerBlock(config) for _ in range(config["n_layers"])])
self.final_norm = LayerNorm(config["emb_dim"])
self.out_head = nn.Linear(
config["emb_dim"], config["vocab_size"], bias=False
)
def forward(self, in_idx):
batch_size, seq_len = in_idx.shape
tok_embeds = self.tok_emb(in_idx)
pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
x = tok_embeds + pos_embeds # Shape [batch_size, num_tokens, emb_size]
x = self.drop_emb(x)
x = self.trf_blocks(x)
x = self.final_norm(x)
logits = self.out_head(x)
return logits