Building Micrograd from Scratch in Rust

Why I even built this

Before anything else I want to be clear about what this project is and what it is not. It is not a library. There is nothing on crates.io, there is no nice API, and honestly nobody should ever use this in real code. It exists for one reason: me.

We all reach for the autograd in PyTorch or TensorFlow basically every day and never stop to think about what is actually happening underneath. You call loss.backward() and gradients just appear. I wanted to build that machinery myself, for plain scalar values, from scratch, so that I actually understood how the graph and the backward pass work instead of treating it as magic.

The other reason is a bit of a rant. Everyone is letting AI write their code now. I get why, its fast and it mostly works. But I am a big believer in “if you dont use it you lose it”. Skills are a muscle, and if you never write the hard parts yourself they quietly fade. So I sat down and wrote this by hand, in Rust, specifically because the borrow checker will not let me hand wave my way through shared, mutable graph nodes. That friction is the whole point.

The reference for all of this is micrograd, Andrej Karpathy’s minimalist autograd engine. It implements backpropagation over a dynamically built computation graph, the same idea the big frameworks use, just shrunk down to scalars.

The Value type

Everything is built on one type. ValueData is a scalar plus the extra bookkeeping needed to run backprop later:

1
2
3
4
5
6
7
8
9

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

#[derive(Debug, Clone)]
pub struct Value(Rc<RefCell<ValueData>>);

The public handle is Value, which wraps Rc<RefCell<ValueData>>. This little combo is the whole reason the thing works in Rust. Rc lets several child nodes point at the same parent (shared ownership), and RefCell lets me mutate a node’s gradient later during the backward pass even though its shared. In a language with a garbage collector you get this for free and never notice, in Rust you have to say it out loud.

The picture below is the bit that took me a while to feel comfortable with. A Value is not the data, its just a lightweight handle, basically a smart pointer. Cloning a Value does not copy the number, it bumps the Rc count and hands back another pointer to the exact same ValueData. So when an expression like d = (a * b) + (a + b) reuses a, both the multiply node adn the add node store their own Value in parents, but those handles point at one single shared ValueData for a:

That sharing is the connection. The graph edges are litterally the Value handles sitting inside each node’s parents vector, and because they are all Rcs into the same boxes, when backprop writes a gradient into a through one path and then again through the other, both writes land on the same ValueData. The RefCell is what makes those writes legal while everyone is holding a shared reference.

Forward pass

A leaf node just stores the number and marks the op as None. Every operation builds a brand new Value and remembers its parents. Addition is the simplest, the local derivative with respect to both inputs is just 1:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

impl<'a, 'b> Add<&'b Value> for &'a Value {
    type Output = Value;
    fn add(self, other: &'b Value) -> Value {
        Value(Rc::new(RefCell::new(ValueData {
            data: self.data() + other.data(),
            op: Operation::Add,
            gradient: 0.0,
            parents: vec![self.clone(), other.clone()],
        })))
    }
}

Mul and Sub follow the exact same pattern, only the op tag and the arithmetic change. The one new operation worth showing is relu, since it is the first non linear thing in the crate:

1
2
3
4
5
6
7
8

pub fn relu(&self) -> Value {
    Value(Rc::new(RefCell::new(ValueData {
        data: self.data().max(0.0),
        op: Operation::Relu,
        gradient: 0.0,
        parents: vec![self.clone()],
    })))
}

That is the whole forward story. Build values, combine them, and every node quietly records how it was made.

Backward pass

The backward pass computes gradients with the chain rule:

The idea is simple. Every node knows which operation created it, so it knows how to hand its own gradient down to its parents. Here is that happening one step at a time for the expression

d is the output so its gradient starts at 1.0. Then we walk the graph backwards. Since d = e + c, addition passes the gradient through unchanged into both e and c. For $e = a \cdot b$, multiplication sends the other input to each parent, so a recieves b and b recieves a. And c = a + b sends 1 to both.

The detail that trips everyone up is that gradients accumulate. a shows up in both e and c, so its final gradient is the sum of the contribution from the multiply node and the one from the add node. That is why every arm below uses += and never =:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

pub fn backward(&mut self) {
    match self.op() {
        Operation::Add => {
            for parent in &self.0.borrow().parents {
                parent.0.borrow_mut().gradient += self.gradient();
            }
        }
        Operation::Mul => {
            let data = self.0.borrow();
            let (p1, p2) = (&data.parents[0], &data.parents[1]);
            let (p1_data, p2_data) = (p1.data(), p2.data());
            p1.0.borrow_mut().gradient += self.gradient() * p2_data;
            p2.0.borrow_mut().gradient += self.gradient() * p1_data;
        }
        Operation::Relu => {
            let grad = if self.data() > 1e-9 { self.gradient() } else { 0.0 };
            self.0.borrow_mut().parents[0].0.borrow_mut().gradient += grad;
        }
        // Sub is the same shape as Add, the second parent just gets the negative
        _ => {}
    }
}

Relu is the interesting one. Since its output is max(0, x), the output is positive only when the input was, so the gradient flows through when the neuron was active and gets killed to zero when it wasnt.

Topological sort

Before running backward on anything I need the nodes in the right order, from the output back through everything it depends on. That is what topo_sort gives me:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

pub fn topo_sort(last: Value) -> Vec<Value> {
    let mut sorted: Vec<Value> = Vec::new();
    let mut visited: Vec<Value> = vec![last.clone()];
    let mut stack: Vec<Value> = vec![last.clone()];
    while let Some(node) = stack.last().cloned() {
        let mut pushed_child = false;
        for parent in &node.0.borrow().parents {
            if !visited.contains(parent) {
                visited.push(parent.clone());
                stack.push(parent.clone());
                pushed_child = true;
            }
        }
        if !pushed_child && let Some(done) = stack.pop() {
            sorted.push(done);
        }
    }
    sorted.reverse();
    sorted
}

The sort deliberately does not seed the output gradient itself. Whoever runs the pass sets it, wich keeps the sort doing exactly one job. So back_propagate is just:

1
2
3
4
5

pub fn back_propagate(last: Value) -> Vec<Value> {
    let mut list = topo_sort(last);
    list.iter_mut().for_each(|node| node.backward());
    list
}

A tiny end to end run looks like this. Note the seed gradient is now set by hand:

1
2
3
4
5
6
7
8

let a = Value::new(1.0);
let b = Value::new(2.0);
let d = &(&a * &b) + &(&a + &b); // d = (a * b) + (a + b)

d.update_gradient(1.0);
back_propagate(d);

// da/dd = b + 1 = 3, db/dd = a + 1 = 2

There is also a to_dot helper that walks the graph and spits out Graphviz, so I can render a real picture of whatever expression I built with dot -Tsvg graph.dot -o graph.svg. Super handy for sanity checking that the graph looks the way I think it does.

From scalars to a neural net

Once the scalar engine works, a neural network is honestly just a lot of Values wired together. A neuron holds a weight per input plus a bias, and its forward pass is the dot product followed by an optional ReLU:

1
2
3
4
5
6
7
8

pub fn forward(&self, data: &[Value], activation: bool) -> Value {
    let mut total = self.bias.clone();
    for (w, x) in self.weights.iter().zip(data.iter()) {
        let product = w * x;
        total += &product;
    }
    if activation { total.relu() } else { total }
}

Every w * x and every += here is building nodes in the same graph from before, so backprop through a neuron comes for free. A Layer is just a Vec<Neuron>, and a Model is a Vec<Layer>. The only real logic in the model is that the hidden layers get an activation and the final output layer does not:

1
2
3
4
5
6
7
8

pub fn forward(&self, data: &[Value]) -> Vec<Value> {
    let mut pred = data.to_vec();
    for (i, layer) in self.layers.iter().enumerate() {
        let is_last = i == self.layers.len() - 1;
        pred = layer.forward(&pred, !is_last); // relu everywhere but the output
    }
    pred
}

Each layer and model also exposes a parameters() method that just collects every weight and bias into one flat Vec<Value>, wich is what the training loop reaches for.

The training loop

This is the part I was most excited to get working. For each sample we do a forward pass, build the loss as more Value nodes, backprop once through the whole thing, and then nudge every parameter against its gradient. Classic stochastic gradient descent:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

pub fn fit(&mut self, epochs: i32, xs: Vec<Vec<f32>>, ys: Vec<Vec<f32>>, lr: f32) {
    // xs and ys get wrapped into Values first (left out here)
    for _ in 0..epochs {
        for (x, y) in xs.iter().zip(ys.iter()) {
            let pred = self.forward(x);

            // mean squared error, built out of Values so it lands in the graph
            let mut loss = Value::new(0.0);
            for (p, t) in pred.iter().zip(y) {
                let diff = p - t;
                loss += &(&diff * &diff);
            }
            loss = &loss * &Value::new(1.0 / pred.len() as f32);

            // one backward pass over the entire network
            loss.update_gradient(1.0);
            back_propagate(loss);

            // sgd step: push each param downhill, then reset for the next sample
            let params = self.parameters();
            for p in &params {
                p.subtract_value(lr * p.gradient());
            }
            for p in &params {
                p.zero_gradient();
            }
        }
    }
}

The zero_gradient at the end is easy to forget and definitly matters. Because gradients accumulate with +=, if you dont reset them between samples every step is polluted by the last one and training just falls apart.

Does it actually learn?

To check that any of this works I made up a boring linear target and asked the network to fit it:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

fn poly(a: f32, b: f32, c: f32) -> f32 {
    2.0 * a + 3.0 * b - c
}

fn main() {
    let mut rng = rand::rng();

    // 64 random samples of the polynomial above
    let mut data = Vec::new();
    let mut truth = Vec::new();
    for _ in 0..64 {
        let a: f32 = rng.random();
        let b: f32 = rng.random();
        let c: f32 = rng.random();
        data.push(vec![a, b, c]);
        truth.push(vec![poly(a, b, c)]);
    }

    // 3 inputs -> 8 hidden (relu) -> 1 output
    let mut model = Model::new(vec![Layer::new(3, 8), Layer::new(8, 1)]);
    model.fit(200, data, truth, 0.05);
}

Running it, the printed loss slides down toward zero over the epochs, wich means the little engine really is computing correct gradients and the weights are moving the right way. Seeing that number drop on a network I wrote every line of, no framework anywhere, was genuinely satisfying.

What I learned

1. Rc<RefCell<T>> is the right tool for a small shared, mutable graph like this. Nodes need shared ownership and delayed mutation and that combo gives you exactly that.
2. Pointer equality matters. PartialEq uses Rc::ptr_eq, so two handles count as equal only when they point at the same node, not when they happen to hold the same number.
3. The topological order is what makes backprop correct. Without it a node can get processed before its gradient has finished accumulating.
4. A neural net really is not a special thing. Its the scalar engine plus a training loop, and once the gradients are right the rest is just plumbing.

Next steps

• Not sure, i may drop it here or extend to deal with tensors and maybe just build on it whenever I read a new paper and want to implement something from scratch.

This is based on my microgradrs project. Again, its just for me, but if reading it demystifies backprop a little then even better.