Tutorial 02: Linear Algebra for Transformers¶
Every operation in a transformer — attention, feed-forward networks, layer normalization — is built on a small set of linear algebra primitives. This tutorial covers the math you actually need, with MLX code for each operation.
1. Vectors, Matrices, and Tensors¶
A vector is a 1D array of numbers:
A matrix is a 2D grid of numbers:
A tensor is a generalization to any number of dimensions:
In MLX:
import mlx.core as mx
vec = mx.array([1.0, 2.0, 3.0]) # 1D tensor, shape (3,)
mat = mx.array([[0.1, 0.2], # 2D tensor, shape (2, 3)
[0.3, 0.4]])
tensor = mx.arange(24).reshape(2, 3, 4) # 3D tensor, shape (2, 3, 4)
2. Matrix Multiplication¶
The single most important operation in a transformer. A linear layer computes:
where x is the input, W is the weight matrix, and b is an optional bias.
The Math¶
If x has shape (batch, in_features) and W has shape (out_features, in_features):
Each output neuron j computes a weighted sum of all inputs, where the weights
are row j of W.
In MLX¶
# Linear layer: maps 3-dimensional input to 4-dimensional output
x = mx.array([[1.0, 2.0, 3.0]]) # shape: (1, 3)
W = mx.random.normal((4, 3)) # shape: (4, 3) — 4 output dims, 3 input dims
# The core operation: x @ W^T
# MLX stores weights as (out_features, in_features), so we transpose
y = mx.matmul(x, W.T) # shape: (1, 4)
print(y.shape) # (1, 4)
This is exactly what our engine's linear() function does:
# core/layers.py — the simplest and most important function in the engine
def linear(x: mx.array, w: mx.array, bias: mx.array | None = None) -> mx.array:
if bias is not None:
return mx.matmul(x, w.T) + bias
else:
return mx.matmul(x, w.T)
Why w.T? MLX (and most frameworks) store weight matrices in a transposed
layout: (out_features, in_features) instead of (in_features, out_features).
This is a convention for memory layout efficiency — rows are contiguous in memory.
3. Batched Matrix Multiply¶
In practice, we never process a single vector at a time. We process batches:
X shape: (batch_size, seq_len, in_features)
W shape: (out_features, in_features)
Y = X @ W^T → shape: (batch_size, seq_len, out_features)
# Batched linear: process 2 sequences of length 5, each with 8 features
batch_size, seq_len, in_features = 2, 5, 8
out_features = 16
X = mx.random.normal((batch_size, seq_len, in_features))
W = mx.random.normal((out_features, in_features))
Y = mx.matmul(X, W.T) # shape: (2, 5, 16)
MLX automatically broadcasts the matrix multiply across the batch and sequence dimensions. This is why transformers are so fast — the entire sequence is processed in parallel.
4. Dot Product and Cosine Similarity¶
The dot product measures how "aligned" two vectors are:
- If the vectors point in the same direction, the dot product is large and positive.
- If they are perpendicular, the dot product is zero.
- If they point in opposite directions, it is negative.
This is the foundation of attention. The query vector "asks a question" and each key vector "describes what it contains." The dot product between a query and a key measures how relevant that key is to the query.
q = mx.array([1.0, 0.0, 0.0]) # points along x-axis
k = mx.array([0.0, 1.0, 0.0]) # points along y-axis
k2 = mx.array([1.0, 1.0, 0.0]) # points at 45 degrees
# Perpendicular: dot product is 0
print(mx.sum(q * k)) # 0.0
# Aligned: dot product is positive
print(mx.sum(q * k2)) # 1.0
5. Softmax: Converting Scores to Probabilities¶
The dot product produces raw scores (logits). To turn these into a probability distribution, we use softmax:
Softmax has two critical properties: 1. All outputs are positive and sum to 1 (valid probability distribution). 2. Larger inputs get exponentially larger outputs (sharpens differences).
logits = mx.array([1.0, 2.0, 3.0])
probs = mx.softmax(logits, axis=-1)
print(probs) # [0.09, 0.24, 0.67]
print(mx.sum(probs)) # 1.0
Why axis=-1?¶
In attention, we compute softmax over the "which token to attend to" dimension, which is always the last axis of the score matrix.
Numerical Stability¶
The naive softmax can overflow (exp(1000) = inf). The stable version subtracts the max:
Subtracting the max doesn't change the result (the exp(x - c) terms cancel in
the numerator and denominator), but prevents overflow. MLX handles this
internally, but it is important to understand.
# Numerically unstable (hypothetical — MLX handles this for us)
# exp(1000) = inf → NaN
# What MLX actually does internally:
x = mx.array([1.0, 2.0, 1000.0])
stable = x - mx.max(x) # [-999, -998, 0]
probs = mx.softmax(x, axis=-1) # MLX subtracts max automatically
6. Reshape and Transpose¶
Transformers constantly reshape tensors to move between "flat" representations and "multi-head" representations.
Reshape¶
x = mx.arange(24).reshape(2, 3, 4) # shape: (2, 3, 4) — 24 elements total
y = x.reshape(2, 12) # shape: (2, 12) — same 24 elements, flattened
z = x.reshape(6, 4) # shape: (6, 4) — same 24 elements
The total number of elements must stay the same. Reshape does not copy data — it just changes how the tensor is "viewed."
Transpose¶
a = mx.arange(6).reshape(2, 3) # [[0, 1, 2], [3, 4, 5]]
b = a.T # [[0, 3], [1, 4], [2, 5]]
# Multi-axis transpose (generalization)
c = mx.arange(24).reshape(2, 3, 4)
d = c.transpose(0, 2, 1) # swap axes 1 and 2: (2, 4, 3)
Why These Matter for Attention¶
The attention mechanism works on individual heads. To go from "all heads packed in one tensor" to "separate heads":
# Input: (batch, seq_len, num_heads * head_dim)
x = mx.random.normal((1, 10, 256)) # 8 heads × 32 dim = 256
# Step 1: Reshape to separate heads
x = x.reshape(1, 10, 8, 32) # (batch, seq, heads, head_dim)
# Step 2: Transpose to put heads in the batch dimension
x = x.transpose(0, 2, 1, 3) # (batch, heads, seq, head_dim)
# Now each head can compute attention independently
7. Softmax in Matrix Form¶
When computing attention, we apply softmax to each row of a score matrix independently:
# Scores: (batch, heads, seq_len, context_len)
# We want softmax over the context dimension (axis=-1)
scores = mx.random.normal((1, 8, 10, 10))
# Each row of scores[b, h, i, :] is a set of relevance scores for token i
probs = mx.softmax(scores, axis=-1) # softmax along context dimension
print(mx.sum(probs, axis=-1)) # each row sums to 1.0
8. Scaled Dot Product: The Attention Formula¶
Putting it all together, scaled dot-product attention is:
where:
- Q (query): shape (batch, heads, seq_len, head_dim)
- K (key): shape (batch, heads, context_len, head_dim)
- V (value): shape (batch, heads, context_len, head_dim)
- 1/sqrt(d_k): scaling factor to prevent dot products from growing too large
The step-by-step:
import math
def scaled_dot_product_attention_simple(q, k, v):
"""The simplest possible attention — building block for everything else."""
head_dim = q.shape[-1]
# Step 1: Compute relevance scores
# (batch, heads, seq, head_dim) @ (batch, heads, head_dim, context)
# → (batch, heads, seq, context)
scores = mx.matmul(q, k.swapaxes(-2, -1))
# Step 2: Scale to prevent gradient vanishing
scale = head_dim ** -0.5 # 1/sqrt(d_k)
scores = scores * scale
# Step 3: Convert scores to probabilities (softmax over context dimension)
weights = mx.softmax(scores, axis=-1)
# Step 4: Weighted sum of values
# (batch, heads, seq, context) @ (batch, heads, context, head_dim)
# → (batch, heads, seq, head_dim)
output = mx.matmul(weights, v)
return output
We will dive deep into attention in Tutorial 06. For now, notice that it is nothing more than three matrix multiplications, a softmax, and a scale.
9. Summary of Key Operations¶
| Operation | Math | MLX | Shape Transform |
|---|---|---|---|
| Linear | y = xW^T |
mx.matmul(x, w.T) |
(B, L, D_in) → (B, L, D_out) |
| Dot product | dot(q, k) |
mx.sum(q * k) |
(D,) → scalar |
| Batch matmul | C = AB^T |
mx.matmul(A, B.swapaxes(-2,-1)) |
(B,L,D) × (B,D,S) → (B,L,S) |
| Softmax | exp(x)/sum(exp(x)) |
mx.softmax(x, axis=-1) |
shape unchanged, rows sum to 1 |
| Reshape | flatten/expand | x.reshape(...) |
total elements preserved |
| Transpose | permute axes | x.transpose(...) |
axis order changes |
What's Next¶
Now that you have the linear algebra toolkit, we will use it to build the first real component of our transformer: the embedding layer, which converts token IDs into the dense vectors that all subsequent operations work with.