AXON: a neural network library for one person's education

i read "backprop is just the chain rule" about fifty times and i still didn't believe it. so i wrote my own network library. from scratch. python, numpy only. because believing requires touching it.

ref: https://www.kdnuggets.com/2022/04/machine-learning-like-brain-part-one-neurons-slow-slow-slow.html

the whole library is a handful of tiny modules. a neuron, a dense layer, a sequential model, an sgd optimizer. keras-shaped, but 100% hand rolled. no tensorflow, no pytorch. just math i could actually read.

here's a neuron in spirit:

and the part everyone makes scary, the update:

that's the whole learning loop. forward, get a guess, measure how wrong, push the blame back through the weights, nudge each weight by a sliver (lr). do it a few thousand times.

the moment i believed it was the softmax gradient. everyone says "it's just the jacobian" and moves on. i derived it by hand.

the library literally has a test that trains a one-neuron network and asserts the final loss is less than the initial loss. i called that proof. because if it can learn one thing, the machinery is honest.

that one line later showed up in my code and i could point at every symbol and say where it came from. that had never happened with the tutorials.

the proof test, 500 epochs. epoch on x, loss on y. a smooth dip, the only graph that matters

$ python -m unittest tests.test_training -v
test_training_reduces_loss (tests.test_training.TestTrainingProcess.test_training_reduces_loss)
Test if training reduces loss over epochs. ... ok
Epoch 1/5 - Loss: 0.6011
Epoch 2/5 - Loss: 0.3631
Epoch 3/5 - Loss: 0.2777
Epoch 4/5 - Loss: 0.2354
Epoch 5/5 - Loss: 0.2125

----------------------------------------------------------------------
Ran 1 test in 0.070s

OK

proof done, so i sized it up. 784 x 128 x 10 on mnist. relu hidden, softmax out. still numpy, still cpu.

mnist, 12 epochs. loss dives to 0.02, train accuracy climbs, 93.4% of the 10000 held-out digits come back right

here's the part i can't stop staring at. each hidden neuron takes 784 inputs, which is exactly one digit's worth of pixels. so each neuron is secretly a little paintbrush. left to themselves, they painted strokes. nobody showed them an edge or a loop or a crossbar. they found them anyway.

the 128 hidden neurons, each reshaped to 28x28. the network drew these, nobody pointed at what to look for

i didn't finish axon. backprop on a cpu with numpy, no gpu, eventually you stare at a wall. the repo says discontinued with the reason: "can't optimize calculations until i own a gpu". honest.

but i'm glad it exists. because now when someone says the gradient flows backwards, i don't take it on faith. i built the wire. and i watched the signal travel.