DATASCI 447 Lecture 9: Training Deep Models

Administrative Stuff

Where We Left Off

• Deep networks learn hierarchical representations
• Depth multiplies expressivity: folds and linear regions grow exponentially
• We saw that deeper networks achieve better accuracy with fewer total units

But: We hand-waved how training actually works

Today: We open the hood

The Optimization Problem

We have a deep network:

\[\hat{y} = f_{\text{out}} \circ f_D \circ \cdots \circ f_1(x)\]

Where:

\[ f_d(\mathbf h_{d - 1}) = \varphi(\mathbf W_d \mathbf h_{d - 1} + \mathbf b_d) \]

And a loss:

\[\mathcal{L} = \frac{1}{N} \sum_{i=1}^N \mathcal{L}(\hat{y}_i, y_i)\]

We need:

\[\theta^* = \arg\min_\theta \mathcal{L}(\theta \mid X, y)\]

The Optimization Problem

Thinking in terms of a nested function:

\[\hat{y} = \sigma_{\text{out}}(\boldsymbol \beta \varphi(\mathbf W_D \varphi(\mathbf W_{D-1} \cdots \varphi(\mathbf W_1 x + \mathbf b_1) \cdots + \mathbf b_{D-1}) + \mathbf b_D) + \boldsymbol \alpha)\]

• Input goes in the first layer

• Project to a new space with \(\mathbf W\) and \(\mathbf b\)

• Apply the activation function \(\varphi\)

• Repeat over and over and over…

The Optimization Problem

We have a set of unknowns:

• \(\mathbf W_1, \mathbf b_1, \ldots, \mathbf W_D, \mathbf b_D\)

• \(\bbeta\) and \(\boldsymbol \alpha\) in the output layer (our regular old linear regression for the final hidden representation)

Goal: Find the values for these parameters given the training data that minimize the loss!

The Optimization Problem

No closed-form solution!

The loss surface is non-convex, high-dimensional, and composed of many nested functions.

We need iterative numerical optimization.

Why Gradient-Based Methods?

We need to minimize:

\[\mathcal{L}(\theta \mid X, y) = \frac{1}{N} \sum_{i=1}^N \mathcal{L}(\hat{y}_i, y_i)\]

over \(\theta = \{\mathbf W_1, \mathbf b_1, \ldots, \mathbf W_D, \mathbf b_D, \bbeta, \boldsymbol \alpha\}\)

• The loss function w.r.t. the parameters is gnarly

Gradient Descent

You’ve seen this in the optimization module — this is a quick review.

The gradient descent “optimal” update:

\[\theta_t = \theta_{t-1} - \eta \nabla_\theta \mathcal{L}(\theta_{t-1})\]

• Why move in the direction of the negative gradient?

• What is \(\eta\)? Why does it matter?

Stochastic Gradient Descent (SGD)

Full gradient descent computes:

\[\nabla_\theta \mathcal{L} = \frac{1}{N} \sum_{i=1}^N \nabla_\theta \mathcal{L}_i\]

This requires a forward pass through all \(N\) observations — expensive!

• Especially when \(N\) is large and each training instance has a lot of features.

• Like images or text documents.

Stochastic Gradient Descent (SGD)

SGD: Approximate the full gradient with a mini-batch of size \(B\):

\[\nabla_\theta \mathcal{L} \approx \frac{1}{B} \sum_{i \in \text{batch}} \nabla_\theta \mathcal{L}_i\]

• Noisy estimate, but unbiased: \(\mathbb{E}[\text{mini-batch gradient}] = \text{full gradient}\)
• Much cheaper: cost per step scales with \(B\), not \(N\)
• The noise can actually help — escaping saddle points, implicit regularization

Question: Why are we more worried about saddle points in high dimensional parameters spaces than local minima?

Second-Order Methods: The Ideal We Can’t Afford

Gradient descent only uses first-order information: the slope.

Newton’s method also uses second-order information: the curvature.

\[\theta_{\text{new}} = \theta - \eta \mathbf H^{-1} \nabla_\theta \mathcal{L}\]

where \(\mathbf H\) is the Hessian — the matrix of second derivatives.

Second-Order Methods: The Ideal We Can’t Afford

The Hessian inverse \(\mathbf H^{-1}\) does two things at once:

1. Rotates the update — accounts for how parameters interact (off-diagonal terms \(\frac{\partial^2 \mathcal{L}}{\partial w_j \partial w_k}\)), straightening the path instead of zig-zagging

2. Rescales the axes — normalizes step sizes per-parameter (diagonal terms \(\frac{\partial^2 \mathcal{L}}{\partial w_j^2}\)), so we move at the same effective speed in all directions

Super efficient!

• Moves directly towards the minimum instead of orthogonal to the contours of the loss surface

• Rescales the update step sizes based on curvature information

Why can’t we directly apply 2nd order optimization methods to deep networks?

Second-Order Methods: The Ideal We Can’t Afford

Why can’t we use this for deep networks?

The Hessian for a network with \(|\theta|\) parameters is a \(|\theta| \times |\theta|\) matrix.

• Our tiny MNIST network: \(3{,}159 \times 3{,}159 \approx 10\) million entries
• A modest modern network with 10M parameters: \(10^7 \times 10^7 = 10^{14}\) entries
• Computing \(\mathbf H\) is \(O(|\theta|^2)\), inverting it is \(O(|\theta|^3)\)

Adam: A Practical Approximation

Adam (Adaptive Moment Estimation) approximates both things the Hessian does, cheaply:

The Engine and the Map

We now have powerful engines: SGD, Momentum, Adam.

They all share the same update structure:

\[\theta_t = \theta_{t-1} - \eta \cdot \text{optimizer}(\nabla_\theta \mathcal{L})\]

The engine (optimizer) determines how we move given gradient information.

The map (gradient) determines where we should go.

The Credit Assignment Problem

For logistic regression: One set of coefficients, clean expression.

For a deep network with \(D\) layers of nested transformations: ???

The Chain Rule Is the Key

Our network is a nested function:

\[\hat{y} = \sigma_{\text{out}}(\bbeta \varphi(\mathbf W_D \varphi(\mathbf W_{D-1} \cdots \varphi(\mathbf W_1 x + \mathbf b_1) \cdots + \mathbf b_{D-1}) + \mathbf b_D) + \boldsymbol \alpha)\]

The chain rule lets us differentiate through compositions:

\[\frac{d}{dx} f(g(x)) = f'(g(x)) \cdot g'(x)\]

For deep networks, this becomes a chain of local derivatives — one per module in the network.

We don’t need to differentiate the entire nested expression at once. We just need each piece to know its own derivative.

The Network as a Chain of Equations

A deep network isn’t one big equation — it’s a sequence of simple steps, each feeding into the next.

Each step is a module: a linear transformation, an activation, a loss computation.

This is exactly how you build a model in PyTorch — stack modules in nn.Sequential:

model = nn.Sequential(
    nn.Linear(P, K1),    # z1 = W1^T x + b1
    nn.ReLU(),           # h1 = ReLU(z1)
    nn.Linear(K1, K2),   # z2 = W2^T h1 + b2
    nn.ReLU(),           # h2 = ReLU(z2)
    nn.Linear(K2, K),    # z_out = beta^T h2 + alpha
)

The Network as a Chain of Equations

Each module knows how to do two things:

1. Compute its output (forward)
2. Compute its local gradient (backward)

The chain of equations flowing up is mirrored by a chain of gradients flowing back down.

The Core Principle

At each module in the chain, we apply:

\[\boxed{\text{Downstream Gradient} = \text{Local Gradient} \times \text{Upstream Gradient}}\]

Gradients vs. Jacobians

When we take derivatives in neural networks, the shape of the result depends on what we’re differentiating.

Gradient: Derivative of a scalar w.r.t. a vector

• \(\frac{\partial \mathcal{L}}{\partial \mathbf w}\) where \(\mathcal{L} \in \mathbb{R}\) and \(\mathbf w \in \mathbb{R}^P\) → result \(\in \mathbb{R}^P\)
• “How much does \(\mathcal{L}\) change if I nudge \(w_j\)?”
• This is what SGD uses to update parameters

Jacobian: Derivative of a vector w.r.t. a vector

• \(\frac{\partial \mathbf h}{\partial \mathbf z}\) where \(\mathbf h \in \mathbb{R}^M\) and \(\mathbf z \in \mathbb{R}^K\) → result \(\in \mathbb{R}^{M \times K}\)
• “How much does \(h_j\) change if I nudge \(z_k\)?”
• Arises at intermediate modules where both input and output are vectors

Gradients vs. Jacobians

Why this matters for backprop:

• The loss \(\mathcal{L}\) is always a scalar, so the final answer for each parameter is always a gradient (same shape as the parameter)

• But the intermediate derivatives between hidden layers are Jacobians

• Backprop chains Jacobians together to produce gradients

The good news: Most Jacobians in neural networks have special structure (diagonal, outer product) so we never actually build the full matrix.

Setting Up a 3-Layer Example

Architecture: \(P\) inputs → \(M_1\) hidden → \(M_2\) hidden → \(M_3\) hidden → \(K\) outputs

\[\underset{(M_1 \times 1)}{\mathbf z_{i1}} = \underset{(M_1 \times P)}{\mathbf W_1^T} \underset{(P \times 1)}{\mathbf x_i} + \underset{(M_1 \times 1)}{\mathbf b_1} \quad \text{[Linear]}\]

\[\underset{(M_1 \times 1)}{\mathbf h_{i1}} = \underset{(M_1 \times 1)}{\varphi(\mathbf z_{i1})} \quad \text{[Activation]}\]

\[\underset{(M_2 \times 1)}{\mathbf z_{i2}} = \underset{(M_2 \times M_1)}{\mathbf W_2^T} \underset{(M_1 \times 1)}{\mathbf h_{i1}} + \underset{(M_2 \times 1)}{\mathbf b_2} \quad \text{[Linear]}\]

\[\underset{(M_2 \times 1)}{\mathbf h_{i2}} = \underset{(M_2 \times 1)}{\varphi(\mathbf z_{i2})} \quad \text{[Activation]}\]

\[\underset{(M_3 \times 1)}{\mathbf z_{i3}} = \underset{(M_3 \times M_2)}{\mathbf W_3^T} \underset{(M_2 \times 1)}{\mathbf h_{i2}} + \underset{(M_3 \times 1)}{\mathbf b_3} \quad \text{[Linear]}\]

Setting Up a 3-Layer Example

\[\underset{(M_3 \times 1)}{\mathbf h_{i3}} = \underset{(M_3 \times 1)}{\varphi(\mathbf z_{i3})} \quad \text{[Activation]}\]

\[\underset{(K \times 1)}{\btheta_i} = \underset{(K \times M_3)}{\bbeta^T} \underset{(M_3 \times 1)}{\mathbf h_{i3}} + \underset{(K \times 1)}{\boldsymbol \alpha} \quad \text{[Linear]}\]

\[\underset{(\text{scalar})}{\mathcal{L}_i} = \text{Loss}(\underset{(K \times 1)}{\sigma_{\text{out}}(\btheta_i)}, \underset{(K \times 1)}{y_i}) \quad \text{[Loss]}\]

Setting Up a 3-Layer Example

Each line is one module. The chain flows up from \(\mathbf x_i\); gradients will flow back down from \(\mathcal{L}_i\).

Setting Up a 3-Layer Example

We can show the chain of partials that will lead to the gradients for each of the unknown weight matrices:

\[\frac{\partial \mathcal{L}_i}{\partial \bbeta} = \frac{\partial \mathcal{L}_i}{\partial \mathbf z_i} \cdot \frac{\partial \mathbf z_i}{\partial \bbeta}\]

\[\frac{\partial \mathcal{L}_i}{\partial \mathbf W_3} = \frac{\partial \mathcal{L}_i}{\partial \mathbf z_i} \cdot \frac{\partial \mathbf z_i}{\partial \mathbf h_{i3}} \cdot \frac{\partial \mathbf h_{i3}}{\partial \mathbf z_{i3}} \cdot \frac{\partial \mathbf z_{i3}}{\partial \mathbf W_3}\]

\[\frac{\partial \mathcal{L}_i}{\partial \mathbf W_2} = \frac{\partial \mathcal{L}_i}{\partial \mathbf z_i} \cdot \frac{\partial \mathbf z_i}{\partial \mathbf h_{i3}} \cdot \frac{\partial \mathbf h_{i3}}{\partial \mathbf z_{i3}} \cdot \frac{\partial \mathbf z_{i3}}{\partial \mathbf h_{i2}} \cdot \frac{\partial \mathbf h_{i2}}{\partial \mathbf z_{i2}} \cdot \frac{\partial \mathbf z_{i2}}{\partial \mathbf W_2}\]

\[\frac{\partial \mathcal{L}_i}{\partial \mathbf W_1} = \frac{\partial \mathcal{L}_i}{\partial \mathbf z_i} \cdot \frac{\partial \mathbf z_i}{\partial \mathbf h_{i3}} \cdot \frac{\partial \mathbf h_{i3}}{\partial \mathbf z_{i3}} \cdot \frac{\partial \mathbf z_{i3}}{\partial \mathbf h_{i2}} \cdot \frac{\partial \mathbf h_{i2}}{\partial \mathbf z_{i2}} \cdot \frac{\partial \mathbf z_{i2}}{\partial \mathbf h_{i1}} \cdot \frac{\partial \mathbf h_{i1}}{\partial \mathbf z_{i1}} \cdot \frac{\partial \mathbf z_{i1}}{\partial \mathbf W_1}\]

Setting Up a 3-Layer Example

The only reason that this is going to work as well as it does in practice is that the partial chains are common until we reach the location in the chain where our parameter of interest lives!

This is what we mean by

\[\boxed{\text{Downstream Gradient} = \text{Local Gradient} \times \text{Upstream Gradient}}\]

• Downstream Gradient: The gradient flowing from the loss back through the network.
• Local Gradient: The gradient of the current module with respect to its inputs.
• Upstream Gradient: The gradient flowing into the current module from the previous module.

The Structure of the Loss

Recall our loss function:

\[\mathcal{L}(\theta \mid X, y) = \frac{1}{N} \sum_{i=1}^N \mathcal{L}_i(\hat{y}_i, y_i)\]

The overall loss is an average of per-observation losses.

The Loss Module

The first step in the backward pass: compute the gradient of the loss w.r.t. the linear predictor \(\btheta_i\).

\[\frac{\partial \mathcal{L}_i}{\partial \btheta_i} = \text{ ???}\]

MSE Loss (Regression)

Output activation: identity (\(\hat{y}_i = \btheta_i\))

Loss:

\[\mathcal{L}_i = \frac{1}{2}(\btheta_i - y_i)^2\]

Binary Cross-Entropy Loss (Binary Classification)

Output activation: sigmoid (\(\hat{y}_i = \sigma(\btheta_i)\))

Loss:

\[\mathcal{L}_i = -[y_i \log \sigma(\btheta_i) + (1 - y_i) \log(1 - \sigma(\btheta_i))]\]

Categorical Cross-Entropy Loss (Multi-Class)

Now \(\btheta_i \in \mathbb{R}^K\) — one score per class.

Output activation: softmax

\[\hat{y}_{ik} = \sigma_{\text{out}}(\btheta_i)_k = \frac{\exp(\btheta_{ik})}{\sum_{g=1}^K \exp(\btheta_{ig})}\]

Loss:

\[\mathcal{L}_i = -\sum_{k=1}^K \mathbb{I}(y_i = k) \log \hat{y}_{ik}\]

Categorical Cross-Entropy Loss (Multi-Class)

Gradient w.r.t. the full vector \(\btheta_i\):

\[\frac{\partial \mathcal{L}_i}{\partial \btheta_i} = \hat{y}_i - e_{y_i} \quad (K \times 1)\]

where \(e_{y_i}\) is a one-hot vector: 1 in position \(y_i\), 0 elsewhere.

Categorical Cross-Entropy Loss (Multi-Class)

\[\frac{\partial \mathcal{L}_i}{\partial \btheta_{ik}} = \hat{y}_{ik} - \mathbb{I}(y_i = k)\]

For the correct class: \(\hat{y}_{ik} - 1\) (negative → push this score up)

For incorrect classes: \(\hat{y}_{ik} - 0 = \hat{y}_{ik}\) (positive → push these scores down)

Again: predicted minus true.

The Unifying Pattern

A Computational Note

Each of these loss gradients is trivially cheap to compute:

• A subtraction of two vectors (or scalars)
• No matrix operations, no inversions, no loops

The Output Linear Module

We’re now one step below the loss in the chain.

The output linear module computes:

\[\underset{(K \times 1)}{\btheta_i} = \underset{(K \times M_3)}{\bbeta^T} \times \underset{(M_3 \times 1)}{\mathbf h_{i3}} + \underset{(K \times 1)}{\boldsymbol \alpha}\]

The upstream gradient arriving from the loss module:

\[\delta_0 = \frac{\partial \mathcal{L}_i}{\partial \btheta_i} \quad (K \times 1)\]

This module needs to produce three things:

1. Gradient for \(\boldsymbol \alpha\) (terminal — used by SGD)
2. Gradient for \(\bbeta\) (terminal — used by SGD)
3. Pass-through gradient for \(\mathbf h_{i3}\) (sent to the module below)

A Note on Notation

We’ll be careful to distinguish two types of multiplication:

Bias Gradient

\[\btheta_i = \bbeta^T \times \mathbf h_{i3} + \boldsymbol \alpha\]

Pass-Through Gradient

\[\btheta_i = \bbeta^T \times \mathbf h_{i3} + \boldsymbol \alpha\]

How does the output change when the input \(\mathbf h_{i3}\) changes?

Weight Gradient — Why It’s Tricky

\[\btheta_i = \bbeta^T \times \mathbf h_{i3} + \boldsymbol \alpha\]

\(\btheta_i \in \mathbb{R}^K\) and \(\bbeta \in \mathbb{R}^{M_3 \times K}\)

Weight Gradient

Write out the elements of \(\btheta_i\):

\[\theta_{i1} = \beta_{11} h_{i3,1} + \beta_{21} h_{i3,2} + \cdots + \beta_{M_3,1} h_{i3,M_3} + \alpha_1\]

\[\theta_{i2} = \beta_{12} h_{i3,1} + \beta_{22} h_{i3,2} + \cdots + \beta_{M_3,2} h_{i3,M_3} + \alpha_2\]

\[\vdots\]

\[\theta_{iK} = \beta_{1K} h_{i3,1} + \beta_{2K} h_{i3,2} + \cdots + \beta_{M_3,K} h_{i3,M_3} + \alpha_K\]

Weight Gradient

\[\theta_{ik} = \beta_{1k} h_{i3,1} + \beta_{2k} h_{i3,2} + \cdots + \beta_{M_3,k} h_{i3,M_3} + \alpha_k\]

\[\frac{\partial \theta_{ik}}{\partial \boldsymbol \beta_{\cdot k}} = \begin{bmatrix} \frac{\partial \theta_{ik}}{\partial \beta_{1k}} \\ \frac{\partial \theta_{ik}}{\partial \beta_{2k}} \\ \vdots \\ \frac{\partial \theta_{ik}}{\partial \beta_{M_3,k}} \end{bmatrix} = \begin{bmatrix} h_{i3,1} \\ h_{i3,2} \\ \vdots \\ h_{i3,M_3} \end{bmatrix} = \mathbf h_{i3}\]

Weight Gradient

Since there’s no cross-talk, the full Jacobian \(\frac{\partial \btheta_i}{\partial \bbeta}\) has a block-diagonal structure.

\[\frac{\partial \btheta_i}{\partial \bbeta} = \begin{bmatrix} \mathbf h_{i3} & \mathbf 0 & \cdots & \mathbf 0 \\ \mathbf 0 & \mathbf h_{i3} & \cdots & \mathbf 0 \\ \vdots & \vdots & \ddots & \vdots \\ \mathbf 0 & \mathbf 0 & \cdots & \mathbf h_{i3} \end{bmatrix}\]

Each block is the same \(\mathbf h_{i3}\) vector — repeated \(K\) times along the diagonal.

We never need to build this full matrix. The structure tells us exactly how to combine things.

Weight Gradient

We need \(\frac{\partial \mathcal{L}_i}{\partial \bbeta}\), which should be \((M_3 \times K)\) — the same shape as \(\bbeta\).

The upstream gradient is \(\delta_0 \in \mathbb{R}^{K \times 1}\), giving us one scalar per output.

Weight Gradient — The Outer Product

Stacking these \(K\) column gradients side by side:

\[\frac{\partial \mathcal{L}_i}{\partial \bbeta} = \begin{bmatrix} | & | & & | \\ \delta_{0,1} \cdot \mathbf h_{i3} & \delta_{0,2} \cdot \mathbf h_{i3} & \cdots & \delta_{0,K} \cdot \mathbf h_{i3} \\ | & | & & | \end{bmatrix}\]

The Linear Module Rules

For any linear module \(\mathbf z = \mathbf W^T \times \mathbf h + \mathbf b\) with upstream gradient \(\delta\):

A Computational Note

What operations did we just use?

• Bias gradient: copy the upstream \(\delta\) (free)
• Weight gradient: outer product (\(\times\)) of two vectors (\(M \times K\) multiplications)
• Pass-through: matrix-vector multiply (\(\times\)) (\(M \times K\) multiplications)

The Activation Module

We’re now below the output linear module in the chain.

The activation module computes:

\[\underset{(M_3 \times 1)}{\mathbf h_{i3}} = \underset{(M_3 \times 1)}{\varphi(\mathbf z_{i3})}\]

Applied elementwise: each element of \(\mathbf z_{i3}\) is transformed independently.

ReLU: The Most Common Activation

\[\varphi(z) = \text{ReLU}(z) = \max(0, z)\]

Applied elementwise: \(h_{i3,m} = \max(0, z_{i3,m})\) for each \(m = 1, \ldots, M_3\)

ReLU Backward

Multiplying by a diagonal Jacobian is just an elementwise multiply — this is where \(\odot\) appears!

ReLU Is a Gate

ReLU’s backward pass is a gate on the gradient flow:

• Active units: the gate is open — gradient flows through
• Dead units: the gate is closed — gradient is blocked

Computationally, this is one elementwise multiply — about as cheap as it gets.

Other Elementwise Activations

The same logic applies to any elementwise activation \(\varphi\).

The Jacobian is always diagonal, so the backward pass is always \(\odot\):

\[\frac{\partial \mathcal{L}_i}{\partial \mathbf z} = \delta \odot \varphi'(\mathbf z)\]

Other Elementwise Activations

Other Elementwise Activations

Other Elementwise Activations

ReLU and its smooth cousins (Leaky ReLU, Swish, GELU) are the default choice for hidden layers:

Leaky ReLU: Prevents dead units by allowing a small gradient when \(z < 0\)

• Use when you suspect dying ReLU is a problem (many units stuck at 0)
• Common in GANs and very deep networks

Swish/SiLU and GELU: Smooth approximations that never fully kill gradients

• GELU is the standard in Transformers (BERT, GPT, etc.)
• Swish used in EfficientNet and mobile architectures
• Both help with vanishing gradients (more next lecture)

Key advantage: All allow unbounded positive activations — the network can represent arbitrarily large values when needed.

The Universal Rule for Activation Gradients

No matter which activation you choose, the backward pass follows the same pattern:

\[\frac{\partial \mathcal{L}_i}{\partial \mathbf z} = \underset{(M \times 1)}{\delta} \odot \underset{(M \times 1)}{\varphi'(\mathbf z)} = \underset{(M \times 1)}{\text{result}}\]

Same shape in, same shape out. Always \(\odot\) (elementwise), never \(\times\) (matrix).

Cheap Gradients!

The computation is always cheap: one elementwise multiply per hidden unit.

Where We’re At

Where We Are in the Chain

We’ve derived backward rules for two module types:

• Linear: bias (\(= \delta\)), weights (\(= \mathbf h \times \delta^T\)), pass-through (\(= \mathbf W \times \delta\))
• Activation: pass-through (\(= \delta \odot \varphi'(\mathbf z)\))

We’ve applied them to the top of our 3-layer network:

\[\mathcal{L}_i \xrightarrow{\delta_0} \text{Linear (output)} \xrightarrow{\delta_1} \text{ReLU 3} \xrightarrow{\delta_2} \quad \cdots\]

We have gradients for \(\bbeta\) and \(\boldsymbol \alpha\), and the upstream gradient \(\delta_2\) is sitting at the entrance to Linear 3 (\(\mathbf W_3, \mathbf b_3\)).

What does \(\delta_2\) actually look like?

The Upstream Gradient Arriving at Layer 3

Unrolling what we’ve computed so far:

\[\delta_0 = \hat{y}_i - y_i^* \quad (K \times 1)\]

\[\delta_1 = \underset{(M_3 \times K)}{\bbeta} \times \underset{(K \times 1)}{\delta_0} \quad (M_3 \times 1)\]

\[\delta_2 = \underset{(M_3 \times 1)}{\delta_1} \odot \underset{(M_3 \times 1)}{\mathbb{I}(\mathbf z_{i3} > 0)} \quad (M_3 \times 1)\]

The Upstream Gradient Arriving at Layer 3

Writing it all out:

\[\delta_2 = \left[\underset{(M_3 \times K)}{\bbeta} \times \underset{(K \times 1)}{(\hat{y}_i - y_i^*)}\right] \odot \underset{(M_3 \times 1)}{\mathbb{I}(\mathbf z_{i3} > 0)}\]

This is the accumulated gradient signal that has flowed down from the loss through the output layer and ReLU 3.

Gradients for the Hidden Layer

From our chain of partials:

\[\frac{\partial \mathcal{L}_i}{\partial \mathbf W_3} = \underbrace{\frac{\partial \mathcal{L}_i}{\partial \btheta_i} \cdot \frac{\partial \btheta_i}{\partial \mathbf h_{i3}} \cdot \frac{\partial \mathbf h_{i3}}{\partial \mathbf z_{i3}}}_{\delta_2 \text{ (already computed!)}} \cdot \frac{\partial \mathbf z_{i3}}{\partial \mathbf W_3}\]

But We Already Know Every Piece!

Each of those partial derivatives is a local gradient we’ve already derived:

Every local gradient is either a \(\times\) by a weight matrix or a \(\odot\) with a binary mask.

We just walk down the chain, applying the same two rules over and over.

Walking Down the Chain

Starting from \(\delta_0 = \hat{y}_i - y_i^* \quad (K \times 1)\):

Output linear module (\(\bbeta, \boldsymbol \alpha\)):

• \(\frac{\partial \mathcal{L}_i}{\partial \boldsymbol \alpha} = \delta_0 \quad (K \times 1)\) ✓
• \(\frac{\partial \mathcal{L}_i}{\partial \bbeta} = \underset{(M_3 \times 1)}{\mathbf h_{i3}} \times \underset{(1 \times K)}{\delta_0^T} \quad (M_3 \times K)\) ✓
• \(\delta_1 = \underset{(M_3 \times K)}{\bbeta} \times \underset{(K \times 1)}{\delta_0} \quad (M_3 \times 1)\)

Walking Down the Chain

Linear 3 (\(\mathbf W_3, \mathbf b_3\)):

• \(\frac{\partial \mathcal{L}_i}{\partial \mathbf b_3} = \delta_2 \quad (M_3 \times 1)\) ✓
• \(\frac{\partial \mathcal{L}_i}{\partial \mathbf W_3} = \underset{(M_2 \times 1)}{\mathbf h_{i2}} \times \underset{(1 \times M_3)}{\delta_2^T} \quad (M_2 \times M_3)\) ✓
• \(\delta_3 = \underset{(M_2 \times M_3)}{\mathbf W_3} \times \underset{(M_3 \times 1)}{\delta_2} \quad (M_2 \times 1)\)

Walking Down the Chain

Linear 2 (\(\mathbf W_2, \mathbf b_2\)):

• \(\frac{\partial \mathcal{L}_i}{\partial \mathbf b_2} = \delta_4 \quad (M_2 \times 1)\) ✓
• \(\frac{\partial \mathcal{L}_i}{\partial \mathbf W_2} = \underset{(M_1 \times 1)}{\mathbf h_{i1}} \times \underset{(1 \times M_2)}{\delta_4^T} \quad (M_1 \times M_2)\) ✓
• \(\delta_5 = \underset{(M_1 \times M_2)}{\mathbf W_2} \times \underset{(M_2 \times 1)}{\delta_4} \quad (M_1 \times 1)\)

Walking Down the Chain

Linear 1 (\(\mathbf W_1, \mathbf b_1\)):

• \(\frac{\partial \mathcal{L}_i}{\partial \mathbf b_1} = \delta_6 \quad (M_1 \times 1)\) ✓
• \(\frac{\partial \mathcal{L}_i}{\partial \mathbf W_1} = \underset{(P \times 1)}{\mathbf x_i} \times \underset{(1 \times M_1)}{\delta_6^T} \quad (P \times M_1)\) ✓

The Dimensions Always Work

Since the dimensionality always matches up perfectly at each layer:

You can stack as many layers as you want — the pattern repeats identically.

This is why nn.Sequential works: each module’s backward output is exactly the right shape for the next module’s backward input.

The Whole Algorithm in One Table

From Chain Rule to Algorithm

What we just derived is the mathematical content of backpropagation.

But there’s an important structural insight beyond “just apply the chain rule.”

The Forward Pass

Given current parameter values, evaluate the chain from bottom to top:

1. \(\mathbf z_{i1} = \mathbf W_1^T \times \mathbf x_i + \mathbf b_1\) — cache \(\mathbf z_{i1}\)
2. \(\mathbf h_{i1} = \varphi(\mathbf z_{i1})\) — cache \(\mathbf h_{i1}\)
3. \(\mathbf z_{i2} = \mathbf W_2^T \times \mathbf h_{i1} + \mathbf b_2\) — cache \(\mathbf z_{i2}\)
4. \(\mathbf h_{i2} = \varphi(\mathbf z_{i2})\) — cache \(\mathbf h_{i2}\)
5. \(\mathbf z_{i3} = \mathbf W_3^T \times \mathbf h_{i2} + \mathbf b_3\) — cache \(\mathbf z_{i3}\)
6. \(\mathbf h_{i3} = \varphi(\mathbf z_{i3})\) — cache \(\mathbf h_{i3}\)
7. \(\btheta_i = \bbeta^T \times \mathbf h_{i3} + \boldsymbol \alpha\)
8. \(\mathcal{L}_i = \text{Loss}(\sigma_{\text{out}}(\btheta_i), y_i)\)

This is just prediction — nothing new. But now we keep all the intermediate values.

The Backward Pass

Starting from the loss, walk backward down the chain:

1. \(\delta_0 = \hat{y}_i - y_i^*\) — loss module backward
2. Store \(\frac{\partial \mathcal{L}_i}{\partial \boldsymbol \alpha}\), \(\frac{\partial \mathcal{L}_i}{\partial \bbeta}\); compute \(\delta_1\) — output linear backward
3. Compute \(\delta_2\) — ReLU 3 backward (uses cached \(\mathbf z_{i3}\))
4. Store \(\frac{\partial \mathcal{L}_i}{\partial \mathbf b_3}\), \(\frac{\partial \mathcal{L}_i}{\partial \mathbf W_3}\); compute \(\delta_3\) — linear 3 backward (uses cached \(\mathbf h_{i2}\))
5. … continue to the bottom

At each module:

• Use the cached forward values and the upstream \(\delta\)
• Produce terminal gradients (if the module has parameters)
• Produce the downstream \(\delta\) (for the module below)

Each Module Has a Forward and a Backward

This is the key design principle:

Forward: takes an input, produces an output, caches what it needs

Backward: takes an upstream \(\delta\), produces parameter gradients (if any) and a downstream \(\delta\)

This Is Automatic Differentiation

Backprop is a special case of reverse-mode automatic differentiation.

The recipe:

1. Define a chain of differentiable modules
2. Run the chain forward (compute and cache)
3. Run the chain backward (chain the \(\delta\)’s)

Any computation that can be expressed as a chain of differentiable modules can be optimized with SGD + autodiff.

This is why the framework extends beyond standard neural networks — anything fully differentiable can be learned.

PyTorch’s Design

Each nn.Module (Linear, ReLU, Sigmoid, BatchNorm, etc.) implements:

• A forward method: computes the output given input, caches intermediates
• A backward method: takes upstream \(\delta\), produces parameter gradients and downstream \(\delta\)

PyTorch Sigmoid — The C Code

PyTorch Sigmoid — The C Code

PyTorch Sigmoid — The C Code

That’s it. The entire sigmoid module is these two functions.

Every module in PyTorch follows this pattern: a forward computation and a local gradient \(\odot\) or \(\times\) the upstream \(\delta\).

The PyTorch Repository

The PyTorch Repository

The PyTorch Repository

Each one implements the same pattern:

1. Forward: compute output, cache what’s needed
2. Backward: local gradient combined with upstream \(\delta\)

C-Optimized and GPU-Ready

Look at the operations each module performs:

• Linear backward: matrix-vector multiply (\(\times\)), outer product (\(\times\))
• ReLU backward: elementwise multiply (\(\odot\))
• Loss backward: vector subtraction

These are simple, uniform, massively parallel operations.

• Written in C and optimized for CUDA
• Run on thousands of GPU cores simultaneously
• The entire mini-batch can be processed in parallel

This is why PyTorch is fast despite networks having millions of parameters.

Why SGD + Backprop Scales

Three properties make deep learning uniquely scalable:

How Does XGBoost Optimize?

Greedy, sequential tree building

• Each tree is added one at a time — cannot parallelize across trees
• Each tree must be fully built before the next one starts

Split search at every node

• Must evaluate candidate splits across all \(P\) features
• Must examine all \(N\) observations to find the best partition
• Finding the best split requires sorting each feature: \(O(N \log N)\)

Complex branching logic

• “If \(x_j < t\), go left; else go right” — inherently sequential, hard to vectorize
• Every split is a different conditional operation

The GPU Analogy

XGBoost operations:

“Read this feature, sort it, evaluate every possible split, pick the best one, branch left or right”

• Complex, conditional, sequential
• Like giving a college student a word problem
• CPUs are good at this: few cores, each very capable

The GPU Analogy

Deep learning’s primitive operations (\(\times\) and \(\odot\)) map perfectly onto GPU architecture.

Tree-based methods do not.

Now the Bad News

We can compute gradients for arbitrarily deep networks — great!

The cost is linear in depth — great!

The Gradient at the Bottom of the Chain

Recall from our summary table, \(\delta\) at the bottom of a \(D\)-layer network passes through:

\[\delta_0 \xrightarrow{\bbeta \;\times} \delta_1 \xrightarrow{\odot \; \mathbb{I}} \delta_2 \xrightarrow{\mathbf W_3 \;\times} \delta_3 \xrightarrow{\odot \; \mathbb{I}} \delta_4 \xrightarrow{\mathbf W_2 \;\times} \delta_5 \xrightarrow{\odot \; \mathbb{I}} \delta_6\]

The Multiplicative Chain Problem

If the terms are on average slightly greater than 1:

\[(1 + \epsilon)^D \to \infty \quad \text{as } D \text{ grows}\]

If the terms are on average slightly less than 1:

\[(1 - \epsilon)^D \to 0 \quad \text{as } D \text{ grows}\]

The deeper the network, the smaller \(\epsilon\) needs to be to cause trouble.

This isn’t a bug in backprop — it’s inherent to chaining many differentiable modules.

Two Failure Modes

Exploding gradients: the products of weight matrices are generally \(> 1\)

• \(\delta\) grows exponentially as it flows down the chain
• Early layers get enormous gradients

Exploding Gradients — What They Look Like

Symptoms:

• Loss suddenly spikes or goes to NaN
• Parameter values diverge to \(\pm \infty\)
• Training was going fine, then suddenly collapses

When it happens:

• Very deep networks
• Poor weight initialization
• Early in training (before parameters settle into reasonable ranges)

Fix 1 — Gradient Clipping

The gradient tells us two things: direction and magnitude.

Key insight: as long as the direction is right, we can cap the magnitude.

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

Fix 2 — Input Normalization

If input features have wildly different scales, the first linear module produces large pre-activations:

\[\mathbf z_{i1} = \mathbf W_1^T \times \mathbf x_i + \mathbf b_1\]

If some \(x_{ij}\) values are in the thousands while others are near zero, \(\mathbf z_{i1}\) can be huge — and that magnitude propagates through the chain.

Fix 3 — The Initialization Problem

Gradient descent requires starting values for all parameters.

Naive approach: \(W_{jk} \sim N(0, 1)\)

Variance Blowup in Both Directions

Forward pass: activation variance scales with fan-in

\[\text{Var}(\mathbf z_l) \propto M_{l-1} \cdot \text{Var}(\mathbf z_{l-1})\]

Each layer amplifies the variance by a factor of \(M_{l-1}\).

Backward pass: gradient variance scales with fan-out

\[\text{Var}(\delta_l) \propto M_{l+1} \cdot \text{Var}(\delta_{l+1})\]

Each step down the chain amplifies the \(\delta\) variance by a factor of \(M_{l+1}\).

With \(N(0,1)\) initialization, both the activations and the gradients blow up exponentially with depth.

Glorot (Xavier) Initialization

Goal: Choose the variance of \(\mathbf W\) so that both forward and backward variance stay at \(\approx 1\).

He Initialization (for ReLU)

ReLU kills roughly half the units (\(z \leq 0\)), cutting the variance in half.

To compensate, double the variance:

\[W_{jk} \sim N\left(0, \frac{2}{M_{\text{in}}}\right)\]

This is the default in PyTorch’s nn.Linear.

The result: activations maintain roughly unit variance across layers, and gradients maintain roughly unit variance as they flow back down the chain.

Proper initialization turns the multiplicative chain from a bomb into a stable signal.

Explosions Are Fixable

We have three simple tools:

• Gradient clipping: cap the step size ✓
• Input normalization: keep initial activations reasonable ✓
• Smart initialization (Glorot/He): prevent variance blowup ✓

Vanishing Gradients Are Different

Look back at the ReLU backward rule:

\[\delta_{\text{below}} = \delta_{\text{above}} \odot \underset{(M \times 1)}{\mathbb{I}(\mathbf z > 0)}\]

Dead ReLUs in the Chain

Imagine a deep network. At each ReLU layer, some fraction of units are dead:

\[\delta_6 = \delta_0 \underset{\times}{\bbeta} \underset{\odot}{[1,0,1,\ldots]} \underset{\times}{\mathbf W_3} \underset{\odot}{[0,1,0,\ldots]} \underset{\times}{\mathbf W_2} \underset{\odot}{[1,0,0,\ldots]}\]

Why Clipping Can’t Help

Vanishing gradients: The signal is gone.

• You can’t “clip up” a zero.
• You can’t amplify a signal that doesn’t exist.

What Would Fix This?

Think about what’s causing the problem:

What Would Fix This?

3. The gradient must pass through every single module

• A chain of \(D\) modules means \(D\) opportunities to shrink \(\delta\)
• What if we gave the gradient a shortcut to skip modules?

Solutions (Scattered In the Next Month When Needed)

Activation functions that never fully kill the gradient:

• Leaky ReLU: \(\max(\alpha z, z)\) — gradient is \(\alpha\) (small but nonzero) when \(z < 0\)
• GELU: smooth, always nonzero gradient