Backpropagation: Applying the Chain Rule

The chain rule is directly applicable to the mechanics of training neural networks. A neural network operates as a large composite function. Its final output (e.g., a classification score or a regression value) depends on the outputs of the preceding layer, which in turn depend on the outputs of the layer before that, extending back to the input data. Each layer computes a function involving weights and biases, which are the parameters that require adjustment.

Training involves minimizing a loss function, let's call it $L$, which measures the difference between the network's prediction and the actual target value. To minimize $L$ using gradient descent (or its variants), we need to calculate the partial derivative of $L$ with respect to every single weight $W$ and bias $b$ in the network: $\frac{\partial L}{\partial W}$ and $\frac{\partial L}{\partial b}$. Given the nested structure, calculating these directly seems daunting. This is where the chain rule becomes indispensable.

Backpropagation is not a new optimization algorithm; it's an efficient algorithm for computing gradients in a neural network. It systematically applies the multivariable chain rule to calculate all the necessary partial derivatives ($\frac{\partial L}{\partial W_{ij}^{[l]}}$ and $\frac{\partial L}{\partial b_i^{[l]}}$ for all layers $l$, neurons $i$, and input connections $j$) by working backward from the loss function.

The Backward Pass Explained

Imagine the forward pass: input data flows through the network, layer by layer, undergoing linear transformations (multiplying by weights, adding biases) and non-linear activations, ultimately producing an output prediction and then a loss value $L$.

Backpropagation reverses this flow:

1. Start at the End: The process begins by calculating the derivative of the loss $L$ with respect to the final output activation of the network, $a^{[L]}$ (where $L$ denotes the final layer). This is usually straightforward, depending on the specific loss function and final activation function used. Let's denote $\frac{\partial L}{\partial a^{[L]}}$ as $\delta a^{[L]}$.

2. Gradient w.r.t. Pre-activation ($z^{[L]}$): Using the chain rule, we find the gradient with respect to the pre-activation linear output $z^{[L]}$ of the final layer. If $a^{[L]} = g(z^{[L]})$, where $g$ is the activation function, then: $\frac{\partial L}{\partial z^{[L]}} = \frac{\partial L}{\partial a^{[L]}} \cdot \frac{\partial a^{[L]}}{\partial z^{[L]}} = \delta a^{[L]} \cdot g'(z^{[L]})$ Let's denote $\frac{\partial L}{\partial z^{[L]}}$ as $\delta z^{[L]}$. The term $g'(z^{[L]})$ is the derivative of the activation function evaluated at $z^{[L]}$.

3. Gradients w.r.t. Parameters ($W^{[L]}, b^{[L]}$): Now that we have $\delta z^{[L]}$, we can find the gradients for the weights $W^{[L]}$ and biases $b^{[L]}$ of the final layer. Since $z^{[L]} = W^{[L]}a^{[L-1]} + b^{[L]}$: $\frac{\partial L}{\partial W^{[L]}} = \frac{\partial L}{\partial z^{[L]}} \cdot \frac{\partial z^{[L]}}{\partial W^{[L]}} = \delta z^{[L]} \cdot (a^{[L-1]})^T$ $\frac{\partial L}{\partial b^{[L]}} = \frac{\partial L}{\partial z^{[L]}} \cdot \frac{\partial z^{[L]}}{\partial b^{[L]}} = \delta z^{[L]} \cdot 1 = \delta z^{[L]}$ (Note: The calculation for $\frac{\partial L}{\partial W^{[L]}}$ involves matrix/vector operations, resulting in a gradient matrix of the same shape as $W^{[L]}$. We often sum $\delta z^{[L]}$ across the batch dimension for $\frac{\partial L}{\partial b^{[L]}}$.)

4. Propagate Gradient to Previous Layer: The truly clever part is propagating the error gradient backward. We need to find $\frac{\partial L}{\partial a^{[L-1]}}$, the gradient with respect to the activations of the previous layer. Again, using the chain rule through $z^{[L]}$: $\frac{\partial L}{\partial a^{[L-1]}} = \frac{\partial L}{\partial z^{[L]}} \cdot \frac{\partial z^{[L]}}{\partial a^{[L-1]}} = \delta z^{[L]} \cdot (W^{[L]})^T$ Let's denote this $\delta a^{[L-1]}$. This tells us how much the activations in layer $L-1$ contributed to the final loss $L$.

5. Repeat for Layer $L-1$: Now we have $\delta a^{[L-1]}$. We can repeat steps 2, 3, and 4 for layer $L-1$:

   • Calculate $\delta z^{[L-1]} = \delta a^{[L-1]} \cdot g'(z^{[L-1]})$.
   • Calculate $\frac{\partial L}{\partial W^{[L-1]}} = \delta z^{[L-1]} \cdot (a^{[L-2]})^T$.
   • Calculate $\frac{\partial L}{\partial b^{[L-1]}} = \delta z^{[L-1]}$.
   • Calculate $\delta a^{[L-2]} = \delta z^{[L-1]} \cdot (W^{[L-1]})^T$.

6. Continue Until Input Layer: This process repeats, stepping backward through each layer until we reach the input layer. At each step $l$, we use the incoming gradient $\delta a^{[l]}$ (calculated from layer $l+1$) to compute $\delta z^{[l]}$, then the gradients for $W^{[l]}$ and $b^{[l]}$, and finally the gradient $\delta a^{[l-1]}$ to pass back to the previous layer.

This systematic backward application of the chain rule ensures that we compute the gradient of the loss $L$ with respect to every parameter in the network efficiently, reusing calculations where possible. The gradients $\frac{\partial L}{\partial W^{[l]}}$ and $\frac{\partial L}{\partial b^{[l]}}$ for all layers are then used in an optimization algorithm like gradient descent to update the parameters:

$W^{[l]} := W^{[l]} - \alpha \frac{\partial L}{\partial W^{[l]}}$ $b^{[l]} := b^{[l]} - \alpha \frac{\partial L}{\partial b^{[l]}}$

where $\alpha$ is the learning rate.

Simplified view of forward propagation (calculating loss) and backward propagation (calculating gradients) in a two-layer network. Backpropagation applies the chain rule layer by layer, starting from the loss, to compute the gradient of the loss with respect to each parameter ($W^{[l]}, b^{[l]}$) and activation ($a^{[l]}$). The red nodes and arrows indicate the flow and calculation of gradients.

Understanding this backward flow and the application of the chain rule at each step is fundamental to comprehending how neural networks learn. The next section will look at how computational graphs can help visualize this process even more clearly.