Calculating Gradients: The Chain Rule

Training a neural network involves adjusting its weights to minimize a loss function, $L$. Gradient descent requires computing the gradient of this loss concerning every single weight and bias in the network, often denoted as $\nabla L$. For a network with potentially millions of parameters spread across many layers, the challenge is calculating the influence of a weight buried deep within the network on the final loss. The loss isn't a direct function of that specific weight; it's a function of the network's output, which depends on intermediate activations, which in turn depend on earlier activations and weights, forming a long chain of dependencies.

This is where a fundamental concept from calculus comes into play: the chain rule. The chain rule provides a way to calculate the derivative of composite functions, functions nested within one another. This is exactly the situation we have in a neural network.

The Chain Rule: Derivatives of Composite Functions

Let's refresh the basic idea. Suppose we have a simple composition of functions. If a variable $y$ depends on a variable $u$, written as $y = f(u)$, and $u$ itself depends on another variable $x$, written as $u = g(x)$, then $y$ indirectly depends on $x$ through $u$: $y = f(g(x))$.

The chain rule tells us how to find the rate of change of $y$ with respect to $x$, denoted as $dy/dx$. It states that this derivative is the product of the derivatives of the "outer" function with respect to its input and the "inner" function with respect to its input:

$$\frac{dy}{dx} = \frac{dy}{du} \times \frac{du}{dx}$$

Think of it as propagating sensitivities. How much does $y$ change if $x$ changes slightly? It depends on how much $u$ changes when $x$ changes ($du/dx$), multiplied by how much $y$ changes when $u$ changes ($dy/du$).

We can extend this to longer chains. If $z = f(y)$, $y = g(x)$, and $x = h(w)$, then $z$ depends on $w$ through this chain. The derivative of $z$ with respect to $w$ is found by multiplying the derivatives along the path:

$$\frac{dz}{dw} = \frac{dz}{dy} \times \frac{dy}{dx} \times \frac{dx}{dw}$$

Applying the Chain Rule to Neural Networks

Now, let's connect this back to our neural network. For example, a very simple network with one input $x$, one hidden neuron $h$, and one output neuron $y$. Let the calculations be:

1. Hidden neuron's pre-activation: $z_1 = w_1 x + b_1$
2. Hidden neuron's activation: $h = \sigma(z_1)$ (where $\sigma$ is an activation function)
3. Output neuron's pre-activation: $z_2 = w_2 h + b_2$
4. Output neuron's activation (prediction): $y = \sigma(z_2)$
5. Loss calculation: $L = \text{Loss}(y, y_{true})$ (e.g., squared error $L = (y - y_{true})^2$)

Forward pass computation in a simple neural network. Arrows indicate dependencies. Calculating the gradient of the Loss $L$ with respect to an early weight like $w_1$ requires propagating the derivative backward through these dependencies using the chain rule.

Suppose we want to find the gradient of the loss $L$ with respect to the weight $w_1$. $L$ depends on $y$, $y$ depends on $z_2$, $z_2$ depends on $h$, $h$ depends on $z_1$, and $z_1$ depends on $w_1$. Using the chain rule, we can write the derivative $\partial L / \partial w_1$ as:

$$\frac{\partial L}{\partial w_1} = \frac{\partial L}{\partial y} \times \frac{\partial y}{\partial z_2} \times \frac{\partial z_2}{\partial h} \times \frac{\partial h}{\partial z_1} \times \frac{\partial z_1}{\partial w_1}$$

Let's break down each term:

• $\partial L / \partial y$: How the loss changes with the final prediction $y$. Depends on the specific loss function used. For $L = (y - y_{true})^2$, this is $2(y - y_{true})$.
• $\partial y / \partial z_2$: How the output activation changes with its pre-activation value $z_2$. This is the derivative of the output activation function, $\sigma'(z_2)$.
• $\partial z_2 / \partial h$: How the output pre-activation $z_2$ changes with the hidden activation $h$. Since $z_2 = w_2 h + b_2$, this derivative is simply $w_2$.
• $\partial h / \partial z_1$: How the hidden activation changes with its pre-activation value $z_1$. This is the derivative of the hidden layer's activation function, $\sigma'(z_1)$.
• $\partial z_1 / \partial w_1$: How the hidden pre-activation $z_1$ changes with the weight $w_1$. Since $z_1 = w_1 x + b_1$, this derivative is simply $x$.

Putting it all together for this specific example:

$$\frac{\partial L}{\partial w_1} = \underbrace{2(y - y_{true})}_{\partial L / \partial y} \times \underbrace{\sigma'(z_2)}_{\partial y / \partial z_2} \times \underbrace{w_2}_{\partial z_2 / \partial h} \times \underbrace{\sigma'(z_1)}_{\partial h / \partial z_1} \times \underbrace{x}_{\partial z_1 / \partial w_1}$$

Notice how the calculation involves terms computed during the forward pass (like $x, h, y, z_1, z_2$) and the derivatives of the activation functions evaluated at the pre-activation values. The chain rule gives us a systematic way to multiply these local derivatives together to find the overall sensitivity of the loss to a specific weight, no matter how deep in the network it is.

In multi-layer networks with many neurons per layer, the dependencies become more complex (a neuron's output affects multiple neurons in the next layer), involving sums of derivatives. However, the core principle remains the same: the chain rule allows us to compute gradients by propagating derivative information backward through the network, layer by layer. This systematic application of the chain rule is precisely what the backpropagation algorithm achieves, which we will detail in the following sections. Understanding the chain rule is the first significant step towards understanding how neural networks learn.