AI ENGINEERING: DEEP DIVE · MODULE 01 — THE MATH UNDERNEATH · 03 A hands-on series · math + code

Module 1.3

Every layer is the same trick

One matrix, one vector, one nonlinearity. Step through each output neuron of a tiny layer and watch that it's just four independent dot products — which is also exactly why this is the operation a GPU is built for.

y = ReLU(Wx + b) — STEP THROUGH EACH OUTPUT
× = ReLU
row 1 · x + b = (0.5)(1.0) + (-0.3)(2.0) + (0.8)(-1.0) + 0.1 = -0.80 ReLU(-0.80) = 0.00

The math

Two moves, repeated at every layer

Matrix-vector product
yᵢ  =  Σⱼ Wᵢⱼxⱼ + bᵢ
= (0.5)(1.0) + (-0.3)(2.0) + (0.8)(-1.0) + 0.1
= -0.80
ReLU nonlinearity
ReLU(z)  =  max(0, z)
= max(0, -0.80)
= 0.00
Scaled attention score
score  =  (q · k) / √d
= 1.02 / √4
= 0.51

Without ReLU (or something like it), stacking matmuls would collapse into one big matmul — the nonlinearity is what makes depth actually mean something.

The code

The whole layer, in four lines

one_layer.py
import numpy as np

W = np.array([[0.5, -0.3, 0.8],
              [-0.6, 0.4, 0.2],
              [0.1, 0.9, -0.5],
              [-0.7, -0.2, 0.6]])
x = np.array([1.0, 2.0, -1.0])
b = np.array([0.1, -0.2, 0.3, 0.0])

y_pre = W @ x + b               # [-0.8, -0.2, 2.7, -1.7]
y     = np.maximum(0, y_pre)     # [ 0.0,  0.0, 2.7,  0.0]  — ReLU

Four rows of W means four completely independent multiply-and-sum jobs. Nothing in row 3's answer depends on row 1's.

Why GPUs matter

A GPU doesn't just have "more cores" in some vague sense — it has thousands of small cores built to do exactly this: many independent multiply-adds, at the same time. Every row above is independent of every other row, so a GPU computes all 4 (or 4 billion, in a real model) in parallel instead of one after another. That parallelism, not core count in the abstract, is the entire reason matrix multiplication had to get fast.

The Nimbus thread

One real attention step, on the query's own words

Tokenize the query into words and treat each as a tiny embedding. Ask how much "export" should attend to each word, including itself.

attention_step.py
import numpy as np

q = np.array([0.9, 0.1, -0.2, 0.4])            # embedding for "export"
keys = np.array([
    [0.9,  0.1, -0.2,  0.4],   # "export" (itself)
    [0.8,  0.0, -0.3,  0.5],   # "fail"
    [-0.5, 0.6,  0.7, -0.1],   # "why"
])

scores  = (keys @ q) / np.sqrt(len(q))                 # [0.51, 0.49, -0.29]
weights = np.exp(scores) / np.exp(scores).sum()   # softmax -> [0.41, 0.40, 0.19]
export
0.41
fail
0.40
why
0.19

"Export" ends up attending almost equally to itself and to "fail," far more than to "why." That pull is exactly how a model's representation of "export" gets blended with the meaning of "fail" — the mechanism Module 2 builds out in full.

So what

A transformer layer isn't mysterious — it's matmul, a nonlinearity, and one more matmul, repeated at every layer, on every token, in parallel.

Go deeper (optional)

The full attention mechanism, visually

The toy example above worked one attention step by hand. 3Blue1Brown's video builds the complete mechanism — multiple heads, why it's Q/K/V and not just one matrix, how it scales to a real sentence — with the same visual style as this page's diagram. It's long (26 min) for a "go deeper" clip, but it's the single best explanation of this exact mechanism available anywhere.

3Blue1Brown — "Attention in transformers, step-by-step," Deep Learning ch. 6. Supplementary, not required to follow this page — Module 2 builds the same mechanism from scratch.