Picking up where I left off

At the end of the last post I said I might extend the thing to deal with tensors. This is that. The scalar engine trained a tiny network fine, but it did it one number at a time, and the moment you want to feed it something real like an image that stops being cute. An MNIST digit is 784 pixels. Doing 784 separate scalar Values per image, per sample, is miserable and slow.

So the plan was simple to say and annoying to do. Make a Value hold a whole matrix instead of one f32, redo the backward pass so every op pushes matrix gradients, and write my own matmul so the operation actually lands on the graph. Same engine as before, just matrix math instead of scalar math. Most of the real work was in the backward pass and in broadcasting.

From one number to a matrix

In the scalar version a node held a single f32. Now it holds an Array2 from ndarray, and the gradient is a matrix of the same shape.

1
2
3
4
5
6

pub struct ValueData {
    pub data: Array2<f32>,
    pub op: Operation,
    pub gradient: Array2<f32>,
    pub parents: Vec<Value>,
}

The nice part is that the graph itself did not change at all. The Rc<RefCell<ValueData>> handle, the parents vector, the topological sort, all of it stayed exactly the same. The only thing that changed is what a node carries and how each op computes its forward value and hands gradients back. That was reassuring, it meant the hard graph stuff I got right the first time was still right.

Writing my own matmul

ndarray gives you two different products. * is elementwise, and .dot() is the actual matrix multiply. The catch is that * already has a trait implementation, and in Rust you can only implement an operator once for a given pair of types, so matrix multiply can not ride on *. That is fine, it just becomes a plain method that builds a node the same way the operators do.

a (1,2)b (2,3)matmuldata = a · b (1,3)op = MatMulparents = [a, b]each op builds a node that records how it was made

The whole trick lives in that output node. It runs .dot() for the forward value, but it also tags itself with op = MatMul and keeps a and b as its parents, wich is exactly what lets backward work out later how it was built. I also assert that the inner dimensions match, so a bad shape fails with a message I can actually read instead of some panic from deep inside ndarray.

Addition and the broadcasting headache

Addition of two equal sized matrices is the easy case. The forward pass is elementwise, and the local derivative of an add is 1, so the output gradient just flows straight into both parents element by element. The gradient sitting at position ij in the output gets copied into each parents gradient at ij. Nothing clever.

It got more involved the moment the shapes did not match. ndarray broadcasts, so a (1, 3) bias added to a (2, 3) matrix stretches its single row across both rows.

$$ \begin{bmatrix} b_0 & b_1 & b_2 \end{bmatrix} \;\longrightarrow\; \begin{bmatrix} b_0 & b_1 & b_2 \\ b_0 & b_1 & b_2 \end{bmatrix} $$

That matters on the way back. That one bias element b_0 now touched two elements of the output, so its gradient is the sum of both of the output gradients it fed into. So the backward pass is just the mirror image of the forward stretch, you take the output gradient and sum the axis that got broadcasted back down to size one. Forward stretches one row into two, backward folds those two rows of gradient back into one.

$$ \begin{bmatrix} g_{00} & g_{01} & g_{02} \\ g_{10} & g_{11} & g_{12} \end{bmatrix} \;\longrightarrow\; \begin{bmatrix} g_{00}+g_{10} & g_{01}+g_{11} & g_{02}+g_{12} \end{bmatrix} $$

I wrapped that up in a little unbroadcast helper, and the pattern I ended up using in litterally every arm is the same. Start with the output gradient in a temp matrix that has the output shape, then unbroadcast it onto the parents shape. If an axis was size one in the parent and bigger in the output, it gets summed back down. If the shapes already lined up, unbroadcast does nothing and you just copy.

Subtraction and multiply

Subtraction is the same shape as addition, the second parent just gets the negative, so a += on the first and a -= on the second, both through unbroadcast.

Multiply is elementwise, and the derivative of $a \cdot b$ is $b$ for a and $a$ for b. So each parent recieves the output gradient times the other parents data, still elementwise, then unbroadcast if one side was broadcasted. Its the same coeffecient idea from the scalar post, just done per element across the matrix.

The dot product is the one that hurt

Matmul backward is not elementwise anymore, its another matrix multiply, and I had to work out the order and wich side gets transposed. Trying to reason about it index by index melted my brain. Thinking in shapes is what made it click.

Say a is (1, 2) and b is (2, 3), so c = a.b is (1, 3). The output gradient has the same shape as c, so its (1, 3). And the gradient of each input has to match its own data shape, so the gradient of a is (1, 2) and the gradient of b is (2, 3).

Now you just play shape tetris. To build a (1, 2) out of the (1, 3) output gradient and b wich is (2, 3), the only way the dimensions line up is (1, 3) times (3, 2), wich is the output gradient times b transposed. For the gradient of b, to get a (2, 3) I need (2, 1) times (1, 3), wich is a transposed times the output gradient.

$$ \frac{\partial L}{\partial a} = \frac{\partial L}{\partial c} \cdot b^\top \qquad\qquad \frac{\partial L}{\partial b} = a^\top \cdot \frac{\partial L}{\partial c} $$ dc (1,3)·bᵀ (3,2)=da (1,2)aᵀ (2,1)·dc (1,3)=db (2,3)inner 3 = 3, okinner 1 = 1, ok
1
2
3
4
5
6
7

Operation::MatMul => {
    let data = self.0.borrow();
    let p1 = &data.parents[0];
    let p2 = &data.parents[1];
    p1.0.borrow_mut().gradient += &(self.gradient().dot(&p2.data().t()));
    p2.0.borrow_mut().gradient += &(p1.data().t().dot(&*self.gradient()));
}

The thing I like about this is you dont actually have to rederive it every time. The shapes only fit one way, so they hand you the answer.

Wiring up a layer

With the ops in place a layer is boring in a good way. Its a matmul with the weights, then add the bias, then relu.

1
2
3
4
5
6
7
8

pub fn forward(&self, data: &Value, activation: bool) -> Value {
    let mut output = data.matmul(&self.weights);
    output += &self.bias;
    if activation {
        output = output.relu();
    }
    output
}

It worked on the toy, then MNIST happened

First I pointed it at the same polynomial from the last post, 2a + 3b - c, three inputs. The loss slid down to basically zero, same as the scalar version. Good, the tensor ops and the backward pass were correct.

Then MNIST. 784 inputs. The first run the loss started at around three million, and the accuracy came out 10 out of 100. Ten percent is exactly chance for ten classes. And the weird part, the loss kept dropping the whole time while the accuracy never moved off chance.

The network just gave up

So I printed the actual predictions, and every single image gave nearly the same output, about 0.1 in every slot. The network had collapsed into a constant function that spits out the mean of the targets and completely ignores the input.

The 0.1 is not a coincidence. The label is one hot, so across the ten slots the average value in any slot is 0.1. If you take that constant 0.1 and measure mse against a one hot target, nine slots are off by 0.1 and one slot is off by 0.9.

$$ \frac{9 \cdot (0.1)^2 + 1 \cdot (0.9)^2}{10} = \frac{0.09 + 0.81}{10} = 0.09 $$

And 0.09 is exactly where the loss parked itself. So the number was basically confessing that the net had given up and gone to the mean.

Why the weights blew up

The weights were drawn uniform from 0 to 1, so all positive, mean 0.5. With 784 positive pixels times positive weights, every term in the sum adds up and nothing ever cancels, so the pre activation grows with the number of inputs.

$$ \mathbb{E}[z] \approx 784 \cdot 0.13 \cdot 0.5 \approx 51 $$

One neuron sits around 50 before relu, the next layer pushes that into the thousands, and the target is 0 or 1. So the outputs were wildly out of scale before a single gradient step even happened. With three inputs on the poly this never showed up because the sum was tiny. The bug scaled with the input size, it was there the whole time and 784 inputs is just what made it loud.

The gradient side of the same story

The forward blowup is only half of it. The part that actually wrecks training is what it does to the gradients. Very roughly, the gradient for a weight is the output error times the input that fed it. When the output is in the thousands, the error is in the thousands, so the gradient is huge too, and then the sgd step weight -= lr * gradient takes a giant leap.

The way I think about it is a loss landscape with more than one valley. The good solution sits in a deep valley where the network actually uses the pixels, and right next to it theres a wide shallow one thats just the network outputting the mean and ignoring everything, wich only costs 0.09. The important thing, and the part I got wrong in my head at first, is that the loss never actually blew up. It went down the whole time. The trouble was that the giant early steps skated it straight over the deep valley and dropped it into the shallow one, and once its stuck in that flat basin theres almost no slope left to climb back out, so it just sat there at ten percent. Scaling the init keeps the steps small enough to fall into the deep valley instead of skipping past it.

one loss surface, two minimalossweightshallow: the meandeep: learnsuniform init: loss still drops, but into the wrong basin, the mean, 10%scaled init: loss drops into the deep basin that actually learns, 86%

Fixing it with better init

Two small changes to how the weights start.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

pub fn new(input_size: usize, output_size: usize) -> Layer {
    let mut rng = rand::rng();
    let scale = 1.0 / (input_size as f32).sqrt();
    let weights = Array2::from_shape_fn((input_size, output_size), |_| {
        (rng.random::<f32>() - 0.5) * scale
    });
    let bias = Array2::zeros((1, output_size));
    Layer {
        weights: Value::new(weights),
        bias: Value::new(bias),
    }
}

First, subtract 0.5 so the weights are centered on zero. Now about half of them are negative, so the terms in the sum cancel instead of piling up, and a mean zero sum grows like the square root of the count instead of the count itself. Second, multiply by 1 over the square root of the input size, wich is the factor that keeps the pre activation around one no matter how wide the layer gets. For 784 inputs thats about 0.036. This is basically Xavier and He init, and the whole job of it is keeping the numbers at a sane scale so the gradients arent insane on the first step.

The bias I just set to zero. A bias is added once per output, its not summed over the inputs, so it doesnt blow up the same way, but a random positive bias still shoves every output up for no reason, so zero is the clean start.

After that the loss started small and the accuracy jumped to 86 out of 100. Same engine, same data, seperate world just from two lines of init.

What I learned

  1. The whole graph layer did not change moving to tensors. Rc<RefCell<..>>, parents, topo sort, all of it survived. Only what a node holds changed, wich felt like a good sign the original design was sound.
  2. Shapes are enough to derive the matmul backward. You dont grind indices, you match shapes and there is exactly one arrangement that fits.
  3. A dropping loss is not the same as learning. Mine fell seven orders of magnitude while the model was useless. Accuracy was the thing that actually told the truth.
  4. Weight init is not something to disregard (your favorite ml framework does it behind the scenes). The same network went from 10 percent to 86 percent from two lines, and the entire reason was just keeping numbers at a sane magnitude.

Next steps

  • probably nothing on this project, it was a quick exercise.

Based on my microgradrs project. Still just for me, but if it makes the jump from scalars to tensors feel less mysterious then good.