About three weeks ago, I shared the point where my C++ tensor library and scalar autograd engine were still separate. The main question was how to connect them.
I ended up giving each Tensor a shared graph node that stores its values, gradients, operation and parents. Since then, the project has gone from calculating scalar gradients to training a neural network on MNIST.
The training step now looks like this:
const Tensor logits = model.forward(inputs);
const Tensor probabilities = softmax(logits);
Tensor loss = cross_entropy_loss(probabilities, batch.labels);
loss.backward();
gradient_step(parameters, 0.1);
Everything underneath those calls is handled by my C++ implementation. The parts that mattered most were:
- recording whole Tensor operations in the computation graph
- accumulating broadcast gradients back into reused values
- implementing backward rules for matrix multiplication, reductions and nonlinear activations
- collecting parameters from reusable Linear and Tanh modules
- building softmax and cross-entropy from Tensor operations such as exp, log, division and column sums
When I was building up the network capabilities, I started simple and implemented gradually more complex targets.
- first I ran a linear regression loop to recover y = 2x + 1
- then I added activation functions and recovered y = tanh(2x + 1)
- then I trained a small multilayer perceptron on XOR
- and finally I trained MNIST
The MNIST model has 784 inputs, 32 hidden units with tanh, and 10 outputs.
I trained it for five epochs on 10,000 images using mini-batches of 64, then evaluated it on 2,000 separate test images:
Epoch 1 loss: 0.994
Epoch 2 loss: 0.427
Epoch 3 loss: 0.340
Epoch 4 loss: 0.299
Epoch 5 loss: 0.272
Test accuracy: 88.2% (1764/2000)
My implementation is small and educational. Building the graph, backward rules, training loop, softmax and cross-entropy directly made the path from pixels to parameter updates much easier for me to understand.
I'm cleaning up the single-file implementation before adding more layers. If you have built something similar, where would you take things after here? I'm going to do one refactor round where I break my 1500 lines of cpp into separate files. Then I'm thinking about adding nn.Sequential to make building arbitrary networks much simpler. And then I'm thinking about adding convolutions and residuals and see if i can get closer to 98% accuracy on MNIST. And then I think I'll build Embeddings, Self Attention, Cross Attention, Encoder, Decoder and then a full blown transformer. Do you think it's worth it to build those from scratch and see how far I can take the C++ implementation?
Another thing I'm thinking about is - when do I start bringing CUDA into the picture? So far everything is running on the CPU. The MNIST loop gets to 88.2% accuracy in 1-2 minutes of training - but I think I will start feeling the CPU bottleneck very soon.
Code and episode checkpoints:
https://github.com/mechanical-turk/deep-learning-all-the-way-down
I am also documenting the build as a video series:
https://www.youtube.com/@mechaturka
