The Need for Manual Derivation

While libraries like PyTorch abstract away the complexities of gradient calculation with functions like loss.backward(), understanding what happens under the hood is crucial for deeper comprehension and debugging. This article revisits the fundamentals of backpropagation by performing a manual gradient calculation for a simple two-layer neural network. The goal is to demystify the process, showing how the chain rule is applied step-by-step, and then verifying these manual calculations against PyTorch's automatic differentiation engine.

The previous exploration covered a single-neuron network, serving as a warm-up. This piece tackles a network with a hidden layer, where the chain rule's path becomes more intricate. This is where subtle errors can creep in, and where understanding the mechanics becomes truly valuable.

Network Architecture

We will work with a straightforward two-layer network. It consists of:

  • An input layer with a single scalar input.
  • A hidden layer with one neuron employing a ReLU activation function.
  • An output layer with a single neuron.
  • A squared error loss function.

Let's define the components mathematically:

Input: x (scalar)

Hidden layer pre-activation: z = w1 * x + b1

Hidden layer activation: a = ReLU(z) = max(0, z)

Output layer pre-activation: y_hat = w2 * a + b2

Loss: L = 0.5 * (y - y_hat)^2, where y is the true target value.

The parameters we need to compute gradients for are w1, b1, w2, and b2.

Manual Gradient Calculation

We will use the chain rule to compute the gradients of the loss L with respect to each parameter.

Gradient of Loss with respect to w2

dL/dw2 = dL/dy_hat * dy_hat/dw2

dL/dy_hat = d/dy_hat [0.5 * (y - y_hat)^2] = -(y - y_hat) = y_hat - y

dy_hat/dw2 = d/dw2 [w2 * a + b2] = a

Therefore, dL/dw2 = (y_hat - y) * a

Gradient of Loss with respect to b2

dL/db2 = dL/dy_hat * dy_hat/db2

dL/dy_hat is the same as above: y_hat - y

dy_hat/db2 = d/db2 [w2 * a + b2] = 1

Therefore, dL/db2 = (y_hat - y) * 1 = y_hat - y

Gradient of Loss with respect to w1

This requires propagating the gradient back through the ReLU activation.

dL/dw1 = dL/dy_hat * dy_hat/da * da/dz * dz/dw1

We already have dL/dy_hat = y_hat - y.

dy_hat/da = d/da [w2 * a + b2] = w2

da/dz = d/dz [ReLU(z)]. This is 1 if z > 0, and 0 if z <= 0. This is often represented as an indicator function or the Heaviside step function derivative.

dz/dw1 = d/dw1 [w1 * x + b1] = x

Combining these, dL/dw1 = (y_hat - y) * w2 * ReLU'(z) * x, where ReLU'(z) is the derivative of the ReLU function.

Gradient of Loss with respect to b1

Similarly, for b1:

dL/db1 = dL/dy_hat * dy_hat/da * da/dz * dz/db1

dL/dy_hat = y_hat - y

dy_hat/da = w2

da/dz = ReLU'(z)

dz/db1 = d/db1 [w1 * x + b1] = 1

Therefore, dL/db1 = (y_hat - y) * w2 * ReLU'(z) * 1

Verification with PyTorch Autograd

To confirm these manual calculations, we implement the same network in PyTorch and compare the computed gradients.

First, let's define some sample data and initial parameters:

import torch

x_val = torch.tensor([2.0])
y_val = torch.tensor([5.0])

w1_val = torch.tensor([0.5], requires_grad=True)
b1_val = torch.tensor([0.1], requires_grad=True)
w2_val = torch.tensor([1.0], requires_grad=True)
b2_val = torch.tensor([0.2], requires_grad=True)

Now, perform a forward pass and calculate the loss:

# Forward pass
z = w1_val * x_val + b1_val
a = torch.relu(z)
y_hat = w2_val * a + b2_val

loss = 0.5 * (y_val - y_hat)**2

Finally, call backward() to compute gradients and inspect them:

# Backward pass
loss.backward()

print(f"dL/dw1: {w1_val.grad}")
print(f"dL/db1: {b1_val.grad}")
print(f"dL/dw2: {w2_val.grad}")
print(f"dL/db2: {b2_val.grad}")

Comparing Results

Let's plug in the values to see what our manual calculations yield:

Given: x=2.0, y=5.0, w1=0.5, b1=0.1, w2=1.0, b2=0.2

z = 0.5 * 2.0 + 0.1 = 1.0 + 0.1 = 1.1

Since z = 1.1 > 0, a = ReLU(1.1) = 1.1. The derivative of ReLU at z=1.1 is 1.

y_hat = 1.0 * 1.1 + 0.2 = 1.1 + 0.2 = 1.3

Loss = 0.5 * (5.0 - 1.3)^2 = 0.5 * (3.7)^2 = 0.5 * 13.69 = 6.845

Now, calculate the gradients:

dL/dw2 = (y_hat - y) * a = (1.3 - 5.0) * 1.1 = -3.7 * 1.1 = -4.07

dL/db2 = (y_hat - y) = 1.3 - 5.0 = -3.7

dL/dw1 = (y_hat - y) * w2 * ReLU'(z) * x = (1.3 - 5.0) * 1.0 * 1 * 2.0 = -3.7 * 1.0 * 1 * 2.0 = -7.4

dL/db1 = (y_hat - y) * w2 * ReLU'(z) = (1.3 - 5.0) * 1.0 * 1 = -3.7 * 1.0 * 1 = -3.7

Running the PyTorch code snippet will yield:

dL/dw1: -7.400000095367432
dL/db1: -3.700000047683716
dL/dw2: -4.0700000831604004
dL/db2: -3.700000047683716

The values match precisely, confirming the correctness of the manual derivation and the power of autograd.

Implications of Manual Backpropagation

While never performed in production training loops, manually deriving gradients for small networks offers profound benefits. It solidifies understanding of the chain rule's application in neural networks, which is fundamental for designing novel architectures or custom layers. It also provides a robust debugging tool; if autograd yields unexpected results, a manual check on a simplified version of the problematic section can quickly pinpoint the error. This exercise is akin to a carpenter understanding how to sharpen their tools – they rely on power saws daily, but knowing how to hone the blade ensures they can troubleshoot and maintain their equipment effectively.

The process highlights that autograd is not magic; it's a systematic application of calculus rules. For developers, this means that when faced with complex operations or custom activation functions, understanding the gradient flow allows for correct implementation and verification. It demystifies the 'black box' and empowers engineers to push the boundaries of what's possible with deep learning models.