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

Full network: Many parameters, most redundant. Trains well.

Winning ticket: The sparse sub-network at the original initialization that performs equally well. Found by training, pruning, rewinding.

Small network from scratch: Same number of parameters, but random initialization. Does not perform as well — the initialization matters.

Why This Matters for Generalization

The lottery ticket hypothesis reframes what overparameterization does:

Old view: Large networks have more capacity → they can memorize → they overfit.

New view: Large networks contain many possible sub-networks. SGD explores this space and finds sub-networks that are small, smooth, and generalizable.

The Practical Takeaway

Overparameterization is not the enemy. Combined with SGD’s implicit bias, it’s a feature.

But SGD’s implicit regularization can be slow — the second descent takes many epochs.

Explicit Regularization

SGD’s implicit bias toward simple solutions is powerful — but slow.

The second descent can take hundreds of epochs. In practice, we want to generalize quickly.

Weight Decay (L2 Regularization)

Add a penalty on the size of all weight matrices:

\[\mathcal{L}_{\text{reg}} = \underbrace{\frac{1}{N}\sum_{i=1}^{N} \ell(f(\mathbf x_i; \btheta), y_i)}_{\text{training loss}} + \underbrace{\lambda \sum_{d=1}^{D} \|\mathbf W_d\|_F^2}_{\text{L2 penalty}}\]

This is Ridge regression applied to every layer.

The penalty \(\lambda \|\mathbf W\|_F^2 = \lambda \sum_{j,k} W_{jk}^2\) pulls every weight toward zero.

Larger \(\lambda\) = stronger pull = smaller weights = smoother function.

What Weight Decay Does to the Loss Surface

Without weight decay (red): parameter norms grow, test loss rises. With weight decay: norms are controlled, test loss stays lower.

Why Weight Decay Works Here

The overparameterized network has infinitely many interpolating solutions.

Weight decay adds an explicit prior: among all solutions that fit the data, pick the one with the smallest weights.

\[\min_{\btheta} \; \mathcal{L}(\btheta) \quad \text{subject to} \quad \|\btheta\| \text{ is small}\]

Early Stopping — The Complexity Path

SGD doesn’t jump to a solution — it walks there along a path through parameter space.

Early Stopping as Regularization

Each SGD step increases the effective complexity of the model:

• At epoch 0: parameters are near the (small) He initialization. The function is nearly linear — low complexity, high bias.
• At epoch \(t\): parameters have moved further from initialization. The function is more complex — lower bias, but potentially higher variance.

Batch Size as a Regularization Knob

Recall: SGD uses a mini-batch of \(B\) observations to estimate the gradient.

\[\hat{\mathbf g}(\btheta) = \frac{1}{B} \sum_{i \in \text{batch}} \nabla_{\btheta} \ell_i\]

Batch Size and Generalization

Small batches find flat minima that tolerate perturbation — these generalize. Large batches converge to sharp minima that are fragile.

Dropout

A different kind of regularization — not about the size of weights, but about how information is distributed across units.

Dropout in Action

Every mini-batch sees a different random half of the hidden units. No unit can afford to be the sole carrier of important information — the network is forced to distribute its representations.

Dropout as an Ensemble

Training: Each forward pass uses a random ~50% of units — a different “sub-network” each time. With \(M\) hidden units, there are \(2^M\) possible sub-networks.

Inference: Use all units, multiply weights by \((1-p)\). This approximates the average prediction over all sub-networks — like a random forest for neural networks. Pytorch handles this for you!

Dropout vs. Dead ReLUs

An important contrast:

Dead ReLU: unit is dead for all observations (or all in a batch). Deterministic. Permanent. Gradient is exactly zero. Caused by the data and weights.

Dropout: unit is dropped for this mini-batch only. Stochastic. Temporary. Same unit is active in the next batch. Caused by random sampling.

Dead ReLUs are a problem — the network loses capacity permanently.

Dropout is a tool — it forces redundancy without permanently destroying anything.

Together with BN: BN prevents dead units by keeping activations healthy. Dropout prevents over-reliance on specific units by randomly killing them. Both encourage distributed representations. They’re complementary.

Dropout and the Lottery Ticket

Recall: the lottery ticket hypothesis says large networks contain small sub-networks that could achieve the same accuracy alone.

Dropout is effectively searching for winning sub-networks during training:

• Each forward pass trains a random ~50% sub-network
• Over many mini-batches, the network learns which units are most consistently useful
• Units that carry important information survive pruning because the network routes around dropped units

The full network provides the search space. Dropout narrows the search toward the sub-networks that are most robust — exactly the kind of solution SGD’s implicit bias prefers.

Putting It All Together

Every technique we’ve covered addresses a specific failure mode:

The Modern Training Recipe

# Architecture
model = nn.Sequential(
    nn.Linear(P, M, bias=False),       # He initialization (PyTorch default)
    nn.BatchNorm1d(M),                 # keep activations healthy
    nn.ReLU(),                         # sparse, efficient
    nn.Dropout(0.3),                   # force distributed representations
    # ... repeat D times ...
    nn.Linear(M, K),                   # output layer
)

# Optimizer
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)

# Training loop
for epoch in range(max_epochs):
    train_one_epoch(model, optimizer, train_loader)
    val_loss = evaluate(model, val_loader)
    if val_loss < best_val_loss:        # early stopping
        best_val_loss = val_loss
        save_checkpoint(model)
    elif patience_exceeded:
        break

# Gradient clipping (inside train_one_epoch)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

Every line maps to a concept from this lecture.

Next Lecture — Convolutional Neural Networks

We got ~55% on CIFAR-10 by flattening 32×32 color images into 3,072-dimensional vectors.

We threw away the fact that nearby pixels are related. A pixel in the top-left corner is treated identically to a pixel in the bottom-right. The network has to learn spatial structure from scratch — with a fully connected weight matrix that has \(3{,}072 \times M\) parameters in the first layer alone.

Next Lecture — Convolutional Neural Networks

Convolutional neural networks encode spatial structure directly into the architecture:

• Local connectivity: each unit sees only a small patch of the input
• Weight sharing: the same filter is applied everywhere
• Translation invariance: a cat in the corner is the same as a cat in the center

Combined with residual skip connections — the gradient highway we previewed today — CNNs take us past 90% on CIFAR-10 and power modern computer vision.