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