Backpropagation and Gradient Descent in Neural Networks

1. The Optimization Problem: Gradient Descent

The process of training a neural network is an optimization problem. We have a loss function, $\mathcal{L}(\theta )$, which depends on all the parameters $\theta$ (weights and biases) of the network. We want to find the set of parameters $\theta$ that makes the loss as small as possible.

The most common algorithm used to solve this is gradient descent. The core idea is simple and intuitive:

1. Calculate the gradient of the loss function with respect to each parameter, ${\nabla }_{\theta }\mathcal{L}$. The gradient is a vector that points in the direction of the steepest ascent of the loss function.
2. To decrease the loss, we take a small step in the opposite direction of the gradient.

This is expressed in the famous update rule, applied to every parameter ${\theta }_i$ in the network:

${\theta }_i\leftarrow {\theta }_i-\eta \frac{\partial \mathcal{L}}{\partial {\theta }_i}$

Here, $\eta$ (eta) is the learning rate, a small positive number that controls the size of our step.

The main challenge, then, is to compute the gradient ${\nabla }_{\theta }\mathcal{L}$, which consists of the partial derivative of the loss with respect to every single weight and bias in the network. For a model with millions of parameters, this is a monumental task. A naive approach, like estimating the derivative for each parameter by slightly perturbing it (finite differences), would be computationally impossible. We need a more efficient way.

2. Backpropagation and Computational Graphs

This is where backpropagation comes in. Backpropagation is a highly efficient algorithm for computing the exact gradient of the loss function with respect to all network parameters. At its core, it is a clever, systematic application of the multivariate chain rule from calculus.

To understand how it works, it's incredibly helpful to visualize the network's calculations as a computational graph. In this graph, nodes represent values (inputs, parameters, intermediate results, the final loss), and edges represent the functions that compute them.

As the video explains, the core operation of backpropagation at any node in the graph is:

$\text{downstream gradient}=\text{upstream gradient}\times \text{local gradient}$

We start at the very end of the graph, where the final loss $\mathcal{L}$ is computed. The "gradient of the loss with respect to itself", $\frac{\partial \mathcal{L}}{\partial \mathcal{L}}$, is trivially 1. This is our initial "upstream gradient". We then work our way backward through the graph, applying this rule at every node, until we have the gradient of the loss with respect to the network's original inputs and parameters.

3. The Mathematics of Backpropagation: Vector-Jacobian Products

The intuitive formula above works for scalar values. Since you're interested in the mathematical foundations, let's formalize this for the vector and matrix operations that dominate neural networks.

When a function takes a vector $\mathbf{x}\in {\mathbb{R}}^n$ and produces a vector $\mathbf{y}\in {\mathbb{R}}^m$, its "local derivative" is the Jacobian matrix, $J_f(\mathbf{x})$, an $m\times n$ matrix of all possible partial derivatives $\frac{\partial y_i}{\partial x_j}$.

The backpropagation rule is then expressed as a Vector-Jacobian Product (VJP).

The key formula is:
${\nabla }_{\mathbf{x}}\mathcal{L}=J_f(\mathbf{x}{)}^T{\nabla }_{\mathbf{y}}\mathcal{L}$

This states that the gradient of the loss with respect to the function's input $\mathbf{x}$ (the downstream gradient) is the transpose of the local Jacobian matrix multiplied by the gradient of the loss with respect to the function's output $\mathbf{y}$ (the upstream gradient).

Backpropagation is simply a chain of these VJP calculations, starting from the loss and moving backward. This method is also called reverse-mode automatic differentiation. Its efficiency comes from the fact that we always perform matrix-vector products, which is significantly cheaper than the matrix-matrix products required by the alternative, forward-mode differentiation. For a typical neural network, where the final output (loss) is a scalar and the parameters are numerous, reverse-mode is vastly more efficient.

4. Backpropagation in Practice: A 2-Layer MLP

Let's apply this to the 2-layer MLP we've been working with. The forward pass is:

1. ${\mathbf{z}}^{(1)}=W^{(1)}\mathbf{x}+{\mathbf{b}}^{(1)}$
2. $\mathbf{h}=\sigma ({\mathbf{z}}^{(1)})$
3. ${\mathbf{z}}^{(2)}=W^{(2)}\mathbf{h}+{\mathbf{b}}^{(2)}$
4. $\hat{\mathbf{y}}=\text{Softmax}({\mathbf{z}}^{(2)})$
5. $\mathcal{L}=-\log ({\hat{y}}_c)$ (Cross-Entropy Loss)

To perform backpropagation, we go in reverse.

Step 1: Gradient of the Loss w.r.t Softmax input (${\mathbf{z}}^{(2)}$)

This is our starting point. We need the derivative of the combined Softmax and Cross-Entropy loss functions. As it turns out, this has a remarkably simple form.

So, our initial upstream gradient is:
${\nabla }_{{\mathbf{z}}^{(2)}}\mathcal{L}=\hat{\mathbf{y}}-\mathbf{y}$
where $\hat{\mathbf{y}}$ is the vector of predicted probabilities and $\mathbf{y}$ is the one-hot encoded true label vector. This vector tells us how much the loss would change for a small change in each logit.

Step 2: Gradients for the Output Layer ($W^{(2)},{\mathbf{b}}^{(2)}$)

Now we backpropagate through the function ${\mathbf{z}}^{(2)}=W^{(2)}\mathbf{h}+{\mathbf{b}}^{(2)}$. Our upstream gradient is ${\nabla }_{{\mathbf{z}}^{(2)}}\mathcal{L}$. We need to calculate the gradients for the parameters $W^{(2)}$ and ${\mathbf{b}}^{(2)}$, and the new upstream gradient to pass back to $\mathbf{h}$.

For this, we use the "VJP trick": we never explicitly form the giant Jacobian matrices. Instead, we use a pre-derived, efficient expression for the VJP of a linear layer.

Using these rules, we get:

• Gradient for weights $W^{(2)}$: (Upstream gradient $\times$ layer input transposed)
  ${\nabla }_{W^{(2)}}\mathcal{L}=(\hat{\mathbf{y}}-\mathbf{y}){\mathbf{h}}^T$
• Gradient for bias ${\mathbf{b}}^{(2)}$: (This is just the upstream gradient, as the bias's local derivative is 1)
  ${\nabla }_{{\mathbf{b}}^{(2)}}\mathcal{L}=\hat{\mathbf{y}}-\mathbf{y}$
• New upstream gradient for $\mathbf{h}$: (Jacobian w.r.t $\mathbf{h}$ is $W^{(2)}$, so the VJP is $(W^{(2)}{)}^T\times$ upstream gradient)
  ${\nabla }_{\mathbf{h}}\mathcal{L}=(W^{(2)}{)}^T(\hat{\mathbf{y}}-\mathbf{y})$

Step 3: Gradients for the Hidden Layer ($W^{(1)},{\mathbf{b}}^{(1)}$)

We continue backward.

• First, we backpropagate through the activation function $\mathbf{h}=\sigma ({\mathbf{z}}^{(1)})$. The local derivative is ${\sigma }^{\prime }({\mathbf{z}}^{(1)})$. Since this is an element-wise operation, the VJP is simply an element-wise product (Hadamard product $\odot$).
  ${\nabla }_{{\mathbf{z}}^{(1)}}\mathcal{L}={\nabla }_{\mathbf{h}}\mathcal{L}\odot {\sigma }^{\prime }({\mathbf{z}}^{(1)})=\left((W^{(2)}{)}^T(\hat{\mathbf{y}}-\mathbf{y})\right)\odot {\sigma }^{\prime }({\mathbf{z}}^{(1)})$
• Now we have the upstream gradient for ${\mathbf{z}}^{(1)}$. We are back at another linear layer, ${\mathbf{z}}^{(1)}=W^{(1)}\mathbf{x}+{\mathbf{b}}^{(1)}$, and we apply the exact same rules as in Step 2:
  ${\nabla }_{W^{(1)}}\mathcal{L}=({\nabla }_{{\mathbf{z}}^{(1)}}\mathcal{L}){\mathbf{x}}^T$
  ${\nabla }_{{\mathbf{b}}^{(1)}}\mathcal{L}={\nabla }_{{\mathbf{z}}^{(1)}}\mathcal{L}$

We have now computed the partial derivative of the loss with respect to every parameter in the network!

5. The Complete Training Algorithm

We can now put everything together into the full training procedure for a neural network. This is the fundamental loop of deep learning.

The algorithm is as follows:

1. Initialize all weights $W$ and biases $\mathbf{b}$ with small random values.
2. Repeat for a set number of iterations (epochs):
   a. Select a batch of training examples ($\mathbf{x},\mathbf{y}$).
   b. Forward Pass: Feed the batch through the network to compute the predicted outputs $\hat{\mathbf{y}}$ and the average loss $\mathcal{L}$ for the batch.
   c. Backward Pass (Backpropagation): Compute the gradient of the loss with respect to every parameter in the network (${\nabla }_W\mathcal{L}$, ${\nabla }_{\mathbf{b}}\mathcal{L}$).
   d. Parameter Update: Adjust each parameter using the gradient descent update rule:
   $W\leftarrow W-\eta {\nabla }_W\mathcal{L}$
   $\mathbf{b}\leftarrow \mathbf{b}-\eta {\nabla }_{\mathbf{b}}\mathcal{L}$

This cycle of forward pass -> loss calculation -> backward pass -> parameter update is repeated thousands or millions of times, gradually nudging the network's parameters towards values that minimize the loss and produce more accurate predictions.

Conclusion

This lesson completes our journey through the fundamental mechanics of a neural network. We have closed the learning loop.

Key Takeaways:

• Gradient Descent is an iterative optimization algorithm that minimizes the loss function by taking small steps in the direction opposite to the gradient.
• Backpropagation is an efficient algorithm that computes the exact gradient of the loss with respect to all network parameters. It does this by applying the chain rule backward through the network's computational graph.
• Mathematically, backpropagation is a sequence of Vector-Jacobian Product (VJP) calculations, which are implemented efficiently without ever materializing the full Jacobian matrices.
• The full training process involves repeatedly executing a loop: a forward pass to get the loss, a backward pass to get the gradients, and an update step to improve the parameters.

You now understand the complete process for training a basic Multi-Layer Perceptron. We have built the essential toolkit that underpins all modern deep learning, including the large language models that are our ultimate goal.

In the next module, we will transition from these general principles to the specific components of the Transformer architecture. We will begin by tackling the first challenge in any language model: how to represent words and their positions in a sequence as numerical vectors that a neural network can process. We will explore embedding layers and positional encodings.