Backpropagation & Optimization

Key Topics

• 1. The Chain Rule and Computational Graphs
• 2. Deriving Backpropagation
• 3. Stochastic Gradient Descent (SGD)
• 4. Momentum
• 5. Adam Optimizer
• 6. Learning Rate Schedules
• 7. Batch Normalization

Computational Graph View

A neural network defines a directed acyclic graph (DAG) of operations. Each node computes a function of its inputs and passes the result to its children. The key insight of backprop is that we can compute gradients by passing "error signals" backward through this graph.

Forward pass: Compute all intermediate values from inputs to loss.

Backward pass: Starting from $\partial L/\partial L = 1$, propagate gradients backward using the chain rule at each node. Each node multiplies incoming gradient by its local Jacobian.

Forward vs Reverse Mode

For a function $f: \mathbb{R}^n \to \mathbb{R}^m$:

• Forward mode (tangent): Computes one column of the Jacobian per pass — efficient when $n \ll m$
• Reverse mode (adjoint): Computes one row of the Jacobian per pass — efficient when $m \ll n$

Since neural network losses are scalar ($m=1$) with many parameters ($n \gg 1$), reverse mode (backpropagation) is optimal: it computes all$n$ partial derivatives in a single backward pass, at roughly the same cost as the forward pass.

Step 2: Output Layer Error

For MSE loss $\mathcal{L} = \frac{1}{2}\|\mathbf{y} - \mathbf{h}^{(L)}\|^2$ with identity output activation:

For cross-entropy loss with softmax output: $\boldsymbol{\delta}^{(L)} = \hat{\mathbf{p}} - \mathbf{y}$ (one-hot).

Step 3: Backward Recursion

For hidden layers $\ell = L-1, L-2, \ldots, 1$, apply the chain rule:

Since $\mathbf{z}^{(\ell+1)} = \mathbf{W}^{(\ell+1)}\mathbf{h}^{(\ell)} + \mathbf{b}^{(\ell+1)}$and $\mathbf{h}^{(\ell)} = \phi(\mathbf{z}^{(\ell)})$:

where $\odot$ is element-wise multiplication (Hadamard product).

Step 4: Parameter Gradients

Once we have $\boldsymbol{\delta}^{(\ell)}$ for each layer, the parameter gradients are:

Each gradient is an outer product of the error signal and the input to that layer. This is why backprop requires storing all forward-pass activations — memory scales linearly with depth.

Computational Complexity

For a network with $P$ total parameters:

• Forward pass: $O(P)$ multiply-add operations
• Backward pass: $O(P)$ multiply-add operations (roughly 2x forward)
• Memory: Must store all $\mathbf{h}^{(\ell)}$ and $\mathbf{z}^{(\ell)}$ from the forward pass

Without backprop, computing gradients by finite differences would cost $O(P^2)$ — intractable for modern networks.

SGD Convergence

For convex functions with learning rate $\eta_t$:

• Fixed LR: SGD converges to a neighborhood of the optimum, with radius $\propto \eta\sigma^2$
• Decreasing LR: If $\sum_t \eta_t = \infty$ and $\sum_t \eta_t^2 < \infty$, SGD converges exactly
• Common schedule: $\eta_t = \eta_0 / (1 + \alpha t)$ satisfies both conditions

Nesterov Accelerated Gradient (NAG)

Nesterov momentum evaluates the gradient at the "look-ahead" position:

This provides a "correction" that improves convergence rate from $O(1/\sqrt{\kappa})$(classical momentum) to $O(1/\kappa)$ steps for condition number $\kappa$ of the loss surface.

Why Bias Correction?

Since $\mathbf{m}_0 = \mathbf{0}$, early estimates are biased toward zero:

Dividing by $(1-\beta_1^t)$ corrects this. For $\beta_1 = 0.9$, the correction is significant for the first ~10 steps, then becomes negligible.

AdamW: Decoupled Weight Decay

Loshchilov & Hutter (2019) showed that L2 regularization and weight decay are not equivalent with Adam. The correct formulation is:

This decouples the weight decay from the adaptive gradient, leading to better generalization. AdamW is the default optimizer for modern transformers.

Backprop Through BatchNorm

The gradient through BatchNorm requires careful derivation. Let $\delta_i = \partial\mathcal{L}/\partial\hat{z}_i$:

This ensures gradients are properly normalized, preventing both vanishing and exploding gradients.

Why Sigmoid Causes Vanishing Gradients

The sigmoid derivative satisfies $\sigma'(z) \leq 0.25$ for all $z$. At each layer, the gradient is multiplied by at most 0.25. After $L$ layers:

For $L = 10$ layers with well-initialized weights ($\|\mathbf{W}\| \approx 1$), gradients shrink by a factor of $0.25^9 \approx 4 \times 10^{-6}$. This is why ReLU ($\phi'(z) = 1$ for $z > 0$) revolutionized deep learning.

Solutions

• ReLU activation: Gradient is exactly 1 for active neurons
• Residual connections: Skip connections add an identity path for gradients
• Batch normalization: Keeps pre-activations in a well-conditioned range
• Gradient clipping: Cap gradient norms to prevent explosion: $\mathbf{g} \leftarrow \mathbf{g} \cdot \min(1, c/\|\mathbf{g}\|)$
• Proper initialization: He/Xavier initialization maintains gradient variance across layers
• Layer normalization: Normalize across features rather than batch, used in transformers

AD Frameworks

Modern deep learning frameworks implement reverse-mode AD automatically:

• PyTorch: Dynamic computation graph (define-by-run), eager execution
• JAX: Functional transformations, supports forward and reverse mode, JIT compilation
• TensorFlow: Static computation graph (with eager mode), XLA compilation

These frameworks compute exact gradients (up to floating point) — not numerical approximations. The user only writes the forward pass; the backward pass is generated automatically.