Linear algebra for AI

1 · Vectors, spaces, span & basis intermediate

Here is the shift in mindset that makes all of modern AI click. When your RAG system says two documents are “similar,” or a model “knows” that king and queen are related, no one wrote a rule for that. Instead, every word, sentence, and image has quietly been turned into a list of numbers — and once meaning lives as numbers, similarity becomes geometry. This section builds the space those numbers live in, and everything else in the lesson is just operations on the points inside it.

Picture a huge room. A vector is a single point in that room — or, equivalently, an arrow from the corner (the origin) out to that point. In daily life you know 2-D (a map) and 3-D (the room you’re in); AI just uses rooms with 384, 768, or 1536 axes instead of 3. You can’t picture 1536 dimensions, and that’s fine — the arithmetic is identical to the 2-D case, so reason in 2-D and trust it scales. The intuition to hold onto: an embedding is a vector, and its direction (which way the arrow points), not where it happens to land, is what carries the meaning.

Two ideas do the heavy lifting. The span of a set of vectors is every point you can reach by scaling them (stretching the arrows) and adding them (laying them tip-to-tail). A basis is a minimal set of arrows whose span fills the whole room — the fewest directions you need to name any point, the way “east” and “north” let you name any spot on a map. The count of arrows in a basis is the dimension. The common mistake is assuming more vectors always means more reach: if one arrow is just a stretched copy of another (they point the same way), it adds nothing to the span — it is linearly dependent, pure redundancy. Finding and squeezing out that redundancy is exactly what PCA (§6) does to shrink bloated embeddings.

Now the precise version, which lands because you have the picture. A vector is an ordered list of numbers — a point (or arrow from the origin) in Rⁿ, the space of all n-tuples of reals. In AI an embedding is a vector: a 384- or 1536-dimensional point whose direction encodes meaning. The span of a set of vectors is every point you can reach by scaling and adding them; a basis is a minimal set whose span is the whole space (n independent vectors span Rⁿ). Dimension is the size of a basis.

python# stdlib only
def add(u, v):        return [a + b for a, b in zip(u, v)]
def scale(c, v):      return [c * a for a in v]
def lincomb(coeffs, vecs):
    out = [0.0] * len(vecs[0])
    for c, v in zip(coeffs, vecs):
        out = add(out, scale(c, v))
    return out

e1, e2 = [1, 0], [0, 1]          # the standard basis of R^2
# Any point in R^2 is a linear combination of e1, e2:
print(lincomb([3, 4], [e1, e2]))     # [3.0, 4.0]

# Two independent vectors also span R^2; a dependent pair does not:
a, b = [1, 1], [2, 2]            # b = 2*a  -> dependent, span is a line
print(lincomb([1, -0.5], [a, b]))    # [0.0, 0.0]  (they cancel)

[3.0, 4.0]
[0.0, 0.0]

The first result rebuilds the point (3,4) from the basis. The second shows a dependent pair: because b = 2a, some non-trivial combination gives the zero vector — the hallmark of linear dependence, and exactly the redundancy PCA (§6) strips out.

2 · Dot product, norms & cosine — the embedding metric advanced

This one operation is the beating heart of every retrieval system, every attention head, every “how similar are these?” you will ever ask a model. When a vector database returns the “most relevant” chunk, when a recommender finds “users like you,” when a transformer decides which earlier token to attend to — underneath, it is the same single number being computed millions of times: the dot product. Get an honest feel for what it measures and half the mystery of modern AI dissolves.

Here is the intuition. Take two arrows from §1. The dot product answers one question: how much do they point the same way? Aligned arrows (small angle) give a big positive number; perpendicular arrows give exactly zero (“unrelated”); arrows pointing opposite give a negative number. That’s it — it is an alignment meter. But it’s a slightly dishonest meter, because it also secretly multiplies in the two arrows’ lengths: a long arrow and a short arrow that point the same way score lower than two long ones. This is the crux, captured in the identity u·v = ‖u‖‖v‖cosθ — the dot product bundles angle and length together.

For text similarity, length is usually noise you want gone. A short tweet and a long essay about the same topic should count as similar; you care about topic (direction), not verbosity (magnitude). The fix is cosine similarity: divide the dot product by both lengths, and the magnitudes cancel, leaving pure cosθ — a number in [−1, 1] that reads as “1 = identical direction, 0 = unrelated, −1 = opposite.” That is exactly why RAG (Ch 3) ranks by cosine and not raw dot product. The common mistake is forgetting the division: comparing raw dot products lets a merely long vector masquerade as a relevant one, so a verbose document out-scores a bang-on short one for the wrong reason. Normalize first, and the comparison is honest.

The precise definitions. The dot product u·v = Σ u_iv_i measures alignment. It equals ‖u‖ ‖v‖ cosθ, so it bundles two facts: the lengths and the angle. The L2 norm ‖v‖ = √(v·v) is a vector's length; the L1 norm is Σ|v_i|. Cosine similarity = (u·v)/(‖u‖‖v‖) throws away length and keeps only direction — which is why RAG (Ch 3) ranks by cosine: two texts are similar when their embedding vectors point the same way, regardless of magnitude.

python# stdlib only
import math

def dot(u, v):   return sum(a * b for a, b in zip(u, v))
def l2(v):       return math.sqrt(dot(v, v))
def l1(v):       return sum(abs(a) for a in v)
def cosine(u, v):
    return dot(u, v) / (l2(u) * l2(v))

u = [1.0, 2.0, 2.0]
v = [2.0, 0.0, 0.0]
print("dot   ", dot(u, v))                    # 2.0
print("l2(u) ", l2(u))                         # 3.0  (sqrt(1+4+4))
print("l1(u) ", l1(u))                         # 5.0
print("cos   ", round(cosine(u, v), 4))        # 0.3333

# Cosine ignores magnitude: scaling v by 10 does not change the angle.
print("cos*10", round(cosine(u, [20.0, 0.0, 0.0]), 4))   # 0.3333

dot    2.0
l2(u)  3.0
l1(u)  5.0
cos    0.3333
cos*10 0.3333

‖u‖ = √(1+4+4) = 3 exactly. The cosine is 2/(3·2) = 1/3 ≈ 0.3333, and scaling v by 10 leaves it unchanged — the property that lets a normalized dot product serve as the similarity metric in every vector store (A7).

3 · Matrices as linear maps & matrix multiplication advanced

If the dot product is the atom, the matrix is the molecule — and matrix multiplication is where nearly all of an LLM’s compute and cost actually go. When you read that GPT-scale training burned millions of dollars of GPU time, or that a model is “compute-bound,” the operation eating those dollars is this one. GPUs were, in effect, redesigned to do it fast. So it is worth slowing down and seeing a matrix for what it really is.

Don’t think of a matrix as a spreadsheet of numbers — think of it as a machine that transforms vectors. Feed it an arrow, it hands back a (usually different) arrow: it can rotate, stretch, squash, or project space. That is what “linear map” means — a transformation that keeps grid lines straight and evenly spaced (it respects addition and scaling). Every weight matrix in a neural network is one of these machines, and learning is just tuning the numbers so the machine transforms inputs into useful outputs. Now, multiplying two matrices means chaining two machines: AB is “first do B, then do A.” The arithmetic rule — entry (i,j) is row i of A dotted with column j of B — is just §2’s dot product run for every row-column pair, so a matmul is a whole grid of dot products at once.

Two facts trip everyone up, and both are worth burning in. First, the shapes must line up: an (m×k) times a (k×n) works only because the inner k’s match, giving an (m×n) result. The mismatched-dimension error is the single most common bug in ML code — the “inner dims must match” assertion in the lab is there precisely because you will hit it. Second, order matters: AB is generally not BA, because “rotate then stretch” is not “stretch then rotate.” That non-commutativity is why the order you stack layers changes the function the network computes. And the cost: the naive algorithm is O(m·n·k) — O(n³) for square matrices — which is why doubling a model’s width roughly octuples the matmul work, and why so much engineering (tiling, BLAS, sub-cubic tricks) exists to shrink the constant.

Precisely: a matrix A (m×n) is a linear map from Rⁿ to R^m: it takes an n-vector and returns an m-vector, preserving addition and scaling. Multiplying two matrices composes their maps. The rule: (AB)_ij = Σ_k A_ikB_kj — row i of A dotted with column j of B. The inner dimensions must match (m×k)(k×n) → m×n, and the naive cost is O(m·n·k) — O(n³) for square matrices.

python# stdlib only
def matmul(A, B):
    m, k = len(A), len(A[0])
    k2, n = len(B), len(B[0])
    assert k == k2, f"inner dims must match: {k} != {k2}"
    C = [[0.0] * n for _ in range(m)]
    for i in range(m):
        for j in range(n):
            s = 0.0
            for t in range(k):          # the O(k) dot product
                s += A[i][t] * B[t][j]
            C[i][j] = s
    return C

A = [[1, 2],
     [3, 4]]              # 2x2
B = [[5, 6],
     [7, 8]]              # 2x2
print(matmul(A, B))       # [[19.0, 22.0], [43.0, 50.0]]

# Composition is NOT commutative: AB != BA in general.
print(matmul(B, A))       # [[23.0, 34.0], [31.0, 46.0]]

# Multiplications for two n x n matrices = n^3 (here 2^3 = 8 per product).
n = len(A)
print("scalar mults:", n**3)   # 8

[[19.0, 22.0], [43.0, 50.0]]
[[23.0, 34.0], [31.0, 46.0]]
scalar mults: 8

(AB)₀₀ = 1·5 + 2·7 = 19; the rest follow. AB ≠ BA proves matrix multiplication is not commutative — order of composition matters, which is why stacking neural layers in a different order gives a different function. The O(n³) count is why matmul dominates both training and inference FLOPs (MS1).

4 · Matrix–vector product = a neural layer advanced

Here is the payoff that connects the math to the models you actually use. A neural network sounds like something exotic and biological — “neurons,” “activations,” “deep learning.” Strip the mystique and the fundamental building block, the thing repeated billions of times inside GPT or Claude, is embarrassingly simple: a matrix times a vector, plus a bias, then bend the result a little. That’s a layer. Once you see it, you can never un-see it — the whole tower is this one Lego brick, stacked.

Think of the input vector x as a description of something (a token, a pixel patch, a hidden state). Each row of the weight matrix W is a little feature detector: it dots itself with x (§2 again) to ask “how much does the input look like the pattern I care about?” and produces one number — one output neuron. Stack m such rows and you get m answers at once, which is exactly the matrix–vector product Wx from §3 with n=1 output columns. The bias b is a per-neuron nudge (“fire a bit more/less eagerly”), and the nonlinearity afterward — here ReLU, which simply clamps negatives to zero — is the one non-negotiable ingredient: without it, stacking layers would collapse into a single linear map (matrix times matrix is just another matrix), and the network could never learn anything a straight line couldn’t.

The common misconception is that the “intelligence” hides in some clever architecture. It mostly hides in the numbers inside W, learned from data — the operation itself never changes. That’s liberating: if you understand y = Wx + b, you understand the mechanical core of every embedding lookup, every attention projection, and every final logit computation in an LLM. The differences between models are width, depth, and how the blocks are wired — not a different atom. This is the exact operation A4 vectorizes for speed and MS1 fits into GPU memory.

Precisely: a fully-connected neural layer is literally y = Wx + b followed by a nonlinearity: W is the weight matrix (out×in), x the input vector, b the bias. The matrix–vector product Wx is just m dot products — each output neuron dots its weight row with the input. Every embedding lookup, projection, and logit computation in an LLM is this shape.

python# stdlib only
def matvec(W, x):
    return [sum(w_i * x_j for w_i, x_j in zip(row, x)) for row in W]

def relu(v):    return [max(0.0, a) for a in v]

# A layer mapping R^3 -> R^2:  y = relu(W x + b)
W = [[0.2, -0.5, 0.1],
     [0.7,  0.3, -0.2]]
b = [0.1, -0.4]
x = [1.0, 2.0, 3.0]

Wx = matvec(W, x)                        # [-0.5, 0.7]
pre = [a + c for a, c in zip(Wx, b)]     # add bias -> [-0.4, 0.3]
y   = relu(pre)                          # [0.0, 0.3]
print("Wx ", [round(v, 3) for v in Wx])
print("pre", [round(v, 3) for v in pre])
print("y  ", [round(v, 3) for v in y])

Wx  [-0.5, 0.7]
pre [-0.4, 0.3]
y   [0.0, 0.3]

Row 0: 0.2·1 − 0.5·2 + 0.1·3 = −0.5. Add bias → −0.4, and ReLU clamps it to 0. Row 1 → 0.7, +bias → 0.3, survives ReLU. Stack thousands of these rows and hundreds of layers and you have a network — but the atom is this matrix–vector product. This is the exact operation A4 vectorizes and MS1 fits into GPU memory.

5 · Transpose, inverse, rank intermediate

Three words — transpose, inverse, rank — and the last one quietly explains why LoRA lets you fine-tune a giant model on a laptop. These feel like dry bookkeeping until you realize rank is a measure of how much a matrix actually does, and that measure is the lever behind the cheapest, most popular fine-tuning trick in the field. So it’s worth the intuition tour.

Transpose (A^T) is the easy one: flip the matrix over its diagonal so rows become columns. It looks trivial, but it’s the piece of plumbing that turns the keys matrix K into K^T so the query-key dot products line up in attention (§8) — a shape adapter, nothing more. The inverse (A⁻¹) is the “undo” button: if A transforms space one way, A⁻¹ transforms it back, so A⁻¹A = I. But undo isn’t always possible — and that’s where rank enters.

Think of rank as the number of genuinely independent directions a matrix’s output can reach. If a 3×3 matrix squashes all of 3-D space flat onto a 2-D plane (because one row is secretly a combination of the others), its rank is 2, not 3 — it has thrown away a dimension. Here is the killer consequence: you cannot un-squash. Once information is flattened, no inverse can recover the lost dimension, so a rank-deficient (“singular”) matrix has no inverse, and its determinant is 0. The common trap is calling inv() or a naive least-squares solve on such a matrix and getting a crash or garbage — the lab deliberately builds a singular matrix (row 2 = 2×row 1) so you feel the failure. The flip side is the opportunity: a low-rank matrix carries little independent information, so it can be stored as a skinny product of two small matrices. That is precisely the LoRA insight — a huge weight update is approximated by a rank-r product with a tiny fraction of the parameters — and the same low-rank compressibility is what makes the SVD in §6 and §7 able to shrink embeddings.

Precisely: the transpose A^T flips rows and columns; it is what turns K into K^T in attention. The inverse A⁻¹ undoes the map (AA⁻¹ = I); it exists only when A is square and full rank. Rank = the number of linearly independent rows (= columns) = the dimension of the output space the map actually reaches. A low-rank matrix squashes space onto a smaller subspace — the idea LoRA exploits (a big weight update approximated by a rank-r product) and the reason SVD can compress.

python# stdlib only
def T(A):
    return [[A[i][j] for i in range(len(A))] for j in range(len(A[0]))]

def inv2(A):
    (a, b), (c, d) = A
    det = a * d - b * c
    if det == 0:
        raise ValueError("singular: det=0, no inverse (rank < 2)")
    return [[ d / det, -b / det],
            [-c / det,  a / det]]

M = [[4.0, 7.0],
     [2.0, 6.0]]
print("Mt  ", T(M))                 # [[4.0, 2.0], [7.0, 6.0]]
inv = inv2(M)
print("inv ", [[round(x, 3) for x in r] for r in inv])

# A rank-deficient matrix: row2 = 2*row1  -> det 0 -> no inverse.
S = [[1.0, 2.0],
     [2.0, 4.0]]
try:
    inv2(S)
except ValueError as e:
    print("singular:", e)

Mt   [[4.0, 2.0], [7.0, 6.0]]
inv  [[0.6, -0.7], [-0.2, 0.4]]
singular: singular: det=0, no inverse (rank < 2)

det(M) = 4·6 − 7·2 = 10, so M is full rank and invertible (inv = [[0.6, −0.7], [−0.2, 0.4]]). S has row 2 = 2×row 1, so det = 0, rank 1, and no inverse exists — the map collapses R² onto a line. Detecting this (singular / ill-conditioned matrices) is exactly what breaks a naive least-squares solve.

6 · Eigenvalues, eigenvectors & power iteration expert advanced

Eigenvectors sound like the most abstract idea in the course, and they are secretly the most practical — they are how Google ranked the web. The word is intimidating (“eigen” is just German for “own” or “characteristic”), but the concept is a matrix confessing its own favorite directions. Once you have that, PageRank, spectral clustering, and the SVD that compresses your embeddings all turn out to be the same idea wearing different hats.

Recall from §3 that a matrix is a machine that transforms arrows — it usually rotates and stretches them, so an arrow comes out pointing somewhere new. But for most machines there are a few special directions that come out pointing the exact same way, only longer or shorter. Those are the eigenvectors: directions the matrix merely scales instead of turning. The scaling factor for each is its eigenvalue λ — how much that direction gets amplified (λ > 1) or shrunk (λ < 1). These are the map’s natural axes, its grain; the eigenvector with the biggest |λ| is the direction the matrix stretches hardest, the one everything drifts toward under repeated application.

That last sentence is the whole trick behind power iteration, and it’s beautifully lazy: take any starting arrow, hit it with the matrix, normalize its length, and repeat. Because the dominant eigen-direction is amplified most each round, the arrow gets pulled toward it and everything else washes out — after a few dozen multiplies you’ve found the top eigenvector without ever solving a scary characteristic polynomial. The common misconception is that you need heavy machinery (determinants, root-finding) to get eigenvectors; for the dominant one you need only repeated matmul-and-normalize, which is exactly why it scales to web-sized matrices. PageRank is the dominant eigenvector of the web’s link matrix; the largest singular value in §7’s SVD comes from the same iteration applied to A^TA.

Precisely: an eigenvector v of A is a direction the map only stretches: Av = λv, where the scalar λ is its eigenvalue. Eigenvectors are the map's natural axes; the largest-|λ| eigenvector is the direction of greatest amplification. Power iteration finds it by repeatedly multiplying a random vector by A and renormalizing — the dominant eigenvector wins. This is the mathematical heart of PageRank, spectral clustering, and (via A^TA) the SVD below.

python# numpy-required
import numpy as np

A = np.array([[2.0, 1.0],
              [1.0, 3.0]])          # symmetric -> real eigenvalues

def power_iteration(A, iters=100):
    v = np.array([1.0, 0.0])        # deterministic start (no RNG, no fake numbers)
    for _ in range(iters):
        Av = A @ v
        v = Av / np.linalg.norm(Av)   # renormalize each step
    lam = v @ (A @ v)               # Rayleigh quotient = eigenvalue
    return lam, v

lam, v = power_iteration(A)
print("lambda ~", round(float(lam), 6))
print("vector ~", np.round(v, 6))

# Cross-check against the exact solver (A is symmetric -> use eigvalsh):
w = np.linalg.eigvalsh(A)               # real, ascending
print("exact eigenvalues:", np.round(np.sort(w)[::-1], 6))

lambda ~ 3.618034
vector ~ [0.525731 0.850651]
exact eigenvalues: [3.618034 1.381966]

The dominant eigenvalue is (5+√5)/2 ≈ 3.618034, and power iteration converges to it and its eigenvector without ever solving a characteristic polynomial — just repeated matmul + normalize. That the iterate lands on the largest |λ| is why the method underlies ranking and the largest singular value in SVD.

7 · SVD & PCA — dimensionality reduction expert expert

The SVD is often called the most important factorization in applied math, and for an AI engineer the reason is money: it is how you store a million 1536-dimensional embeddings for a fraction of the cost with barely any quality loss. It sounds heavy — three matrices, orthonormal bases, singular values — but the story it tells about any matrix is a clean one, and PCA, the workhorse of dimensionality reduction, is just the SVD with a haircut.

Here is the picture. Every matrix, no matter how messy, can be broken into three simple stages: rotate, stretch, rotate — that is A = UΣV^T. V^T spins the input onto a clean set of perpendicular axes; Σ stretches each of those axes by a singular value (all non-negative, and helpfully sorted biggest-first); U spins the result into the output space. The singular values are the payoff: each one tells you how much this direction matters. A big singular value is a direction carrying lots of the data’s spread; a tiny one is a direction that’s almost flat — nearly noise. The Eckart–Young theorem makes it official: throw away the smallest singular values and you get the best possible lower-rank approximation of the matrix.

That throwing-away is PCA: keep only the top-k directions of greatest variance and project your data onto them. Concretely, a 1536-dim embedding whose real information lives in ~100 directions can be squeezed to ~100 numbers, and the “variance kept” ratio (sum of kept singular values squared over the total) tells you exactly how much fidelity you traded for the savings — a dial you can turn. The common mistake is skipping the centering step: PCA measures variance about the mean, so you must subtract the column mean first, or the dominant “direction” the SVD finds is just the offset of your data from the origin, not its real shape. Done right, this compresses embeddings, denoises them, reveals latent structure, and connects straight back to the low-rank adapters of §5 — cheaper, faster vector search (A7) for minimal quality loss.

Precisely: every matrix factors as A = UΣV^T (the SVD): orthonormal directions V^T in the input, singular values Σ (non-negative, sorted) that scale them, and orthonormal directions U in the output. Keeping the top-k singular values gives the best rank-k approximation (Eckart–Young) — that is PCA: project data onto the few directions of greatest variance. In AI this compresses embeddings, denoises them, and reveals latent structure; it is also the theory behind low-rank adapters.

python# numpy-required
import numpy as np

# 5 points that live (almost) on a plane in 3-D:
X = np.array([[2.0, 0.0, 0.0],
              [0.0, 1.0, 0.0],
              [2.0, 1.0, 0.0],
              [4.0, 0.0, 0.0],
              [0.0, 2.0, 0.0]])

Xc = X - X.mean(axis=0)              # center (PCA works on variance about the mean)
U, S, Vt = np.linalg.svd(Xc, full_matrices=False)
print("singular values:", np.round(S, 4))     # 3rd ~ 0 -> data is 2-D

k = 2
proj = Xc @ Vt[:k].T                 # coordinates in the top-2 principal axes
var_kept = (S[:k]**2).sum() / (S**2).sum()
print("reduced shape:", proj.shape)            # (5, 2)
print("variance kept:", round(float(var_kept), 6))   # 1.0

singular values: [3.617  0.9577 0.    ]
reduced shape: (5, 2)
variance kept: 1.0

The third singular value is exactly 0 — the data has no spread in the z-direction, so it truly lives in 2-D and PCA keeps 100% of the variance while dropping a dimension. On real 1536-dim embeddings the tail singular values are small (not zero), so keeping the top ~50–100 components preserves most meaning at a fraction of the storage — faster, cheaper vector search (A7) with minimal quality loss.

8 · Why attention is QK^T expert expert

This is the summit — the mechanism that powers every transformer, and the reason this whole lesson existed. “Attention” is written up with an intimidating formula, softmax(QK^T/√d)V, and it gets described in mystical terms: the model “decides what to focus on.” But you have now built every single piece of that formula by hand, and the promise of this section is that it will read like a sentence you already speak, not a hieroglyph.

Here is the plain-language version. Every token produces three vectors: a query (“what am I looking for?”), a key (“what do I offer?”), and a value (“what I’ll contribute if chosen”). To decide how much token i should listen to token j, you ask how well i’s query aligns with j’s key — and alignment, as you learned in §2, is a dot product. Doing that for every pair at once is exactly the matrix QK^T from §3: row i, column j is “how relevant is j to i.” Then softmax turns each row of raw scores into clean weights that sum to 1 (a probability distribution over “who to attend to”), and multiplying those weights by V takes a weighted average of the value vectors — another matmul. That’s the entire mechanism: score alignment, normalize into weights, average the values. Dot product plus weighted average — nothing you haven’t already coded.

The one piece that looks arbitrary, the √d division, is the section’s “common mistake” made into a design choice. A dot product of two d-dimensional vectors sums d little products, so its typical size grows with d — feed those inflated scores into softmax and it saturates into a near one-hot spike, killing the gradients the model needs to learn (and risking numerical overflow). Dividing by √d rescales the scores back to a sane range so softmax stays soft and trainable. Skip it and big models simply fail to train — a beautiful example of a whole lesson’s worth of intuition (dot products, norms, variance, matmul) converging on one √ in the denominator.

Precisely: scaled dot-product attention is softmax(QK^T/√d)·V. Read it as linear algebra: Q (queries) and K (keys) are matrices whose rows are vectors; QK^T is the matrix of all pairwise dot products — row i, column j is how much query i aligns with key j (§2). Dividing by √d keeps those scores from exploding as dimension grows (variance control), softmax turns each row into weights that sum to 1, and multiplying by V takes a weighted average of the value vectors (a matrix–matrix product, §3). Attention is nothing but dot-product similarity + weighted averaging — the two operations this whole lesson built.

python# numpy-required
import numpy as np

def attention(Q, K, V):
    d = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d)             # (i) all pairwise dot products, scaled
    scores = scores - scores.max(axis=-1, keepdims=True)   # (ii) stable softmax
    w = np.exp(scores); w /= w.sum(axis=-1, keepdims=True)
    return w, w @ V                           # (iii) weighted average of values

Q = np.array([[1.0, 0.0]])                    # one query
K = np.array([[1.0, 0.0],                     # key 0 aligns with the query
              [0.0, 1.0]])                    # key 1 is orthogonal
V = np.array([[10.0, 0.0],
              [ 0.0, 5.0]])
w, ctx = attention(Q, K, V)
print("weights:", np.round(w, 4))             # more weight on key 0
print("context:", np.round(ctx, 4))

weights: [[0.6698 0.3302]]
context: [[6.6976 1.6512]]

The query aligns with key 0 (dot 1/√2 ≈ 0.707) and is orthogonal to key 1 (dot 0), so softmax gives key 0 the larger weight (0.6698 vs 0.3302), and the output leans toward V’s first row. That is the entire mechanism: QK^T scores alignment, softmax normalizes, V is averaged. For the numerical-stability details of that softmax see A4; for the GPU kernels that make it fast see MS3.

Checkpoint expert

import math
def dot(u, v): return sum(a*b for a, b in zip(u, v))
def l2(v):     return math.sqrt(dot(v, v))
def cosine(u, v): return dot(u, v) / (l2(u) * l2(v))

q  = [1.0, 1.0, 0.0]
c1 = [2.0, 2.0, 0.0]     # same direction as q
c2 = [0.0, 0.0, 5.0]     # orthogonal
s1, s2 = cosine(q, c1), cosine(q, c2)
print(round(s1, 4), round(s2, 4))        # 1.0 0.0
print('winner: c1' if s1 > s2 else 'winner: c2')
print(round(cosine(q, q), 4))            # 1.0

1.0 0.0
winner: c1
1.0

def matmul(A, B):
    m, k = len(A), len(A[0])
    k2, n = len(B), len(B[0])
    assert k == k2, f'inner dims {k} != {k2}'
    return [[sum(A[i][t]*B[t][j] for t in range(k))
             for j in range(n)] for i in range(m)]

A = [[1, 2, 3], [4, 5, 6]]     # 2x3
B = [[1, 0], [0, 1], [1, 1]]   # 3x2
print(matmul(A, B))            # [[4.0, 5.0], [10.0, 11.0]]
try:
    matmul(A, A)               # 2x3 @ 2x3 -> mismatch
except AssertionError as e:
    print('error:', e)

[[4.0, 5.0], [10.0, 11.0]]
error: inner dims 3 != 2

def layer(W, b, x):
    pre = [sum(w*xi for w, xi in zip(row, x)) + bi
           for row, bi in zip(W, b)]
    return pre, [max(0.0, v) for v in pre]

W = [[0.2, -0.5, 0.1], [0.7, 0.3, -0.2]]
b = [0.1, -0.4]
x = [1.0, 2.0, 3.0]
pre, y = layer(W, b, x)
print('pre', [round(v, 3) for v in pre])   # [-0.4, 0.3]
print('y  ', [round(v, 3) for v in y])     # [0.0, 0.3]

pre [-0.4, 0.3]
y   [0.0, 0.3]

# numpy-required
import numpy as np
A = np.array([[2.0, 1.0], [1.0, 3.0]])

def power_iteration(A, iters=200):
    v = np.array([1.0, 0.0])
    for _ in range(iters):
        Av = A @ v
        v = Av / np.linalg.norm(Av)
    return float(v @ (A @ v)), v

lam, v = power_iteration(A)
print(round(lam, 6), np.round(v, 6))       # 3.618034 [0.525731 0.850651]
print(np.round(np.sort(np.linalg.eig(A)[0])[::-1], 6))   # [3.618034 1.381966]

3.618034 [0.525731 0.850651]
[3.618034 1.381966]

# numpy-required
import numpy as np
X = np.array([[2.0,0.0,0.0],[0.0,1.0,0.0],[2.0,1.0,0.0],
              [4.0,0.0,0.0],[0.0,2.0,0.0]])
Xc = X - X.mean(axis=0)
U, S, Vt = np.linalg.svd(Xc, full_matrices=False)
k = 2
proj = Xc @ Vt[:k].T
kept = (S[:k]**2).sum() / (S**2).sum()
print(np.round(S, 4))              # [3.617  0.9577 0.    ]
print(proj.shape, round(float(kept), 6))   # (5, 2) 1.0

[3.617  0.9577 0.    ]
(5, 2) 1.0

# numpy-required
import numpy as np

def top_k(query, docs, k=3):
    q = query / np.linalg.norm(query)
    D = docs / np.linalg.norm(docs, axis=1, keepdims=True)
    sims = D @ q                          # (N,) cosine scores
    idx = np.argpartition(-sims, k-1)[:k] # top-k unordered, O(N)
    return idx[np.argsort(-sims[idx])], sims

rng = np.random.default_rng(0)
docs = rng.random((8, 5))
idx, sims = top_k(docs[2], docs, k=3)     # doc 2 IS the query
print('top-3:', idx.tolist())             # 2 leads
print('score0:', round(float(sims[idx[0]]), 6))  # 1.0
assert idx[0] == 2

top-3: [2, 1, 0]
score0: 1.0