Now, let's build your very own AI.

The whole structure — from token IDs to next-token probabilities

XSLM’s GPT does the calculation below once for each input token. The structure is a decoder-only Transformer like GPT-2, just with 1 block and 1 head.

• Take a vector (8 slots) from the embedding table by token ID, take a position vector by position number, and add them.
• Normalize, mix in information from earlier positions with causal attention, and add it to the original vector (residual).
• Normalize again, process with the MLP, and add it to the original vector again (residual).
• Finally, normalize and use the output head (head) to give a score (logit) for each of the 40 vocabulary tokens.
• Turning the scores into probabilities with softmax gives the next-token probabilities.

A real GPT stacks the middle block (attention + MLP) many times. GPT-2 Small stacks 12, GPT-3 stacks 96. Each block takes the vectors made by the previous block and refines them a little more. A block’s input and output have the same shape, so you can stack as many as you like.

Embedding — an ID picks one row of a table

An embedding is a table of vocabulary size × width. In XSLM it is 40 rows × 8 slots. If the token ID is 5, it just takes the 8 numbers in row five (counting from 0). It is a table lookup with no multiplication and no search.

Why turn IDs into vectors instead of using them directly? An ID has only size and order, so it cannot hold relations like "9 and 23 are similar". A vector of 8 numbers can be near or far from others in many directions, so training can move the vectors of tokens used in similar ways closer together.

The embedding table is also a model weight. It starts as random numbers, and training adjusts it in the direction that lowers the loss. No separately built embedding model is imported. In the experiment below, look at the randomly initialized vector of token 5 and the vector of position 0, and count the sizes of the 9 matrices.

from model import GPT

m = GPT(40, 8, 12)
print('token 5   :', [round(v.data, 3) for v in m.weights['token'][5]])
print('position 0:', [round(v.data, 3) for v in m.weights['position'][0]])
for name, matrix in m.weights.items():
    rows, cols = len(matrix), len(matrix[0])
    print(name, rows, 'x', cols, '=', rows * cols)
print('total', len(m.parameters))

cd ~/xslm
mkdir -p lab
micro lab/embedding.py
python3 -m lab.embedding

token 5   : [-0.081, 0.024, 0.043, -0.149, 0.005, 0.131, -0.201, -0.032]
position 0: [-0.102, -0.036, 0.089, 0.01, -0.085, -0.084, 0.267, 0.114]
token 40 x 8 = 320
position 12 x 8 = 96
query 8 x 8 = 64
key 8 x 8 = 64
value 8 x 8 = 64
project 8 x 8 = 64
expand 16 x 8 = 128
contract 8 x 16 = 128
head 40 x 8 = 320
total 1248

Embedding space — close means similar

An embedding writes a thing or a word as a few numbers. Each number is one axis. The map below writes things with 2 numbers: speed (across) and size (up and down). 0 is the slowest and smallest, and 1 is the fastest and biggest. A cow is [0.1, 0.8], so it is slow and big. An eagle is [0.9, 0.4], so it is fast and small.

Once things are numbers, you can measure "how similar" as a distance. The Euclidean distance between two points [x₁, y₁] and [x₂, y₂] is √((x₁ − x₂)² + (y₁ − y₂)²). The shorter the distance, the more similar they are. Measured from the alarm clock [0, 0], it looks like this.

• Bicycle [0.3, 0.3]: √(0.09 + 0.09) = 0.424
• Cow [0.1, 0.8]: √(0.01 + 0.64) = 0.806
• Eagle [0.9, 0.4]: √(0.81 + 0.16) = 0.985

So the thing most like the alarm clock is the bicycle. On the map, click a point to make it the starting point, or click an empty spot to start from there. The closest-first order is worked out again right away.

AI embeddings use the same idea, with two differences. There are not 2 axes but hundreds to thousands (XSLM has 8), and training decides what each axis means, not people. So the axes have no names like "speed", but the vectors of tokens that are used in similar ways gather close together.

Measuring distance is used all over AI services. A vector database finds the documents closest to the embedding of a question, and RAG hands those documents to the language model along with the question. Databases like MariaDB do the same math with a VECTOR column and the VEC_DISTANCE_EUCLIDEAN function. The experiment below works out the same map in Python.

import math

# A 2-number embedding: [speed, size]. 0 is the slowest and smallest, 1 the fastest and biggest.
things = {
    'alarm clock': [0.0, 0.0],
    'cow': [0.1, 0.8],
    'bicycle': [0.3, 0.3],
    'eagle': [0.9, 0.4],
    'elephant': [0.1, 0.9],
    'me': [0.2, 0.65],
    'horse': [0.7, 0.7],
    'car': [0.9, 0.85],
    'ship': [0.3, 0.9],
    'rocket': [0.99, 0.95],
}


def distance(a, b):
    return math.sqrt((a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2)


for query in [[0.0, 0.0], [0.1, 0.8], [0.5, 0.5]]:
    print('from', query)
    nearest = sorted(things, key=lambda name: distance(things[name], query))
    for name in nearest[:4]:
        print(f'  {name:12} {distance(things[name], query):.3f}')

cd ~/xslm
mkdir -p lab
micro lab/distance.py
python3 -m lab.distance

from [0.0, 0.0]
  alarm clock  0.000
  bicycle      0.424
  me           0.680
  cow          0.806
from [0.1, 0.8]
  cow          0.000
  elephant     0.100
  me           0.180
  ship         0.224
from [0.5, 0.5]
  horse        0.283
  bicycle      0.283
  me           0.335
  eagle        0.412

Position embedding — tells which position a token is in

Attention ignores order when mixing earlier positions. Without position information, the token vectors of "dog can run" and "run can dog" form the same set. So each position gets one learned position vector that is added to the token vector.

XSLM’s position table is context (12) rows × 8 slots. With only 12 rows, at most 12 tokens can go in at once. The assert on the first line of forward checks this, and that is why generate passes only the last 12 tokens, as in ids[-12:]. This limit is called the context length (context window). GPT-2 uses 1,024 tokens; recent models use hundreds of thousands.

GPT-2 and XSLM learn a vector for each position. The original Transformer paper computed position vectors with sin and cos, and recent models often use RoPE, which rotates Query and Key according to position. The methods differ, but the goal is the same: to let the model know the order.

Linear layers, normalization and ReLU — the three calculations a model is made of

Three functions in engine.py do most of the model’s math.

• linear(x, W): takes the dot product of x with each row of matrix W (multiply matching slots and add them all). If W has 3 rows, the result has 3 slots. In the experiment, the dot product of x = [1, 2] and the row [0.5, 0.5] is 0.5 + 1.0 = 1.5. Changing the number of rows lets you grow or shrink the number of slots in a vector.
• norm(x): subtracts the mean of the slots and divides by the standard deviation, giving mean 0 and variance 1. [2, 4, 6] becomes [-1.225, 0, 1.225]. It keeps numbers from getting too big or too small as calculations pile up. It adds 1e-5 to the variance to avoid dividing by 0.
• relu(v): turns negatives into 0 and leaves positives as they are.

A nonlinear function like ReLU is essential. Two linear layers joined together are still just one linear layer, so no matter how many you stack, they can only express straight-line relations. A bending function in between lets the model learn complex relations.

XSLM’s linear has no added value (bias), and its norm has no learned scale and shift (γ, β). PyTorch’s nn.Linear and nn.LayerNorm have both.

from engine import linear, norm

x = [1.0, 2.0]
W = [[1.0, 0.0], [0.5, 0.5], [-1.0, 1.0]]
print('linear:', linear(x, W))
print('norm  :', [round(v, 3) for v in norm([2.0, 4.0, 6.0])])
print('relu  :', [max(0.0, v) for v in [-2.0, 0.5, 3.0]])

cd ~/xslm
mkdir -p lab
micro lab/linear.py
python3 -m lab.linear

linear: [1.0, 1.5, 1.0]
norm  : [-1.225, 0.0, 1.225]
relu  : [0.0, 0.5, 3.0]

Attention — mixing earlier positions with Query, Key and Value

Attention decides which earlier positions the current position should consult and how much, and mixes their information. Each position makes three vectors.

• Query (q): what the current position is looking for.
• Key (k): the feature each position shows, saying "I have this kind of information".
• Value (v): the content each position actually passes on.

The calculation has four steps.

• Score: the dot product of the current q with each position’s k, divided by √width (√8 ≈ 2.83 in XSLM).
• Weight: softmax turns the scores into weights that add up to 1.
• Mix: multiply each position’s v by its weight and add them up.
• Tidy: transform the mixed result once more with the project matrix and add it to the original vector.

As a formula, it is softmax(q·kᵀ / √d)·v. The experiment below does this calculation with 3 handmade keys and one query (width 2, so it divides by √2). The query [2, 0] looks for the first feature, so positions 1 and 3, which have it, get 0.446 each, and position 2, which lacks it, gets 0.108. The result is [6.687, 3.313], leaning toward the values of positions 1 and 3.

For example, in "the cat sat because it was tired", attention at the position of it can learn to give a large weight to the position of cat. Real models use several heads (12 in GPT-2 Small), and each head looks at a different relation (what a word refers to, the word just before, format markers and so on).

from engine import Value, softmax

keys = [[1, 0], [0, 1], [1, 1]]
values = [[10, 0], [0, 10], [5, 5]]
query = [2, 0]
scores = [sum(q * k for q, k in zip(query, key)) / 2 ** 0.5 for key in keys]
weights = softmax([Value(s) for s in scores])
mixed = [sum(w.data * v[i] for w, v in zip(weights, values)) for i in range(2)]
print('scores ', [round(s, 3) for s in scores])
print('weights', [round(w.data, 3) for w in weights])
print('mixed  ', [round(m, 3) for m in mixed])

cd ~/xslm
mkdir -p lab
micro lab/attention.py
python3 -m lab.attention

scores  [1.414, 0.0, 1.414]
weights [0.446, 0.108, 0.446]
mixed   [6.687, 3.313]

Causal mask — no peeking at future tokens

During training, the whole window goes in at once and each position must predict its next token. If a position could see the tokens after it, it would be peeking at the answer and would learn nothing. So each position sees only itself and earlier positions. This is called causal attention.

The PyTorch implementation fills the upper part of the score table above the diagonal (the future) with −∞ so it becomes 0 after softmax (a lower-triangular mask). XSLM is simpler. It processes tokens one by one from the front and keeps only what has come so far in the keys and values lists, so future tokens are not in the lists yet.

So adding tokens at the end must not change the results for earlier positions. The experiment below feeds [1, 2] and [1, 2, 3] and checks that the scores for the first two positions are identical. True is evidence that it does not look at the future.

These keys and values lists are the same idea as the KV cache in real models. If you store the keys and values of earlier tokens while generating, you only need to compute the one new token. XSLM’s generate keeps it simple and recomputes the whole context for every token.

from model import GPT

m = GPT(10, 8, 12)
short = m.forward([1, 2])
long = m.forward([1, 2, 3])
same = [[v.data for v in r] for r in short] == [[v.data for v in r] for r in long[:2]]
print('first two positions unchanged:', same)
print('scores at position 3:', [round(v.data, 3) for v in long[2][:4]], '...')

cd ~/xslm
mkdir -p lab
micro lab/causal.py
python3 -m lab.causal

first two positions unchanged: True
scores at position 3: [0.159, 0.011, -0.239, 0.428] ...

softmax — scores to probabilities, and numerical stability

softmax turns a list of scores into probabilities between 0 and 1 that add up to 1. It applies exp to each score and divides by the total. exp is always positive and makes big scores much bigger, so it keeps the order while widening the gaps. In the experiment, the scores [2, 1, 0.1] become the probabilities [0.659, 0.242, 0.099].

The problem is that exp grows very fast. exp(1000) goes beyond the range a computer’s floating-point number can hold and raises OverflowError (the math range error in the experiment). So the softmax in engine.py first subtracts the largest score (peak) from every score. Subtracting the same number from every score does not change the probabilities, and the largest value becomes exp(0) = 1, so nothing overflows. [1000, 999] also comes out correctly as [0.731, 0.269].

softmax is used twice: to get the attention weights and to turn output head scores into next-token probabilities.

import math
from engine import Value, softmax

probs = softmax([Value(s) for s in [2.0, 1.0, 0.1]])
print([round(p.data, 3) for p in probs], round(sum(p.data for p in probs), 3))
try:
    print([math.exp(s) for s in [1000.0, 999.0]])
except OverflowError as error:
    print('without max:', error)
print('with max:', [round(p.data, 3) for p in softmax([Value(1000.0), Value(999.0)])])

cd ~/xslm
mkdir -p lab
micro lab/softmax.py
python3 -m lab.softmax

[0.659, 0.242, 0.099] 1.0
without max: math range error
with max: [0.731, 0.269]

MLP — processing information at each position

If attention gathers information by deciding "where to look", the MLP (multilayer perceptron, also called FFN) handles "how to transform what was gathered". It computes each position separately, with the same weights.

• expand: widens 8 slots to 16.
• ReLU: bends negatives to 0.
• contract: narrows 16 slots back to 8 and adds the result to the original vector.

The reason to widen and then narrow is to give room, in a wider space, to switch combinations of features on and off separately. Most of a big GPT’s parameters are in this MLP, and knowledge such as learned facts is thought to build up mainly here.

XSLM widens by 2 times and uses ReLU. GPT-2 widens by 4 times and use GELU. GELU is a function that bends smoothly near 0 and trains a little better in big models. XSLM chose ReLU because its derivative is easy to read by hand.

Residual connections and pre-normalization — trainable even when stacked deep

The line x = [a + b for a, b in zip(x, …)] in model.py is a residual connection. Instead of replacing x with the layer’s result, it adds the result to the original x.

• The original information stays, so the layer only needs to learn the difference to fix.
• In backward, addition passes the gradient through unchanged, so gradients flow by a shortcut all the way to earlier layers. Even with dozens of blocks stacked, gradients do not vanish.

Where normalization happens also matters. Like GPT-2, XSLM applies norm before attention, before the MLP and before the output head (Pre-LN). It normalizes only the values going into a layer and leaves the residual path untouched, so even deep models train stably. The original Transformer normalized after the addition (Post-LN).

As formulas, one block is this.

• x = x + Attention(norm(x))
• x = x + MLP(norm(x))

Output head — a score (logit) for each vocabulary entry

After the block, the vector is normalized one last time and multiplied by the head matrix (40 rows × 8 slots), giving one score for each of the 40 vocabulary tokens. These raw scores are called logits. Logits can be negative and do not add up to 1; they become probabilities only after softmax.

forward returns one score list for each input position. Feed 2 tokens and you get 2 × 40 scores. Training uses the scores of every position (each position tries to predict its next token), while generation uses only the scores of the last position, [-1].

GPT-2 does not keep a separate output head matrix; it shares the token embedding table (weight tying). This reduces parameters and uses the same vectors for reading and producing tokens. XSLM keeps them separate to show their different roles.

Automatic differentiation — following the computation graph backward

To train, you need to know how much the loss changes when each weight changes a little (the gradient). Wiggling all 1,248 one at a time would be far too slow, so the chain rule of derivatives finds them all at once. This is called backpropagation, and the tool that does it automatically is automatic differentiation (autograd). Value in engine.py is exactly that tool.

• Every time Value computes, it records the result along with (input, local derivative with respect to that input) in parents. Addition gives 1, multiplication the other value, exp the result itself, log 1/x, and ReLU 1 if positive, otherwise 0.
• backward starts from the loss with gradient 1.0 and walks the graph in reverse, adding "my gradient × local derivative" to each input’s gradient (the chain rule).
• If one value is used in several places, the gradients from every path are added up. That is why grad accumulates with +=, not =.

In the experiment below, y = x² + 2x differentiates by hand to 2x + 2, which is 8 at x = 3. The 8.0 from autograd matches the numerical derivative of 8.0 measured by moving x a tiny bit. When a is used twice, as in c = a·b + a, dc/da = b + 1 = 6. The gradients from the two paths were added.

PyTorch’s loss.backward() does the same thing. It just records and computes in tensors instead of single numbers.

from engine import Value

x = Value(3)
y = x * x + 2 * x
y.backward()
f = lambda v: v * v + 2 * v
h = 1e-6
print('autograd:', x.grad)
print('numeric :', round((f(3 + h) - f(3 - h)) / (2 * h), 6))
a, b = Value(2), Value(5)
c = a * b + a
c.backward()
print('dc/da =', a.grad, ' dc/db =', b.grad)

cd ~/xslm
mkdir -p lab
micro lab/grad.py
python3 -m lab.grad

autograd: 8.0
numeric : 8.0
dc/da = 6.0  dc/db = 2.0

Counting parameters — where are the 1,248?

Parameters are the individual numbers that training adjusts. They are also called weights. XSLM holds all its parameters in 9 matrices, and adding up the slots of every matrix gives 1,248. Compare the table below with the output of the earlier embedding experiment.

You can read a few things from the table.

• The two vocabulary-bound tables, token and head, hold 640 numbers, more than half. Growing the vocabulary grows these two tables.
• The 4 attention matrices (query, key, value, project) hold 256, and the 2 MLP matrices hold 256. Doubling the width makes this part 4 times bigger.
• In real models, blocks (attention + MLP) are stacked dozens of times and take up most of the parameters.

The B in "8B model" means billion: 8 billion parameters. Storage size is the number of parameters × the size of one number. With 32-bit floats (4 bytes) that is about 32GB, with 16-bit 16GB, reduced to 8-bit (quantization) 8GB, and 4-bit about 4GB. Running it needs extra memory on top for intermediate values and the KV cache. XSLM’s model.json is JSON with numbers written as text, so even with 1,248 numbers it is about 35KB.

What XSLM cut — differences from a real GPT

XSLM removed or reduced a few things to stay readable. Here is how it compares with GPT-2 Small.

• Number of blocks: XSLM 1, GPT-2 Small 12. The more blocks you stack, the more complex the relations it can learn.
• Attention heads: XSLM 1, GPT-2 Small 12 (each head 768 ÷ 12 = 64 slots wide). Multiple heads look at different relations at the same time.
• MLP: XSLM 2 times with ReLU, GPT-2 4 times with GELU.
• bias and the normalization γ and β: none in XSLM, present in GPT-2.
• dropout: none in XSLM. It is an overfitting guard that randomly sets some of the values (usually 10%) to 0 during training so the model does not rely on any particular path. It is switched off when generating.
• Initial values: a normal distribution with standard deviation 0.1 in XSLM and 0.02 in GPT-2.
• How it computes: XSLM one number at a time; PyTorch computes tensors shaped (batch B, length T, width C) all at once on the GPU.

Everything removed is machinery for training bigger and more stably. The core calculation that produces next-token probabilities is the same.

Next-token prediction is not vector search

If you know embeddings only as vectors for search, it is easy to misunderstand that "GPT picks the next word by finding the closest vector". In reality it works like this.

• Vector search (RAG and the like): documents are turned into vectors and stored, and the documents closest to the question vector are found and taken out. The result is a document that already existed.
• GPT’s prediction: the blocks transform the input vectors by calculation, then the output head scores the whole vocabulary, and one token is sampled from those probabilities. The result is a newly calculated distribution.

Attention is not search either. It only computes weights among the earlier positions inside the input sentence and mixes them; it does not look up outside documents.

The two are used together. Real services use the question to find related documents by vector search, attach those documents to the prompt, and have GPT generate the answer (RAG). Search handles up-to-date information; GPT handles writing sentences.