DATASCI 447 Lecture 10: Training Deep Models That Work and Generalize

Administrative Stuff

Gradient Problems and Generalization

\[ \newcommand\hbb{{\hat{\boldsymbol \beta}}} \newcommand\bb{{\boldsymbol \beta}} \newcommand\expn{{\frac{1}{N} \sum \limits_{i = 1}^N}} \newcommand\sumk{\sum \limits_{k = 1}^K} \newcommand\argminb{\underset{\bb}{\text{argmin }}} \newcommand\argmaxb{\underset{\bb}{\text{argmax }}} \newcommand\gtheta{\mathbf g(\boldsymbol \theta)} \newcommand\htheta{\mathbf H(\boldsymbol \theta)} \]

Last time: we derived how deep networks learn.

Today: why training deep networks is hard — and why they generalize despite being massively overparameterized.

Where We Left Off

Last lecture we derived backpropagation — the algorithm for computing gradients in deep networks.

Why Backprop Is So Powerful

This algorithm has three properties that make deep learning possible:

1. Computational efficiency

One forward pass + one backward pass — regardless of how many parameters.

Cost is \(\approx 2\times\) a single forward pass, linear in depth. Compare to finite differences: one forward pass per parameter.

PyTorch Is an Autodiff Machine

These three properties are exactly what PyTorch implements:

model = nn.Sequential(
    nn.Linear(P, M1),     # forward: z = W^T x + b
    nn.ReLU(),             # forward: h = max(0, z)
    nn.Linear(M1, M2),    # forward: z = W^T h + b
    nn.ReLU(),             # ...
    nn.Linear(M2, K),
)

Two Operations, One Algorithm

Every backward step in a fully-connected network uses one of two operations:

These are simple, uniform, massively parallel — perfect for GPUs.

This is why deep networks scale to billions of parameters while tree-based methods hit walls with high-dimensional data.

The Price of Depth

The \(\delta\) arriving at the bottom of a \(D\)-layer chain passes through every module:

\[\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

When you multiply many terms together:

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}\]

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, destructive 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 with many multiplicative steps
• Poor weight initialization (\(\mathbf W\) entries too large)
• Unlucky regions of the loss surface — steep cliffs

The geometry: the loss surface has steep cliffs. One SGD step lands on a cliff, the gradient is enormous, and the update shoots parameters into an entirely different region.

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.

\[g'(\theta) = \min\left(1, \frac{c}{\|g(\theta)\|}\right) \cdot g(\theta)\]

• If \(\|g\| \leq c\): do nothing — the gradient is fine
• If \(\|g\| > c\): scale it down to have norm \(c\), preserving direction

\(c \approx 1\)–\(5\) is typical. Simple, widely used, especially for recurrent models.

One line of PyTorch:

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

Gradient Clipping Is Insurance, Not Strategy

Clipping prevents catastrophic single steps, but it doesn’t address the root cause.

If you find yourself relying on clipping to keep training stable, something deeper is wrong:

• Learning rate too high
• Initialization too large
• Architecture too deep without normalization

Think of clipping as a safety net — you want it there, but you shouldn’t be hitting it often.

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 up 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}\).

Variance Blowup in Both Directions

Glorot (Xavier) Initialization

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

Forward constraint: \(\text{Var}(W) \cdot M_{\text{in}} = 1 \implies \text{Var}(W) = \frac{1}{M_{\text{in}}}\)

Backward constraint: \(\text{Var}(W) \cdot M_{\text{out}} = 1 \implies \text{Var}(W) = \frac{1}{M_{\text{out}}}\)

Glorot (Xavier) Initialization

He Initialization (for ReLU)

ReLU kills roughly half the units (\(z \leq 0\)), cutting the effective 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.

He Initialization (for ReLU)

Adam Helps Too

Adam’s adaptive per-parameter learning rates provide additional protection:

\[\theta_t = \theta_{t-1} - \eta \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}\]

The denominator \(\sqrt{\hat{v}_t}\) tracks the running variance of each parameter’s gradient.

Adam Helps Too

If a parameter’s gradient is consistently large (explosive), \(\hat{v}_t\) is large, and the effective step size shrinks.

If a parameter’s gradient is consistently small, \(\hat{v}_t\) is small, and the effective step size grows.

Exploding Gradients

These all address the same underlying issue: scale.

The \(\delta\)’s are too big? Make them smaller. The activations are too big? Normalize them. The initial weights amplify too much? Shrink them.

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)}\]

Every time a unit is dead (\(z_m \leq 0\)), its gradient is not small — it’s exactly zero.

\[\delta_{\text{above},m} \odot 0 = 0\]

This isn’t a scale problem. The information is genuinely destroyed.

No amount of rescaling can recover a signal that has been multiplied by zero.

Visualizing the Gate

Active units: gate is open — gradient passes through. Dead units: gate is shut — gradient is zeroed.

Dead ReLUs Accumulate Across Layers

Each ReLU \(\odot\) zeros out some dimensions of \(\delta\). Once a dimension is zeroed, it stays zero through all subsequent layers. After enough layers, most of the gradient signal is gone.

Dead ReLUs in the Chain — The Math

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]}\]

Each \(\odot\) with a ReLU mask zeros out some dimensions of \(\delta\).

After enough layers, the \(\delta\) signal can be completely extinguished before reaching the bottom of the chain.

Early layers receive no gradient → their weights don’t update → they don’t learn.

The network “forgets” about its early layers.

Permanent Death

If a unit is dead for every observation in the training set:

• \(\mathbb{I}(z_m > 0) = 0\) for all \(i = 1, \ldots, N\)
• Gradient is zero for all \(i\)
• Weight never updates
• Unit stays dead forever

You can’t “clip up” a vanishing gradient — there’s no signal to amplify.

Turning down volume is easy. Amplifying silence is impossible.

This requires fundamentally different solutions.

Activation Function Fixes

The Alternatives

Leaky ReLU: \(\varphi(z) = \max(\alpha z, z)\) where \(\alpha \approx 0.01\)–\(0.1\)

• Gradient is \(\alpha\) (small but nonzero) when \(z < 0\)
• No unit is ever fully dead

ELU: \(\varphi(z) = z\) if \(z > 0\), \(\alpha(e^z - 1)\) if \(z \leq 0\)

• Smooth negative side, pushes mean activation toward zero
• Slightly more expensive (requires \(\exp\))

GELU: \(\varphi(z) = z \cdot \Phi(z)\) where \(\Phi\) is the standard normal CDF

• Smooth everywhere, never fully kills the gradient
• Used in transformers (BERT, GPT)

The Sparsity Tradeoff

ReLU’s zeros are a feature, not just a bug.

Each observation activates a sparse subset of hidden units. The zeros mean that:

• Each hidden unit corresponds to a coherent subfunction of the input space
• The representation is interpretable — you can ask “which units fired?”
• Less complex models — analogous to how LASSO zeros out unimportant coefficients

The Sparsity Tradeoff

Leaky ReLU and GELU sacrifice this sparsity for gradient flow.

Everything is a little bit active for every input — a muddier representation.

Three Sources, Three Fixes

The Drift Problem

He initialization puts pre-activations near zero at step 0. But as weights update, pre-activations drift — systematically shifting negative, killing more and more units.

The Drift Problem

Changing variance is also an issue since this causes the gradients for individual observations to become bigger or smaller (most of the time the smaller ones are the biggest problem since the gradient gets shoved to zero!)

Batch Normalization

The idea: If the problem is that pre-activations drift away from zero during training, force them back at every step.

Batch Normalization

Given a mini-batch of \(B\) observations, the pre-activations at layer \(d\) are a \(B \times M_d\) matrix.

For each hidden unit \(m\) (each column):

\[\mu_m = \frac{1}{B} \sum_{i \in \text{batch}} z_{i,m} \qquad \sigma^2_m = \frac{1}{B} \sum_{i \in \text{batch}} (z_{i,m} - \mu_m)^2\]

\[\hat{z}_{i,m} = \frac{z_{i,m} - \mu_m}{\sqrt{\sigma^2_m + \epsilon}}\]

Now \(\hat{z}_{i,m}\) has mean 0 and variance 1 across the batch.

What BN Does to the Pre-Activations

BN centers the pre-activations at zero. With a symmetric distribution around zero, roughly half the units are active — plenty of gradient signal flows through.

The Learnable Scale and Shift

Forcing all pre-activations to mean 0 and variance 1 might be too restrictive.

Maybe some units should have a different mean or spread!

The Learnable Scale and Shift

If the optimal pre-activation distribution is \(N(0, 1)\), the network can keep it there.

If the optimal distribution is \(N(2, 0.5)\), the network can learn \(\gamma_m = 0.5\), \(\beta_{\text{bn},m} = 2\).

Batch Normalization

Any intuition as to why this is okay to do in the hidden layers?

Why BN Helps Gradient Flow

Without BN: dead units pile up in early layers, gradient norms decay exponentially.

With BN: dead unit fraction is stable (~50%) at every layer, gradients maintain consistent magnitude.

Where BN Goes in the Chain

The modern default ordering:

\[\mathbf x \xrightarrow{\text{Linear}} \mathbf z \xrightarrow{\text{BN}} \tilde{\mathbf z} \xrightarrow{\text{ReLU}} \mathbf h\]

nn.Sequential(
    nn.Linear(M_in, M_out, bias=False),  # no bias — BN handles it
    nn.BatchNorm1d(M_out),
    nn.ReLU(),
)

Training vs. Inference

During training: \(\mu_B\) and \(\sigma^2_B\) are computed from the current mini-batch.

• Different batches → different statistics → slight noise in the normalized values
• This noise acts as a mild regularizer (more on this soon)

During inference: We don’t have a mini-batch — we might predict on a single observation.

• PyTorch keeps a running average of \(\mu\) and \(\sigma^2\) across all training batches
• At inference time, use these fixed running averages

model.train()   # BN uses batch statistics (training)
model.eval()    # BN uses running averages (inference)

This is one reason model.train() and model.eval() matter in PyTorch — they change how BN behaves.

BN Does Not Solve Overfitting

BN fixes gradient flow — the network can train. But training loss → 0 while test loss rises. The network memorizes the training data.

• And batch norm allows the network to quickly approach the 0 loss solution!

From Optimization to Generalization

BN + He initialization = we can train deep networks.

From Optimization to Generalization

The remaining question — and it turns out to be the deeper question — is:

How do we get deep networks to generalize?

The Generalization Puzzle

A 20-layer, 256-wide network on CIFAR-10 has ~600,000 parameters and only 50,000 training observations.

Classical statistics says: catastrophic overfitting. More parameters than observations means a system with infinitely many solutions — and no reason to expect that any of them generalizes.

The Classical View — Bias-Variance Tradeoff

Think back to linear regression.

If we fit a degree-1 polynomial to nonlinear data: high bias, low variance. Underfitting.

If we fit a degree-\(N\) polynomial to \(N\) points: zero bias, high variance. Overfitting — the curve passes through every point but oscillates wildly.

This predicts a U-shaped test error curve: too simple → bad, just right → good, too complex → bad again.

For decades, this was the guiding principle of model selection.

The Overparameterization Problem

Neural networks violate this story completely.

A 20-layer MLP with 256 units per layer has ~600K parameters for 50K training observations. That’s a 12:1 parameter-to-observation ratio.

In the linear regression analogy, this is like fitting a degree-600,000 polynomial to 50,000 points.

SGD Is Doing the Heavy Lifting

A large part of the generalization story comes from SGD itself — completely separate from the neural network architecture.

SGD and Implicit Regularization →

The key takeaway: SGD preferentially finds minimum-norm interpolating solutions — the simplest function (in a specific sense) that fits the training data perfectly.

SGD Doesn’t Find Arbitrary Solutions

Left: Too few parameters — can’t capture the pattern. Middle: Exactly \(N\) parameters — one solution, and it’s wiggly. Right: Many more parameters than needed — infinitely many interpolating solutions, and SGD finds a smooth one.

Double Descent

This explains the surprising test error curve:

Why More Parameters Help

More parameters = more ways to interpolate = more chances that one of those ways is simple.

Overparameterization provides searchability, not just capacity.

The Lottery Ticket Hypothesis

This leads to a striking conjecture:

Frankle & Carlin (2019): A large, randomly-initialized network contains small sub-networks that — if trained in isolation from the same initialization — could reach the same test accuracy as the full network.

Finding the Winning Ticket