A Simple Neural Network (<200loc, rust)

The very first decision was figuring out how to store the neural network – my very first attempt at doing this involved a lot of pure functions and state being passed around all over the place which quickly became messy – it was much more maintainable to have a single place to represent a neural network.

struct Net {
}

Happily enough, with the decision to have a single input and output, a single hidden layer means that there are only so many weights I need to store.

The overall structure is fairly simple:

    o
  /   \
x - o  - y
  .   .
   . .
    o 

Before drawing in the calculations that need to happen, I'm going to introduce a system for indexing nodes which should make this much easier to talk about and parse:

	  Layer ->
   +------------------------->
   |    
N  |               0       1
e  |               |       |
u  |          0 -  o       .
r  |       /         \     .
o  | 0 - x -  1 -  o - 0 - y
n  |       .         .
|  |        .       .
v  |         ns -  o
   v               

It's a fairly simple system for drawing neurons – a layer index that starts at the first hidden layer; and a per-layer index into the neuron itself.

For each weight in the system, there are going to be 3 coordinates: the neuron it's at, and then the input it's specified for – this is a little bit different than most online courses, because I'm not directly using vectors to be very explicit about the underlying calculations (both avoiding matrix calculations and paying the price in terms of speed as the price).

Drawing in the calculations for evaluating the network forward to make this a little bit more concrete:

	 0    1
	 |    |                   layer input   
				    |     |
     0 - o <------- n[0,0] = relu(w[0, 0, 0]x + b[0, 0, 0])
  /        \                           |
x -  1 - o -  y                      neuron   
  .        .   \
   .      .     relu(sum(w[1, x, i]*n[0, i] + b[1, x, i]))
    ns - o

I also have somewhat stronger constraints on the potential values for my indexes:

• for layer = 0
  • the neuron's index can go from 0 -> ns (number of neurons)
  • the input can only be 0
• for layer = 1
  • the neuron's index can only be 0 (just one output)
  • and the input can go from 0 -> ns (# of neurons in previous layer)

In the code, I often refer to the [layer, neuron, input] index as [x, y, z].

Having this written out explicitly helped me a lot in implementing the network, and I hope the constraints should make sense to you given the structure of the network.

Finally, instead of dealing with initializing and taking care of multi-dimensional vectors or arrays, I'm going to use a trick I often rely on in Advent of Code for complex grid structures and keep a single vector – with a function to translate my custom indexing system into offsets.

struct Net {
    ws: Vec<f64>,
    bs: Vec<f64>,
    ns: usize,
}

impl Net {

    pub fn new(ns: usize) -> Net {
	let size = ns * 2;
	let ws: Vec<f64> = vec![0; size]; // This will change
	let bs: Vec<f64> = vec![0; size]; 

	Net { ws, bs, ns }
    }

    pub fn pt(self: &Self, x: usize, y: usize, z: usize) -> usize {
	match x {
	    0 if z == 0 && y < self.ns => y,
	    1 if y == 0 && z < self.ns => self.ns + z,
	    _ => panic!("Invalid location: {}, {}, {}", x, y, z),
	}
    }
}

Starting the network with all weights initialized to 0 is a good way to get stuck with 0 gradients during gradient descent: and a good experiment to run while running the code yourself.

For this first iteration, I'm hard coding random initialization of weights into the program – but I'll expose another function that accepts custom initial weights in the next iteration for exploration.

+ use rand::thread_rng;
+ use rand::{distributions::Standard, Rng};

  impl Net {

      pub fn new(ns: usize) -> Net {
          let size = ns * 2;
-         let ws: Vec<f64> = vec![0; size]; // This will change
+         let ws: Vec<f64> = thread_rng().sample_iter(Standard).take(size).collect();
-         let bs: Vec<f64> = vec![0; size]; 
+         let bs: Vec<f64> = thread_rng().sample_iter(Standard).take(size).collect();

          Net { ws, bs, ns }
      }

  }

And a corresponding dependency in Cargo.toml:

[dependencies]
rand = "0.7.3"

I admit being disappointed about having to use an external crate that does have some magic, but RNGs are another black box to open and play with for another day.

Up next is evaluating the value predicted by the network at a single point, which is a straight forward traversal through the network. I'm going to use a leaky ReLU as my activation function, so I'll quickly define it:

fn relu(v: f64) -> f64 {
    if v >= 0.0 {
	v
    } else {
	0.01 * v
    }
}

First, I'm going to throw in several utility functions to access or calculate values in different parts of the net – in my first implementation, I extracted these out after implementing back-prop, but there's no harm in putting them up here first.

impl Net {

    /// Relu(w * input + b) for coordinates x, y, z with input val
    fn rwxb(self: &Self, val: f64, x: usize, y: usize, z: usize) -> f64 {
	relu(self.wxb(val, x, y, z))
    }

    /// w * input + b for coordinates x, y, z with input val
    fn wxb(self: &Self, val: f64, x: usize, y: usize, z: usize) -> f64 {
	self.w(x, y, z) * val + self.b(x, y, z)
    }

    fn w(self: &Self, x: usize, y: usize, z: usize) -> f64 {
	self.ws[self.pt(x, y, z)]
    }

    fn b(self: &Self, x: usize, y: usize, z: usize) -> f64 {
	self.bs[self.pt(x, y, z)]
    }

}

Iterating through the net and calculating values lends itself comfortably to Rust's beautiful iterator/expression based system:

impl Net {
    pub fn eval(self: &Self, val: f64) -> f64 {
	relu(
	    (0..self.ns)
		.map(|i| self.rwxb(self.rwxb(val, 0, i, 0), 1, 0, i))
		.sum(),
	)
    }
}

If you look at the diagram, I need to calculate the values of the individual neurons – \(Relu(w * x + b)\) – which can be expressed as rwxb(x, 0, i, 0) with x as the input and i ranging from 0 to ns (number of neurons).

Then, the value of y is the value of these neurons as inputs added up, and passed through Relu again. Which is exactly what the function above does.

(I admit to messing up the naming convention a little bit with x playing double duty as both the input and the first coordinate; I'll try to clean this up later if I find anyone actually reading this post.)

Keeping loss as simple as possible, I'm simply calculating it as \((y - y')^2\) at a data point. The cost, or loss across a given set of data is then the average of the total loss. The implementation is as straightforward as you would expect, given data in the form of tuples.

This is yet another decision to play around with in the future with alternative functions.

impl Net {
    pub fn cost(self: &Self, data: &[(f64, f64)]) -> f64 {
	let mut loss = 0.0;
	for (x, y) in data {
	    let val = self.eval(*x);
	    loss += (y - val).powi(2);
	}
	loss / self.ns as f64
    }
}

Training involves calculating the cost for a given set of inputs, determining the gradients of the cost for that set of inputs in terms of all the weights and biases.

And then updating the weights and biases with the given learning rate – which involves even more decisions and hyperparameters.

This was the hardest part of the whole exercise: packages like Pytorch allow automating the gradient calculation – by swapping out the implementation of all the other mathematical functions, and transparently converting that to the gradient with autodiff.

Instead of building a system to do automatic differentiation, I decided to do things by hand for my simple function. I'll start off with getting some utility functions out of the way to start – the differential of Relu depends on the value of the original function, so I extracted relu_ish out of the implementation.

/// Leaky relu
fn relu(v: f64) -> f64 {
    relu_ish(v, v)
}

/// Leaky relu based on another variable, useful for derivatives
fn relu_ish(v: f64, point: f64) -> f64 {
    if point >= 0.0 {
	v
    } else {
	0.01 * v
    }
}

Some rough calculations for the basis of the next function:

The source code below is after some clean up – after directly writing out the calculations, I simply extracted and re-used some common variables and loops.

The adjustments to be made are first calculated and then applied, to prevent the order of evaluation affecting the rest of the calculation.

impl Net {
  fn backprop(self: &mut Self, data: &[(f64, f64)], learning_rate: f64) {
      let mut dws: Vec<f64> = vec![0.0; self.ns * 2];
      let mut dbs: Vec<f64> = vec![0.0; self.ns * 2];

      for i in 0..self.ns {
	  let pt1 = self.pt(0, i, 0);
	  let pt2 = self.pt(1, 0, i);

	  for (x, y) in data {
	      let yy = self.eval(*x);

	      dws[pt2] += -2.0 * (y - yy) * relu_ish(self.rwxb(*x, 0, i, 0), yy);
	      dbs[pt2] += -2.0 * (y - yy) * relu_ish(1.0, yy);

	      dws[pt1] += -2.0 * (y - yy) * relu_ish(self.ws[pt2] * relu_ish(*x, self.wxb(*x, 0, i, 0)), yy);
	      dbs[pt1] += -2.0 * (y - yy) * relu_ish(self.ws[pt2] * relu_ish(1.0, self.wxb(*x, 0, i, 0)), yy);
	  }
      }

      for i in 0..self.ns {
	  for pt in &[self.pt(1, 0, i), self.pt(0, i, 0)] {
	      self.ws[*pt] -= dws[*pt] * learning_rate;
	      self.bs[*pt] -= dbs[*pt] * learning_rate;
	  }
      }
  }
}

Finally, training is pretty simple: break the input data into pieces to train each batch, and loop through it all. I like to record the total cost at a reasonable interval – to get a total of 10 data points around how cost is proceeding to give me a sense of the nets behavior.

impl Net {
    pub fn train(
	self: &mut Self,
	training_data: &[(f64, f64)],
	epochs: usize,
	batch_size: usize,
	learning_rate: f64,
    ) {
	let log_interval = epochs / 10;

	for epoch in 0..epochs {
	    let mut point = 0;
	    while point <= training_data.len() {
		let limit = min(point + batch_size, training_data.len());
		self.backprop(&training_data[point..limit], learning_rate);
		point += batch_size;
	    }

	    if log_interval > 0 && epoch % log_interval == 0 {
		eprintln!("Epoch {}: {}", epoch, self.cost(training_data));
	    }
	}
    }
}