autograd.js — annotated walkthrough

What this file is. backprop.html derives the backpropagation algorithm in math. This file shows how that math maps to the ~380 lines of autograd.js in PicoGPT-JS — line by line.

The math defines two kinds of gradient that matter:

Activation gradient δ_l = ∂Loss/∂h_l — how sensitive the loss is to each layer's output. Travels backward through the network (right ← left in the diagram). Stored in .grad on intermediate tensors (h₁, h₂, …).
Weight gradient ∂Loss/∂W_l — how much each weight should change. Computed as (activation from forward pass) × (activation gradient from backward pass). Stored in .grad on parameter tensors (Wqkv, Wout, …). Used directly by the optimizer to update weights.

The two types share the same .grad field — what it means depends on whether the tensor is an intermediate activation or a learnable parameter.

How the graph is built. There is no separate "tape" object. As the forward pass runs, each operation writes two fields into its output tensor: .inputs (which tensors fed into it) and .backward (a closure that can push gradients back into those inputs). This list of closures is the tape — it is implicit in the tensor graph, not a separate data structure. When backward(loss) is called, it traverses the graph, collects all closures in reverse order, and calls them one by one. Each closure performs one step of the chain rule.

1 · Tensor — what every value in the network carries

structure Every scalar, vector, or matrix in PicoGPT-JS lives in a Tensor. It holds two data arrays and two graph fields. The data arrays carry the forward computation; the graph fields build the backward pass automatically. .data — the forward value, a Float32Array. For a weight matrix W this is the current weight values; for an activation h_l this is the output of layer l. .grad — the gradient, a Float32Array of the same shape, filled during the backward pass. For an intermediate tensor (h_l), .grad holds the activation gradient δ_l = ∂Loss/∂h_l. For a parameter tensor (W_l), .grad holds the weight gradient ∂Loss/∂W_l. Same field, different meaning. null until backward() reaches this tensor. .inputs — the list of tensors this one was computed from. This is the backward edge of the computation graph: "to push gradient into me, push through these." .backward — a closure registered when this tensor is created by an operation. It reads this.grad (already filled by the node closer to the loss) and accumulates the correct gradient into each tensor in .inputs. Each closure implements one application of the chain rule (equations 2–3 in backprop.html). ensureGrad() lazily allocates the gradient array on first use. Tensors that are never differentiated (e.g. targets, masks) stay null and waste no memory.

10export class Tensor { 11 constructor(data, rows, cols, requiresGrad = false) { 12 this.data = data; // Float32Array — forward values (activations or weights) 13 this.rows = rows; 14 this.cols = cols; 15 this.requiresGrad = requiresGrad; 16 this.grad = null; // Float32Array — activation gradient δ or weight gradient ∂L/∂W 17 this.inputs = []; // tensors this was computed from — graph edges for backward 18 this.backward = null; // closure: reads this.grad, accumulates into inputs' .grad 19 } 21 static zeros(rows, cols, requiresGrad = false) { 22 return new Tensor(new Float32Array(rows * cols), rows, cols, requiresGrad); 23 } 26 ensureGrad() { 27 if (!this.grad) this.grad = new Float32Array(this.rows * this.cols); 28 return this.grad; 29 } 30}

2 · matmul — how one operation computes both gradient types

activation gradient weight gradient Every operation follows a three-part pattern: (1) compute the forward output, (2) record .inputs so the graph knows where gradient should flow, (3) register .backward as a closure that computes and accumulates gradients. matmul is the clearest example of how both gradient types arise from one closure. For out = matmul(a, b):

If a is an activation h_{l−1}: then a.grad after backward = δ_{l−1} = ∂Loss/∂h_{l−1} — the activation gradient that will be used by the layer below to compute its own weight gradient.
If b is a weight matrix W_l: then b.grad after backward = ∂Loss/∂W_l = h_{l−1}ᵀ · δ_l — the weight gradient used by the optimizer to update W_l.

The same closure produces both. The math from backprop.html equation (6) is implemented directly here.

60export function matmul(a, b) { 61 if (a.cols !== b.rows) 62 throw new Error(`matmul shape mismatch: ${a.rows}×${a.cols} · ${b.rows}×${b.cols}`); 65 const out = Tensor.zeros(a.rows, b.cols, true); 66 for (let i = 0; i < a.rows; i++) // forward: compute a·b 67 for (let k = 0; k < a.cols; k++) { 68 const aVal = a.data[i * a.cols + k]; 70 for (let j = 0; j < b.cols; j++) 71 out.data[i * b.cols + j] += aVal * b.data[k * b.cols + j]; 74 out.inputs = [a, b]; // graph edge: out ← a, b 75 out.backward = () => { // closure stored now, called during backward() 76 const outGrad = out.grad; // out.grad = δ_out = ∂Loss/∂out (already filled) 79 const gradA = a.ensureGrad(); // a.grad will hold δ_{l−1} = ∂Loss/∂a (activation gradient) 85 gradA[i * a.cols + k] += outGradVal * b.data[k * b.cols + j]; // δ_{l−1} = δ_out · Wᵀ — activation gradient flowing to the layer below 90 const gradB = b.ensureGrad(); // b.grad will hold ∂Loss/∂W (weight gradient) 96 gradB[k * b.cols + j] += aVal * outGrad[i * b.cols + j]; // ∂Loss/∂W = h_inᵀ · δ_out — weight gradient used by optimizer to update W // += not = : gradients from multiple uses of this tensor accumulate (eq 3, sum of routes) 99 }; // end of closure 100 return out; 101}

3 · crossEntropy — the loss and the gradient at the output layer

surprisal The forward pass of crossEntropy does three things: convert logits to probabilities via softmax, compute surprisal (−log q_θ) for each token position, then average. The average is the cross-entropy H(p, q_θ) — a single scalar, the loss. Why softmax first. The logit for token j is the model's raw score; it can be any real number. Softmax converts all V logits into a proper probability distribution (positive, sums to 1). The subtracted maxLogit is a numerical stability trick — it does not change the output.

315export function crossEntropy(logits, targets) { 316 const numPositions = logits.rows; // T — one prediction per input token position 317 const vocabSize = logits.cols; // V — one logit per vocabulary entry 319 const probs = new Float32Array(numPositions * vocabSize); 320 let totalSurprisal = 0; 322 for (let i = 0; i < numPositions; i++) { // for each token position i 324 let maxLogit = -Infinity; 325 for (let j = 0; j < vocabSize; j++) maxLogit = Math.max(maxLogit, logits.data[i * vocabSize + j]); 326 let sumExp = 0; 327 for (let j = 0; j < vocabSize; j++) { 328 const e = Math.exp(logits.data[i * vocabSize + j] - maxLogit); // subtract max for stability 329 probs[i * vocabSize + j] = e; sumExp += e; 330 } 331 for (let j = 0; j < vocabSize; j++) probs[i * vocabSize + j] /= sumExp; // probs[i,*] is now q_θ(·|context_i) — the model's probability distribution over tokens 334 totalSurprisal += -Math.log(probs[i * vocabSize + targets[i]] + 1e-12); // −log q_θ(correct token at position i) = surprisal: how surprised the model was 335 } 338 const loss = new Tensor(new Float32Array([totalSurprisal / numPositions]), 1, 1, true); // average surprisal over all T positions = H(p, q_θ) = the cross-entropy loss scalar

activation gradient The backward closure here is the starting point for all gradient flow. The loss tensor is the root of the graph. Its .grad is seeded to 1 in backward() (dLoss/dLoss = 1, equation (9) seed). This closure then computes the activation gradient at the logit layer: logits.grad[i,j] = ∂Loss/∂logits[i,j]. The formula is q_θ(j) − p(j) where p is the one-hot true distribution. When the model is correct (q_θ(correct)≈1), this is near zero — no gradient, no update. When the model is wrong, it is large, pointing toward the correct token. This gradient at the logit layer is δ for the output matmul. It will flow backward through every matmul, producing activation gradients for each layer and weight gradients for every weight matrix, via the closure in §2.

339 loss.inputs = [logits]; 340 loss.backward = () => { // called first by backward() — it is the seed 343 const logitGrad = logits.ensureGrad(); // will hold δ at the logit layer = ∂Loss/∂logits 344 const scaleByBatch = loss.grad[0] / numPositions; // loss.grad[0] = 1 (seeded), /T for average 345 for (let i = 0; i < numPositions; i++) 346 for (let j = 0; j < vocabSize; j++) 347 logitGrad[i * vocabSize + j] += 348 scaleByBatch * (probs[i * vocabSize + j] - (j === targets[i] ? 1 : 0)); // q_θ(j) − p(j) — activation gradient at the logit layer (δ for the output matmul) // zero when the model is correct; nonzero when it is wrong — this is the training signal 349 }; 350 return { loss, probs }; 351}

4 · backward() — propagate gradients through every layer

propagation backward() is the traversal engine. It does two things: 1. Topological sort. Depth-first visit starting from the loss. Each tensor is added to topoOrder only after all its inputs, so the list runs from the loss's direct inputs back to the raw embeddings and weight matrices. Reversed, it processes tensors from loss → inputs — i.e. the backward direction. 2. Seed and walk. loss.ensureGrad()[0] = 1 seeds dLoss/dLoss = 1, then the loop calls each closure in reverse order. Each closure reads the activation gradient already placed in its output tensor's .grad, and writes two things:

into the activation input's .grad: the activation gradient δ for the layer below (equation 8: δ_{l−1} = δ_l · J_l)
into the weight input's .grad: the weight gradient ∂Loss/∂W (equation 6: ∂Loss/∂W = h_inᵀ · δ_out)

By the time backward() returns, every weight matrix has its .grad filled with ∂Loss/∂W_l = h_{l−1}ᵀ · (δ_L · J_L ··· J_{l+1}) — the full product from backprop.html equation (10). The optimizer in model.js reads these and subtracts η·grad from each weight.

357export function backward(loss) { 360 const topoOrder = []; 361 const visited = new Set(); 362 function visit(node) { 363 if (visited.has(node)) return; 364 visited.add(node); 365 for (const input of node.inputs) visit(input); // inputs before self 366 topoOrder.push(node); // self last — inputs → outputs order 367 } 368 visit(loss); 371 loss.ensureGrad()[0] = 1; // seed: dLoss/dLoss = 1 (eq 9 starting point) 375 for (let i = topoOrder.length - 1; i >= 0; i--) // walk loss → inputs 376 if (topoOrder[i].backward) topoOrder[i].backward(); // each closure: reads .grad, fills inputs 377}

weight gradient After backward() returns, every parameter tensor (Wqkv, Wout, Wfc1, Wfc2, gain, bias, …) has its .grad holding ∂Loss/∂W. This is the gradient of the cross-entropy loss with respect to that weight — the signal the optimizer uses. It is nonzero because the model's distribution q_θ does not yet match the data distribution p, meaning the activation gradient at the logit layer (§3) is nonzero, which propagates backward through every layer (§4), giving each weight a nonzero weight gradient. When q_θ = p everywhere, all weight gradients become zero and training has converged.

The optimizer in model.js then reads each param.grad and applies the Adam update to param.data.