Neural Network From Scratch

In this post we will go through the mathematics of neural networks and code from scratch, in Python, a small modular library to build and train such models. The only import we'll ever need is NumPy for dealing with n-dimensional arrays:

import numpy as np

I will assume you already have some knowledge about neural networks. The purpose here is not to explain why, but how.

You can find the entire code of the article on my GitHub:

Preamble

A neural network is nothing but a mathematical function. It's a function with generally many parameters which can be tweaked to change the behavior of that function. For example:

This is the equation of a line, where [begin-latex-inline]a[end-latex-inline] and [begin-latex-inline]b[end-latex-inline] are the parameters that will determine which line exactly this equation represents. In the same way, a neural network is a very generic function that maps [begin-latex-inline]n[end-latex-inline] inputs to [begin-latex-inline]m[end-latex-inline] outputs and has some number of parameters, usually denoted as [begin-latex-inline]W[end-latex-inline], for weights.

Those weights determine what the function does. A certain set of [begin-latex-inline]W[end-latex-inline] will make the function turn an image of a cat or dog (images are just pixels, which are just numbers) into a label 0 for cats and a label 1 for dogs. Another set of [begin-latex-inline]W[end-latex-inline] could make the function predict the next word in a sentence, etc.

So it's really a matter of finding the [begin-latex-inline]W[end-latex-inline] that makes the function behave the way you want. For this, we need 2 things: math and data.

In a nutshell, the idea is to define some loss function that evaluates how bad the network prediction is compared to what you really wanted, and then try to minimize that function using Gradient Descent.

In pseudo-Python code, that would be:

for _ in range(1000):
    y_pred = f(x, weights) # predict output from input
    error = loss(y_true, y_pred) # compute loss between actual output and desired output
    weights_gradient = error.gradient_wrt(weights) # compute derivative of error with respect to weights
    weights -= learning_rate * weights_gradient # update weights such that error is minimized (gradient descent)

Layer by Layer

Our neural network [begin-latex-inline]f_W[end-latex-inline] is usually a composition of many other small functions, i.e. [begin-latex-inline]f_W = f_1(f_2(f_3(x)))[end-latex-inline], each with its own set of parameters.

Those intermediate [begin-latex-inline]f_1[end-latex-inline], [begin-latex-inline]f_2[end-latex-inline], [begin-latex-inline]f_3[end-latex-inline] functions are what we call layers. Since they're a fundamental building block of neural network, we want to model them properly using an abstract class.

Forward Propagation

All layers must be able to return an output, given a certain input:

We can already emphasize one important point which is: since our neural network is a composition of functions, the output of one layer is the input of the next one.

Backward Propagation

Remember that once the forward propagation is complete and we have calculated [begin-latex-inline]E(Y^*, Y)[end-latex-inline], we need each layer to update its parameters. This can only be done if each layer can compute [begin-latex-inline]\frac{\partial E}{\partial W}[end-latex-inline] if [begin-latex-inline]W[end-latex-inline] is the set of parameters of that layer.

Therefore, during backward propagation, when want a layer to be able to return [begin-latex-inline]\frac{\partial E}{\partial X}[end-latex-inline] given [begin-latex-inline]\frac{\partial E}{\partial W}[end-latex-inline]:

Here, I'm abusing the notation of derivatives, because we normally derive with respect to a variable, not a vector. However [begin-latex-inline]X \in \R^n[end-latex-inline] and [begin-latex-inline]Y \in \R^m[end-latex-inline]. This notation should be read as:

Forget about [begin-latex-inline]\frac{\partial E}{\partial X}[end-latex-inline] for now. The reason we want this, is that if we have access to [begin-latex-inline]\frac{\partial E}{\partial Y}[end-latex-inline] then we can very easily calculate [begin-latex-inline]\frac{\partial E}{\partial W}[end-latex-inline] (for gradient descent) without knowing anything about the network architecture! We simply use the chain rule:

Another abuse of notation... What I really mean is:

Since [begin-latex-inline]\frac{\partial E}{\partial y_i}[end-latex-inline] would be given, the layer simply needs to calculate [begin-latex-inline]\frac{\partial y_i}{\partial w}[end-latex-inline] which it knows how to, since for the layer: [begin-latex-inline]Y=f_W(X)[end-latex-inline].

Remember that the output of one layer is the input to the next one! Which means returning [begin-latex-inline]\frac{\partial E}{\partial X}[end-latex-inline] from a given layer, will give the previous layer its own [begin-latex-inline]\frac{\partial E}{\partial Y}[end-latex-inline]. That is how we propagate the gradient through the network. We can again use the chain rule:

This is really the key to understand back-propagation!

Abstract Base Class: Module

class Module:
    """Abstract class for a module in a neural network."""

    def __init__(self):
        self.inputs: np.ndarray = None
        self.outputs: np.ndarray = None

    def forward(self, inputs: np.ndarray) -> np.ndarray:
        raise NotImplementedError()

    def backward(self, output_grad: np.ndarray) -> np.ndarray:
        raise NotImplementedError()

    def parameters(self) -> list[np.ndarray]:
        return []

    def gradients(self) -> list[np.ndarray]:
        return []

As expected we have a forward() and a backward() method, but we also have two extra methods for the parameters() and the gradients() of those parameters. This is because it won't be layer itself that will do the update, but a dedicated training routine. The layer will simply accumulated the weight gradients during the backward method (+=) and should return gradient in the same order as it provides the parameters.

Now let's implement some layers!

Dense Layer

The dense layer is one of the fundamental pieces in neural networks. You'll often find it represented that way:

Where this means to say that each input "neuron" is connected to each output "neuron".

Forward Propagation

The actual value of each output neuron is calculated as the following:

Where [begin-latex-inline]w_{ji}[end-latex-inline] and [begin-latex-inline]b_j[end-latex-inline] are learnable parameters, commonly called weights and biases.

As you can see, this is just a linear transformation, so it can be expressed more succinctly for all output neurons as a simple matrix multiplication:

In effect, this is mapping [begin-latex-inline]n[end-latex-inline] inputs to [begin-latex-inline]m[end-latex-inline] outputs.

Backward Propagation

As we said earlier, suppose we have a matrix containing the derivative of the error with respect to that layer's output [begin-latex-inline]\frac{\partial E}{\partial Y}[end-latex-inline]. We need to compute:

• The derivative of the error with respect to the parameters [begin-latex-inline]\frac{\partial E}{\partial W}[end-latex-inline], [begin-latex-inline]\frac{\partial E}{\partial B}[end-latex-inline]
• The derivative of the error with respect to the input [begin-latex-inline]\frac{\partial E}{\partial X}[end-latex-inline]

Let's calculate [begin-latex-inline]\frac{\partial E}{\partial B}[end-latex-inline]. This matrix should be the same size as [begin-latex-inline]B[end-latex-inline] itself ([begin-latex-inline]m \times 1[end-latex-inline]):

Using the chain rule:

Everything cancels out since the [begin-latex-inline]b_i[end-latex-inline] only appears in [begin-latex-inline]y_i[end-latex-inline]. Also, we can rewrite this more simply as:

Let's move on to the [begin-latex-inline]\frac{\partial E}{\partial W}[end-latex-inline], which should be the same size as [begin-latex-inline]W[end-latex-inline] itself ([begin-latex-inline]m \times n[end-latex-inline]):

Using the chain rule:

Substituting this result in the [begin-latex-inline]\frac{\partial E}{\partial W}[end-latex-inline] matrix:

Lastly, we need to compute [begin-latex-inline]\frac{\partial E}{\partial X}[end-latex-inline] which will be passed to the previous layer during back-propagation!

Again, [begin-latex-inline]\frac{\partial E}{\partial X}[end-latex-inline] should be the same size as [begin-latex-inline]X[end-latex-inline] itself ([begin-latex-inline]n \times 1[end-latex-inline]):

Using the chain rule:

Substituting in [begin-latex-inline]\frac{\partial E}{\partial X}[end-latex-inline]:

Dense Layer Code

We can now create a Python class for the Dense layer:

class Linear(Module):
    """Linear transformation applied to input column vector."""

    def __init__(self, input_size: int, output_size: int):
        super().__init__()
        self.weights = np.random.randn(output_size, input_size)
        self.bias = np.random.randn(output_size, 1)
        self.weights_grad = np.zeros_like(self.weights)
        self.bias_grad = np.zeros_like(self.bias)

    def forward(self, inputs: np.ndarray) -> np.ndarray:
        self.inputs = inputs
        self.outputs = np.dot(self.weights, inputs) + self.bias
        return self.outputs

    def backward(self, output_grad: np.ndarray) -> np.ndarray:
        self.weights_grad += np.dot(output_grad, self.inputs.T)
        self.bias_grad += output_grad
        return np.dot(self.weights.T, output_grad)

    def parameters(self) -> list[np.ndarray]:
        return [self.weights, self.bias]

    def gradients(self) -> list[np.ndarray]:
        return [self.weights_grad, self.bias_grad]

We have completed a milestone in our neural network library! The rest should be a lot easier!

So far, we are able to perform linear transformation on our data. However, the effect of applying many linear transformations in a row is nothing but a linear transformation itself! We need to add non-linearity in the neural network by applying non-linear functions to the data. This is where activation layers come in the picture.

Activation Layer

An activation layer simply passes each of its inputs into a non-linear activation function that we define. Some commonly used activation functions are [begin-latex-inline]tanh(x)[end-latex-inline] or [begin-latex-inline]max(0, x)[end-latex-inline].

Forward Propagation

For a given input [begin-latex-inline]X[end-latex-inline], the output is the activation function applied to every element of [begin-latex-inline]X[end-latex-inline].

Backward Propagation

Similarly as for the Dense layer, we are given [begin-latex-inline]\frac{\partial E}{\partial Y}[end-latex-inline] and we want to compute [begin-latex-inline]\frac{\partial E}{\partial X}[end-latex-inline]. This time, there's no no trainable parameters so that's all we need.

Here, [begin-latex-inline]\odot[end-latex-inline] denotes element-wise multiplication.

Activation Layer Code

We can now create a Python class for the Activation layer:

class Activation(Module):
    """Applies an activation function to the input, element wise."""

    def activation(self, inputs: np.ndarray) -> np.ndarray:
        raise NotImplementedError()

    def activation_prime(self, inputs: np.ndarray) -> np.ndarray:
        raise NotImplementedError()

    def forward(self, inputs: np.ndarray) -> np.ndarray:
        self.inputs = inputs
        return self.activation(inputs)

    def backward(self, output_grad: np.ndarray) -> np.ndarray:
        return output_grad * self.activation_prime(self.inputs)

And implement one such activation, for example [begin-latex-inline]tanh(x)[end-latex-inline]:

class Tanh(Activation):
    """Hyperbolic tangent activation layer."""

    def activation(self, inputs: np.ndarray) -> np.ndarray:
        return np.tanh(inputs)

    def activation_prime(self, inputs: np.ndarray) -> np.ndarray:
        return 1.0 - np.power(np.tanh(inputs), 2)

Sequential Layer

Remember how we said that neural networks are often composition of functions?

The composition itself is of course also a function, so we could implement a Sequential layer who's only job is to do that composition! This will make the code super clean and will allow us to represent the entire neural network as a single layer.

class Sequential(Module):
    """Sequential neural network."""

    def __init__(self, modules: list[Module]):
        super().__init__()
        self.modules = modules

    def forward(self, inputs: np.ndarray) -> np.ndarray:
        outputs = inputs
        for module in self.modules:
            outputs = module.forward(outputs)
        return outputs

    def backward(self, output_grad: np.ndarray) -> np.ndarray:
        input_grad = output_grad
        for module in reversed(self.modules):
            input_grad = module.backward(input_grad)
        return input_grad

    def parameters(self) -> list[np.ndarray]:
        return [param for module in self.modules for param in module.parameters()]

    def gradients(self) -> list[np.ndarray]:
        return [grad for module in self.modules for grad in module.gradients()]

As you can see, we now wrote very explicitly two ideas:

• Forward propagation: the output of one layer is the input of the next one
• Backward propagation: the input gradient of one layer is the output gradient of the previous one

With this class in hand, a neural network can be expressed as easily as this:

# A neural network that maps 10 inputs to 3 outputs
network = Sequential([
    Linear(10, 50),
    Tanh(),
    Linear(50, 3),
    Tanh()
])

Loss Function

Until now, for a given layer, we supposed that [begin-latex-inline]\frac{\partial E}{\partial Y}[end-latex-inline] was given by the next layer. But what happens to the last layer? How does it get the output gradient?

During the training loop, we will pass that gradient manually to the last layer and then propagate the value to all previous layers. However, we still haven't defined what [begin-latex-inline]E[end-latex-inline] is! This is the loss function.

Just like activation functions, there's a panoply of loss functions that you can pick from depending on your problem. In all cases, the loss is computed based on the output of the network and the true desired output.

For now, we will see one such function: MSE, or Mean Squared Error.

Mean Squared Error

For a given output of the neural network [begin-latex-inline]Y = \{y_1, y_2, \ldots, y_n\}[end-latex-inline] and the expected output [begin-latex-inline]Y^* = \{y^*_1, y^*_2, \ldots, y^*_n\}[end-latex-inline], MSE is defined as the squared distance between the two points:

This single number will tell us how good or bad the model is doing and will eventually steer the training of the neural network. When [begin-latex-inline]E=0[end-latex-inline] then both the predicted output and the expected output are equal. This is a suitable loss for regression tasks, for example predicting stock prices.

At last, let's compute [begin-latex-inline]\frac{\partial E}{\partial Y}[end-latex-inline]:

Mean Squared Error Code

Since the idea of a loss function is generic we'll have a base class and one implementation of it for MSE.

class Loss:
    """Abstract class for a loss function."""

    def loss(self, y_true: np.ndarray, y_pred: np.ndarray) -> float:
        raise NotImplementedError()

    def loss_prime(self, y_true: np.ndarray, y_pred: np.ndarray) -> np.ndarray:
        raise NotImplementedError()


class MSE(Loss):
    """Mean squared error loss."""

    def loss(self, y_true: np.ndarray, y_pred: np.ndarray) -> float:
        return np.mean(np.power(y_true - y_pred, 2))

    def loss_prime(self, y_true: np.ndarray, y_pred: np.ndarray) -> np.ndarray:
        return 2.0 * (y_pred - y_true) / np.size(y_true)

Optimizer

So far we have built all the pieces that allow us to make predictions with our neural networks, compute gradients in the network, evaluate a given loss function, but we're never updating the parameters of the network to decrease the error. This too can be abstracted into a general concept: the optimizer.

The optimizer's job is to update the parameters of the network using their associated gradients. Since the network itself is abstracted away as a Module, the optimizer can take that as a parameter. The step() function is where the update logic happens.

class Optimizer:
    """Abstract class for updating the parameters of a module."""

    def __init__(self, module: Module):
        self.module = module

    def step(self):
        raise NotImplementedError()

    def zero_gradients(self):
        for grad in self.module.gradients():
            grad.fill(0)

Stochastic Gradient Descent

One such optimizer is stochastic gradient descent. For each parameter of the network it applies the following update:

[begin-latex-inline]\alpha[end-latex-inline] is a non-trainable scalar parameter that controls the update size. We call it the learning rate.

class SGD(Optimizer):
    def __init__(self, module: Module, learning_rate: float = 0.01):
        super().__init__(module)
        self.learning_rate = learning_rate

    def step(self):
        parameters = self.module.parameters()
        gradients = self.module.gradients()
        for i in range(len(parameters)):
            parameters[i] -= self.learning_rate * gradients[i]

Note: there are 3 variants of gradient descent which have different names:

• Gradient Descent: the original method goes through the entire data, accumulates the gradients, and then does an update of the parameters (the average of the gradients at each step). This is slow in practice.
• Stochastic Gradient Descent: a variation where we do an update on the parameters after every datapoint seen. This is computationally expensive in practice.
• Mini-Batch Gradient Descent: a mix of 1 and 2 where we update the parameters after a predefined number of samples (mini-batch).

The optimizer class is agnostic of when it's being called and therefore this logic will be implemented upstream, in the next section. This is why we have this zero_gradients() method which will help us reset the gradients when we want to so we don't accumulate (+=) indefinitely. Although in practice, you almost certainly want to reset the gradients after an update.

Trainer

We now have all the pieces we need to write the training routine of our neural networks! Let's put this in a function:

def train(
    module: Module,
    x_train: np.ndarray,
    y_train: np.ndarray,
    loss: Loss,
    optimizer: Optimizer,
    epochs: int,
) -> list[float]:
    """Trains the module on the given data, loss, and optimizer."""
    errors = []
    for i in range(epochs):
        error = 0
        for x, y in zip(x_train, y_train):
            output = module.forward(x)
            error += loss.loss(y, output)
            module.backward(loss.loss_prime(y, output))
            optimizer.step()
            optimizer.zero_gradients()

        error /= len(x_train)
        errors.append(error)
        print(f"{i+1}/{epochs} error={error:.5f}")
    return errors

XOR

The hello world of neural networks is to solve XOR - Exclusive Or.

XOR (often denoted as [begin-latex-inline]\oplus[end-latex-inline]) is a function that maps [begin-latex-inline]\{0, 1\} \to \{0, 1\}[end-latex-inline] in the following way:

Because the XOR function is non-linear:

In other words, you can never draw a line that will separate points [begin-latex-inline]\{(0,0), (1,1)\}[end-latex-inline] from points [begin-latex-inline]\{(1,0), (0,1)\}[end-latex-inline].

Therefore, if the neural network correctly learns a function that behaves like XOR for the given on the graph above, then it has learned a non-linear function, which is a very good sign.

x_train = np.array([[0, 0], [0, 1], [1, 0], [1, 1]]).reshape((4, 2, 1))
y_train = np.array([[0], [1], [1], [0]]).reshape((4, 1, 1))

model = Sequential([
    Linear(2, 3), # Map 2 inputs of XOR to 3 temporary units
    Tanh(),       # Activate those 3 temporary units
    Linear(3, 1), # Maps 3 units to 1 output
    Tanh(),       # Activate the output
])

trainer.train(
    model,
    x_train,
    y_train,
    MSE(),
    SGD(model),
    epochs=1000,
)

for x in x_train:
    print(x.tolist(), model.forward(x).tolist())

Output:

1/1000 error=2.19182
2/1000 error=2.10703
3/1000 error=1.88604
...
998/1000 error=0.00022
999/1000 error=0.00022
1000/1000 error=0.00022
[[0], [0]] [[0.0006479610416383107]]
[[0], [1]] [[0.9820342750260771]]
[[1], [0]] [[0.9766218725199783]]
[[1], [1]] [[0.0010051941355460458]]

The model has been trained successfully!

I am sure this is not a satisfying result after going through all that math. So let's solve some more interesting problems with our library... in part 2!