Backpropagation: Computational Graph Derivation — The algorithm behind modern deep learning

The Mathematics Behind `loss.backward()

This article has a youtube video that visually demonstrates and discusses the mathes through animations. Click Here to view the youtube video.

When training deep learning models in PyTorch, what does the line loss.backward() do and how? Through this one line of code, every parameter in your network has a gradient, that can be used to optimise it. But what is this line actually doing underneath the hood to obtain these gradients.

The answer is backpropagation, an algorithm published back in 1986 (well before most of us were born so basically old). This algorithm rightly won Geoffrey Hinton a nobel prize last year because without this algorithm none of the modern deep learning would work. It allows the gradients of scaler loss with respect to billions (or even over a trillion if the speculation about the size of GPT-4 was true) to computed in a single backward pass. Its importance cannot be overstated and is, without exaggeration, the engine that drives the entire field.

In this article I want to work through backprop from scratch; starting with why we need it in the first place, see how the chain rule enables the gradients to be pushed back through any sequence of operations, and subsequently derive the forward and backward rules for the layers like linear, activations, MSE . The same logic can be extended beyond simple chains to include complex branches and residual connections that make up most modern deep learning architectures. Not only is it powerful for ones understanding of ML but hopefully by the end of this article loss.backward() should feel less like a method you could implement yourself.

Why we need backpropagation

Here’s how typical neural network is trained. We take some data, process it through the network to produce an output, and check how the output compares to the actual ground truth to calculate loss. The weights are subsequently optimised to reduce the loss. That optimisation of the weights is done through gradient descent.

Whilst this looks very simple, look closely at what this equation requires. We need the partial derivative of the loss with respect to each parameter, ∂L/∂θ. For a modern network, the number of parameters is in the billions, meaning we would have to calculate billions of partial derivatives to compute the loss. Computing these derivatives symbolically or estimating through finite difference one at a time would take a tremendous amount of time. Fortunately you don’t need to compute the each partial derivative from scratch, backpropagation provides an extremely efficient way of computing these derivatives. By performing a single forward and single backward pass through the network, the gradient of the loss with respect to every model parameter is computed. This is possible thanks to two key principles, the chain rule which allows the derivates of compositions to be decomposed and the fact that neural networks is usually a composition of layers.

One layer at a time- the chain rule

Let’s consider the simplest case. We have a single layer that computes some output x{out} from an input x{in} and has some parameters θ. This output is used as an input for other layers in the neural network, and eventually leads to a loss function, L. We want to compute ∂L/∂θ, which is how the loss will be impacted by changing the layer’s parameters.

The chain rule says we can write this as:

This one line of equation is the complete idea of backpropagation. The first component is the gradient of the loss function w.r.t this layer’s output; this is what we receive from the subsequent layers in the network. The second component is knowledge of this layer regarding itself, since this layer is the one that computed x_{out} from θ.

If you think that this particular explanation is overly simplified, good. That is exactly the intention. The essence of backpropagation is that there isn’t anything particularly complicated about any one step. Rather, it is the fact that one step, when applied recursively, allows the method to be applied to networks of arbitrary depth and complexity.

A chain of layers, and a useful pattern

Now, let’s make things more realistic. A neural network is not a single layer but rather it consists of a series of layers. The input x_0 is sent to layer 1, which generates the output x_1, and this output is sent to layer 2, which generates x_2, and this process continues until we obtain the final output x_n, which is then used to compute the loss, L.

We want the gradient with respect to every parameter θ. We can write out a few of these gradients, starting from the last layer and working backwards towards the beginning of the network:

If you look at these three expressions carefully, you can see a pattern? Each expression begins with the derivative as computed for the previous value of θ — gradient at θ_{n-1} begins with the same ∂L/∂x_n already computed for θn, gradient at θ{n-2} multiplies it by another term, and so forth going all the way back to θ_1.

And here’s the insight behind the efficiency of backpropagation: if we simply store our running product as we trace back down through the graph, nothing needs to be recomputed. Each layer is visited once, performs a constant number of operations, and sends its gradient upstream. Total cost per layer visit is linear in the number of layers.

Every layer has two functions: forward and backward

Once you spot the pattern above, there’s something truly elegant you will notice. At their core, each layer of any feed-forward neural network consists of two functions. The first function is the forward function: it receives inputs x{in} and parameters θ, and computes an output x{out}.

The second function is the backward function that takes the input gradient flowing in from downstream, g{out} = ∂L/∂x{out}, and outputs two gradients: first, gradient flowing upstream to the next layer, g{in} =∂L/∂x{in} and second, gradient with respect to its own parameters θ.

That is it. If every layer in your network are able to do these two steps, you can chain them up arbitrarily and backprop just works. And its important to bear in mind that this is not just a conceptual convenience it is literally how PyTorch works. Each nn.Module has a both forward method (that has to be defined by us) and a backward computation (generated automatically using the chain rule). The gradients flow through the network backwards based on the rules we just wrote down.

The Complete Algorithm

We can now combine all we have discussed and write a pseudo-code or procedure of the algorithm:

1. Forward Pass: Push the input data through every layer in the network, and cache each intermediate activation x_i at each layer.

2. Initial (seed) Gradient: Compute g_n=∂L/∂x_n, the gradient of the loss with respect to the final layer. This is the starting point of the backward pass.

3. Backward pass: For i going from final layer, n down to the first, 1:

• compute the parameter gradient ∂L/∂θ_i=g_i⋅∂x_i/∂θ_i

• pass the gradient back to next layer g_(i-1)=g_i⋅∂xi/∂x(i-1)

4. Update: For each parameter, use the computed gradient to perform one step of gradient descent.

This process is repeated during the training loop until the model paramters stop varying.

Why use Batches

In practice models are never trained on one example at a time but rather on batchs of (16, 32, 64 and so on). The model loss is calculated from the average over the complete batch:

Because the derivative of a sum is the sum of the derivatives so we can write:

This is very simple — we can compute the gradient for each sample separately and then perform backpropagation on the average of the gradients. This helps with computational efficiency of performing model training on GPUs with parallel computation.

The batching is important beyond efficient GPU utilisation, it helps with model training stability. Each sample gradient is an estimate of the “true” gradient, and averaging over a batch reduces the amount of noise in the estimate from each individual sample and gives a cleaner gradient descent direction. It is this smoothing effect that makes modle training so much more stable than single samples.

Backpropagation through a linear layer

The linear layer is a basic building block of neural networks, so it is worth doing derivation in complete, starting with the forward pass:

where W is N x M, x{in} is M x 1, and x{out} and b are both N x 1. It is worth spending some time on the dimensions of each matrix/vector and visualising the multiplications before we start to think about the gradient computations during the backpass. The aim is to derive three parameters: g_in , ∂L/∂W, and ∂L/∂b during the backpass operation.

The input gradient, g_in

We want ∂L/∂xin[j] for every index j. Since x{in}[j] influences the loss only through its effect on x_{out}, the chain rule gives:

From the forward pass:

The only term in that sum that contains x_in[j] is the one with l=j. Every other term has derivative zero. So:

Substituting this back in:

Which can be written in a cleaner vector form:

The weight gradient

The weight gradient is where things get geometrically interesting and require a some time to visualise correctly to follow . We want ∂L/∂W[i,j] for every entry of W. Lets start with chain rule again:

Looking back at the forward equation, we can say two things about this equation:

• If k≠i, the weight W[i,j] does not appear in x_{out}[k] at all, so the derivative is zero

• If k = i, only the l=j term in the sum contains W[i,j], and it has a coefficient of x_{in}[j].

If we put these together and then collapse the delta:

Every value/entry of ∂L/∂W is the product of a single vlaue of g{out} and a single valueof x{in} . That is the outer product multiplication so we can rewrite this more cleanly as:

g{out} is N×1, x^⊤{in} is 1×M, and their outer product is N×M, which is exactly the right shape for a gradient with respect to W. The dimensional analysis is consistent as one would hope.

The bias gradient

Compared to the weight gradient, the bias is very easy. Starting from the same forward equation

The gradient with respect to the bias turns out to be just the gradient at the output with any additional transformations or calculation since the bias is simply an added on.

Non-Linearity Through Activation Functions

Activation functions like ReLU, sigmoid and tanh add non-linearities to deep learning models otherwise it is just a set of linear operations that collapses to a linear transformation and depth of network adds no additional value to models ability to be more expressive or learn more complicated patterns. The gradient computation for these activation functions are extremely simple, they apply the same scalar function to every entry of the input without any mixing, so the Jacobian is diagonal:

Because the Jacobian is diagonal, multiplying it against any vector just picks up the diagonal entries so no matrix multiplication is required and can be written as a vector operation with elementwise multiplication:

These activation functions along with their derivatives look like:

The activation functions have no parameters, so there are no ∂L/∂θ to compute. The backward pass just transforms g{out} into g{in} and passes it along. If you have read my previous article on RNNs and backpropagation through time, the 1-tanh² (x) term might look familiar since it is the same derivative that appears in the recursive BPTT gradient at every time step in the RNN.

Putting together an MLP

We have now defined everything we need to define an MLP. An MLP is just a chain of linear layers with nonlinearities between them, followed by a loss function at the end. A typical MLP might look like:

In the forward pass: feed x_0 through each layer and cache the activations. The seed gradient for the backward pass is g=2(y_pred -y_true) from the MSE loss function. Then feed back this gradient through the chain of linear layers where each one uses W^⊤ g to pass gradient back and g x^⊤ to accumulate its own weight gradient, each ReLU uses 1[x>0]⊙g. After one full traversal we have a gradient for every parameter in the network and applying one step of the gradient descent results in parameter optimisation. And then repeat for many iterations (batchs and epochs) until the loss stops reducing and stablises.

Here is a typical example of what a linear layer looks like in python:

class Linear:
    def __init__(self, in_dim, out_dim):
        self.W = np.random.randn(out_dim, in_dim) * 0.01
        self.b = np.zeros((out_dim, 1))

    def forward(self, x_in):
        self.x_in = x_in                  # cache for backward
        return self.W @ x_in + self.b

    def backward(self, g_out):
        self.dW = g_out @ self.x_in.T     # outer product
        self.db = g_out                  
        g_in   = self.W.T @ g_out         # W transposed g_out
        return g_in

Three lines inside backward() function are the three gradients we derived earlier. This is what pytorch is doing for every layer type that you use.

Beyond Linear Chains: Branches and Residuals

One final bit of complexity to add on top of what we have discussed is branches. Most networks are not just a set chains, for example transformers have attention, where one vector gets used as a query, a key, and a value all at the same time. The ResNet architecture has skip connections. This non-linearity in the architecture is requires an additional rule that we haven’t discussed so far.

Consider this example: y = f(x) + g(x) then the derivative, using the chain rule can be written as

The key insight here is that when a variable is used in more than one place, the gradients from each use are summed together. This equation falls straight out of the chain rule applied to each path independently. We can apply this rule to the residual connection where the forward pass is:

Because x{in} appears in two places, both inside f and as a skip, the gradient with respect to x{in} has contributions to from both paths:

The second term comes from the skip connection and is the reason residual networks train so well. No matter how badly the gradient through f has decayed (vanishing gradient problem), the residual connection term makes sure that at least some gradient always reaches x_{in}. It is a beautifully simple fix for a problem that stumped researchers for years. You can also think of the residual connection as a first order differential equation with an euler integration but I will let you write it down in equation form and convince yourself that it is the case.

We can now go back to the question we started with: what does loss.backward() do?

It does what we have derived, it walks the computation graph in reverse, applying each layer’s backward function in turn, caching results along the way, summing contributions for variables that has multiple uses, and determining the gradient for every parameter along the way. Regardless of the layer and model architecture, the central idea is the same: the chain rule plus caching plus reverse traversal

What I personally find remarkable about backprop is the simplicity. My university lecturer on optimisation would go on and on about a lot of different optimisation alogorithms and their robustness against getting stuck in local minimas. In comparison backprop is not a clever optimisation algorithm. It is not learning anything about the model or data or problem. It is just the chain rule and yet this single procedure is what lets us train models with hundreds of billions of parameters. The whole AI models that is taking the world by storm rests on the chain rule.

Thanks for reading, and I hope this article gave you a clearer picture of what loss.backward() does and helped you understand the mathes behind backpropagation.

Unless otherwise noted, all images are by the author.