Algebra and gradients

MiniGrad · View code on GitHub ↗

I started with something I could work out on paper. Take two additions, X + Y = Z and A + B = C. Multiply their results to get Z × C = W. With X = 2, Y = 3, A = 4 and B = 1, the answer is 25.

Now the useful question. If I change X a little, how much does W change? X changes Z by the same amount, and Z gets multiplied by C, which is 5. Nudging X by 0.1 raises W by 0.5. The derivative is 5. Click through the graph to trace that effect backward.

That is the basic idea behind backpropagation and how the chain rule comes into play. We just solve the derivative in our current local expression, then multiply it by the derivative passed back from the operation after it.

My first autograd engine, MiniGrad, recorded those operations and worked backward through them automatically. If a value was used twice, I had to add both contributions. For a network, the final number is its loss, a measure of how wrong the prediction was. Backprop gives each weight a gradient. The optimizer uses it to change the weight.

Autograd and tensors

TensorGrad · View code on GitHub ↗

The scalar engine worked! But a network processes whole arrays of numbers. I didn't want to create a separate Python object for every pixel and multiplication, so the next step was TensorGrad.

A tensor is an array. Mine stored its values, a gradient array of the same shape, and a backward function. NumPy handled the array calculations; I wrote the rules for passing gradients back.

The calculus was the same, but now each operation handled a whole matrix.

That carries over to a full layer. Matrix multiplication uses each weight across multiple examples, so backward adds those contributions into its gradient. Broadcasting does the same with a shared bias. Reuse a value going forward, add its contributions going backward.

I checked the backward rules against small numerical changes to the inputs. Then I used TensorGrad to recover y = 2x₁ − 3x₂ + 1 from 32 examples. The learned weights came close to 2, −3, and 1. Now I could start building actual networks.

Neural networks and optimizers

MiniNN · View code on GitHub ↗

TensorGrad could tell me how to change a weight. Now I needed a way to organize a whole network without manually tracking every weight myself. So I built MiniNN.

The basic layer was y = xW + b. W and b are trainable tensors, which I called Parameters. I wrapped the calculation in a Module so a layer knew how to run itself and return the parameters it owned.

from tensorgrad import Tensor

from mininn.parameter import Parameter


class Module:
    def __call__(self, inputs: Tensor) -> Tensor:
        raise NotImplementedError("Must implement __call__ function")

    def parameters(self) -> list[Parameter]:
        return []

    def zero_grad(self) -> None:
        for parameter in self.parameters():
            parameter.grad[...] = 0.0

That's most of the base class! __call__() runs the calculation, parameters() finds the trainable tensors, and zero_grad() clears their previous gradients. I needed that last step because backward adds contributions. Forget to clear them and the next update includes the last batch too.

I could then stack Linear → activation → Linear with Sequential. The activation matters. Without a nonlinearity between them, two linear layers still describe one linear transformation. Adding Tanh or SiLU lets the network learn more interesting functions.

The training loop finally came together. Predict, calculate a loss, clear the gradients, run backward, then update the weights. For basic gradient descent, the update is just weight -= learning_rate × gradient. Backprop calculates the direction of increasing loss, so we take a small step the other way.

I tested MiniNN by training a small MNIST classifier: 784 inputs → 128 hidden features → 10 digit probabilities. The network contained 101,770 trainable parameters and reached 95.56% test accuracy after 30 epochs.

Seeing that MNIST result was reassuring. These weren't just classes that looked like a framework anymore. The gradients and updates were actually training a network. I also built Adam, which tracks recent gradients and their scale to adapt each weight's update. The MNIST check above used SGD.

Autoencoders

Now I could train a network. But could I squeeze a game frame into a small set of numbers and still recover the objects? That became my first autoencoder. The encoder compresses the screen, and the decoder tries to rebuild it. I trained both together.

The encoder was just Linear, Tanh, then Linear, ending in 64 numbers. Those numbers are the latent. The decoder went back the other way, with a sigmoid turning its outputs into probabilities that each cell was occupied.

Here x is the frame, E encodes it, and D decodes the latent z. The sigmoid, σ, turns the output into probabilities. I used binary cross-entropy to compare them with the actual occupied or empty cells. If a cell contains a projectile and the decoder gives it a tiny probability, the loss penalizes that prediction.

It could reconstruct a screen! But that didn't mean it understood what would happen next. For that, I needed actions and history.

Attention and transformers

TinyTransformer · View code on GitHub ↗

One frame wasn't enough to tell me where something was going. I needed a way to read a sequence. Before putting a transformer into the game model, I built a character-level one, TinyTransformer. It was easier to check what was happening when the inputs were letters.

First, To be becomes token IDs, [51, 54, 1, 40, 43]. Each ID selects a row from a learned embedding table. It's a lookup, not a little network converting the number 51 into a vector. I added a learned position vector so the model could also tell where each character appeared.

Then came attention. Each vector goes through three linear projections to make a query, key, and value. I compare a position's query with the available keys. Softmax turns those scores into weights, and I use them to mix the value vectors. That's how an earlier character can affect the current prediction.

There was one obvious way to cheat. During training, the full sequence is right there. So I masked future positions before softmax. The final e can read itself and the earlier characters; the first T can only read itself. Each position has to predict the next character without seeing it.

I ran several attention heads in parallel, joined their outputs, and projected them back to the model width. With width 64 and four heads, each head handles 16 features. Joining them gives 64 again. I didn't assign the heads jobs. Their weights learned from the prediction loss.

A full block adds two more pieces. LayerNorm keeps each token's features normalized, and a residual connection adds the input back to the layer's output. After attention, a small feed-forward network transforms each position separately. I repeated that block to build the transformer.

At the end, a linear layer gives one score per possible next character. Cross-entropy compares those scores with the actual next character, and my autograd engine carries the error back through the whole thing. Embeddings, attention, and feed-forward layers all learn together. For NanoDreamer, I reused the causal sequence idea with game latents and actions, then predicted states, rewards, and continuation instead of letters.

Actor and critic basics

Finally, someone had to choose a move. The actor was another MLP. It read the state and produced a score for each action. Softmax turned those scores into probabilities, and I sampled a move. The critic was a second MLP, estimating how much future reward to expect from that state.

Here's the smallest example. The critic expects 1 point of future return. After the rollout, I calculate a target of 3. That action did 2 better than expected. I train the actor to make it more likely and the critic to move its estimate toward 3.

G is the return target, V(s) is the critic's guess, and A is the difference. The actor multiplies that feedback by the negative log probability of its chosen action. A positive advantage pushes the action's probability up; a negative one pushes it down. The sg means stop-gradient. I treat the feedback as fixed during the actor update.

For the simple critic, I squared the difference between its guess and the target. One detail confused me early on. The target is a return, not just the next reward. It adds discounted rewards across the rollout and can use the critic to estimate what comes after. The final agent predicts a distribution over returns, but it's answering the same question.

The actor update can use rewards from a real game or from our model. But if the model invents the wrong rewards, the actor learns the wrong game. Figuring that out became a large part of NanoDreamer, and it's where the main write-up picks up.