AXIS — Back-propagation: How Neural Network Assigns Blame

A step by step intuitive walkthrough of how neural networks figure out what is causing the error

Photo by Google DeepMind on Unsplash

This is part of the AXIS series (AI eXplained through Intuitive Systems), where we break AI systems into intuitive, visual steps to make it super easy for anyone to learn the magic of AI in the simplest way possible.

As a part of this exercise we will go through simple linear model with one node in the neural network. Discussion starts with basic math and then moves into how pytorch’s autograd calculates sensitivity of the error to parameters through auto derivative calculations.

Such calculations are properly understood with graph of operations. So I leveraged AI model to generate UI code to help show the forward pass and back propagation in an intuitive way. With some propt engineering and ideas the forward pass and back propagation UI got working pretty fast. The screen shots below are from that application.

If you want the github link for that app please comment and I’ll share the link with you.

1. The Model

Let’s say you use a linear model to predict house price and model looks like

y = wx + b

You predicted house price to be $400k, Actual is $350k. You are off by $50k.

Who is responsible? w or b?

In this model w and b are the parameters.

A deep learning model trying to learn this relationship will need to adjust these parameters.

How does it figure out how much to adjust as a part of its learning?

2. A simple neural network

The above equation models a very simple neural network with one input and one output.

w is slope of the equation or you can think of it like sensitivity

b is offset. It’s also known as bias or baseline

3. What are we optimizing?

A neural network outputs a value, in this case its output is a value for y but then how does it learn?

It learns by understanding what it needs to optimize for. In this case we need to tell the neural network that we want the predicted value of y to be close to the true value.

The error in predicted and actual value in this case is difference between those values and to highlight larger differences we square the error.

And that’s simple mean squared error loss function.

We square the error so big mistakes hurt more

So how do we find the error?

Take a look at the forward pass for the linear model.

The forward pass, i.e., calculating the value of output for given input is given below for

x = 3, w = 2, b = 1, y = 10

For these input

y_hat = wx + b = 3 * 2 + 1 = 7

and

Loss (also written as L below) = (y_hat — y )² = (7 — 10)² = 9

The image below shows the forward pass as calculations done in a way similar to how libraries like pytorch would do to track various steps that will later be used during backpropagation (More on backpropagation later in this section).

In computational graph created by pytorch, mathematical expressions are broken down and stored as a graph of nodes where nodes contain simple mathematical operations.

Loss is computed in the final node in the graph show below.

4. The Core Problem

We know the prediction is wrong (predicted value is 7 but the actual is 10), we also know that we have a loss of 9 from the above image.

But how does the neural network adjust w and b to make better prediction to optimize (or reduce) the loss?

Welcome to backpropagation!

5. Backpropagaion Intuition

The error (or loss) calculated at the output in the forward pass needs to be propagated backward through the computational graph to understand how each variable is contributing to the loss.

In a nutshell if a variable is tweaked a bit (either increased or decreased), how will that impact the loss?

Think of it like the following analogy:

Prediction = Final product

Loss = Quality check failure

Back-propagation = tracing defect backward

Idea is to move backward from the loss to understand contribution of each parameter (w, b) and then adjust them to make better prediction next time.

6. Basic math behind it

We need to measure how sensitive “loss” is to “w”. In other words if we change “w” by a very small amount then what would be the change in “loss”.

Mathematically, this is what is called derivative — dy/dx is derivative of y with respect to x and measures how much y changes if there is a small change in x.

So, the goal is to find dL/dw. , But w doesn’t directly affect loss. It affects prediction, which in turn affects loss.

During forward pass:

w → y_hat → Loss

and in backward pass

Loss (L)→ y_hat → w

So to find d(Loss)/dw or dL/dw we use derivative’s chain rule

*dL/dw = dL/d(y_hat) d(y_hat)/dw**

7. Backpropagation is Chain Rule in Action

Full mathematical derivative of loss with respect to w and b is in given below:

The above math is just for those who really like to find derivative from the first principle.

But libraries like pytorch doesn’t do symbolic derivate as done in the picture above. It uses built-in automatic differentiation engine called torch.autograd.

The diagram below shows how autograd would work step by step in propagating error at the right most node to the determine gradient at each step as it move backward in the computational graph.

As gradient from the very last output is passed to the previous node the incoming value to a node is multiplied by its own change and changes from multiple paths incoming are accumulated to arrive at the final change for that node.

This is roughly how pytorch backpropagation works under the hood.

In the following image the forward pass has finished and state of each node with their values are shown.

For the final node the incoming gradient defaults to 1 as there is no node after it in the forward pass. At this point we are ready to start the backpropagation.

Grad changes from pending to 1 for node h

h contributes -6 based on it’s current value of g and it’s derivative with respect to g

dh/dg = d(g²)/dg = 2g

For g = -3, dh/dg = -6

-6 is what’s passed to the previous node.

Similar math is explained in each screen shot below as the derivative value propagates backward from node to its parents all the way to the nodes that represent parameters w and b.

Step by step screenshots are as follows:

As you can see the gradient for node a that present w is -18, just as determined by hand in the first image where derivatives were calculated by symbolic differentiation.

You can also see that

as w changes from 2 to 2.01,

loss changes from 9 to 8.8209 so

dL/dw = (change in loss)/(change in w) = -17.91 which is pretty close to -18.

If you choose even smaller increment for w (such as 0.0001) the calculation will come very very close to -18.

This means a very small change in w leads to a change of -18 when w is 2 and other parameters at their given values as shown in the image.

8. Meaning of the Gradient

Large gradient for a parameter means it strongly responsible for the loss

Small gradient for a parameter means it has a weak influence on the loss

Positive of negative sign indicates the direction of influence

9. What We Achieved

We started with a simple model.

Defined what to optimize: error, aka loss

Performed a forward pass to get the value of loss

Traced the loss backward to quantify responsibility of w and b on the loss.

10. So what now?

Now we know ho much w is responsible, but how do we use this information to improve the prediction?

This is where parameter update comes into picture using algorithm such as gradient descent, which will be covered in the next article.

Please share your feedback (likes/dislikes) and anything else you would like to be explained using intuitive visuals in the deep learning domain.