Try first

Ten classes, not two. The sigmoid gives one number between 0 and 1.

You need ten numbers, each between 0 and 1, that together say how likely each digit is. What must be true of those ten numbers, and how would you build them?

Softmax

They must be positive and they must add to 1, because exactly one digit is correct.

Positive is easy: Math.exp of anything is positive. Adding to 1 is easy too: divide each by the total.

function softmax(z) {
  const m = Math.max(...z);
  const e = z.map(v => Math.exp(v - m));
  const s = e.reduce((a, b) => a + b, 0);
  return e.map(v => v / s);
}

That is the whole derivation. Exponentiate, then normalize.

Subtracting the largest value first changes nothing mathematically, because the same factor cancels top and bottom. It matters numerically: Math.exp(900) is Infinity, and the subtraction guarantees the largest exponent is exp(0). Every library does this and it is worth knowing why.

Softmax with two classes reduces exactly to the sigmoid. It is the same idea with more outputs.

Labels as vectors

Lesson 3070 warned that labeling ten classes 0 to 9 tells the model that 8 is near 9 and far from 1. Here is the fix.

function oneHot(d, k = 10) {
  const v = new Array(k).fill(0);
  v[d] = 1;
  return v;
}

The digit 3 becomes [0,0,0,1,0,0,0,0,0,0]. Every label is the same distance from every other, which is the truth about digits.

The loss is log loss extended to k outputs, and only the correct class contributes, because the others are multiplied by zero.

function crossEntropy(p, target) {
  let s = 0;
  for (let i = 0; i < p.length; i++) s -= target[i] * Math.log(p[i] + 1e-12);
  return s;
}

The gradient is the same as it has always been

Work the chain rule through softmax and cross entropy together, and the same cancellation as lesson 3073 happens.

delta at the output = p - target

Predicted probabilities minus the one-hot target. No softmax derivative to compute, no special cases. The output delta has been “prediction minus target” for every model in this course: linear regression, logistic regression, the XOR network, and now this.

The network

function makeMLP(sizes) {
  const layers = [];
  for (let k = 0; k < sizes.length - 1; k++) {
    const fanIn = sizes[k], fanOut = sizes[k + 1];
    const sd = Math.sqrt(2 / fanIn);
    layers.push({
      W: Array.from({ length: fanOut }, () => Array.from({ length: fanIn }, () => gaussian() * sd)),
      b: new Array(fanOut).fill(0),
      last: k === sizes.length - 2,
    });
  }
  return layers;
}

function forwardMLP(layers, x) {
  const acts = [x];
  let a = x;
  for (const l of layers) {
    const z = l.W.map((row, i) => dotProduct(row, a) + l.b[i]);
    a = l.last ? softmax(z) : z.map(relu);
    acts.push(a);
  }
  return acts;
}

function backwardMLP(layers, x, target) {
  const acts = forwardMLP(layers, x);
  let delta = acts[acts.length - 1].map((p, i) => p - target[i]);
  const grads = new Array(layers.length);

  for (let k = layers.length - 1; k >= 0; k--) {
    const aIn = acts[k];
    grads[k] = { dW: delta.map(d => aIn.map(v => d * v)), db: delta.slice() };
    if (k > 0) {
      const next = new Array(aIn.length).fill(0);
      for (let i = 0; i < aIn.length; i++) {
        let s = 0;
        for (let j = 0; j < delta.length; j++) s += delta[j] * layers[k].W[j][i];
        next[i] = aIn[i] > 0 ? s : 0;
      }
      delta = next;
    }
  }
  return grads;
}

relu in the hidden layers and He initialization to match, exactly as lessons 3100 and 3101 require. Softmax on the output only.

The one new line is aIn[i] > 0 ? s : 0, which is the relu slope: pass the gradient through where the unit was active, block it where it was not.

Train it

const s = split(digits, 0.7, 0.15);

function trainMLP(layers, data, epochs = 30, lr = 0.1, batch = 16) {
  for (let e = 0; e < epochs; e++) {
    for (const b of batches(shuffled(data), batch)) {
      const acc = layers.map(l => ({
        dW: l.W.map(r => r.map(() => 0)), db: l.b.map(() => 0)
      }));
      for (const r of b) {
        const g = backwardMLP(layers, r.x, oneHot(r.y));
        g.forEach((gk, k) => {
          gk.dW.forEach((row, i) => row.forEach((v, j) => acc[k].dW[i][j] += v));
          gk.db.forEach((v, i) => acc[k].db[i] += v);
        });
      }
      layers.forEach((l, k) => {
        l.W.forEach((row, i) => row.forEach((_, j) => l.W[i][j] -= lr * acc[k].dW[i][j] / b.length));
        l.b.forEach((_, i) => l.b[i] -= lr * acc[k].db[i] / b.length);
      });
    }
    log("epoch", e, " train", accuracy(layers, s.train).toFixed(3),
        " val", accuracy(layers, s.val).toFixed(3));
  }
}

function predictDigit(layers, x) {
  const p = forwardMLP(layers, x)[forwardMLP(layers, x).length - 1];
  let best = 0;
  for (let i = 1; i < p.length; i++) if (p[i] > p[best]) best = i;
  return { digit: best, confidence: p[best], probs: p };
}

function accuracy(layers, data) {
  let right = 0;
  for (const r of data) if (predictDigit(layers, r.x).digit === r.y) right++;
  return right / data.length;
}

const net = makeMLP([256, 48, 10]);
trainMLP(net, s.train);
log("test accuracy", accuracy(net, s.test).toFixed(4));

Expect high nineties within about twenty epochs. That is a model with roughly 12800 parameters, trained by code you wrote from scratch, recognizing images.

What to watch

Check the loss before training anything. Ten classes with random weights means every probability is about 0.1, so cross entropy should start near -log(0.1), which is 2.303. If it starts far from that, something is wrong with the output layer, and knowing the expected starting value catches it in one line. This is the first item on the checklist in the next lesson.

Exercises

  1. Print a 10 by 10 confusion matrix. Which digits get confused, and does it match what you saw in lesson 3111?
  2. Try hidden sizes 8, 48 and 200. Where does validation accuracy stop improving?
  3. Add a second hidden layer. Does it help at this size?
  4. Find the ten test images with the lowest confidence and draw them. Would you have got them right?
PreviousNext