Week 13: Nonconvex Deep Learning

The Nonconvex Reality#

Every deep neural network — from a 3-layer MLP to a GPT-scale transformer — defines a nonconvex loss landscape. The loss $J(\theta)$ as a function of all parameters $\theta$ has:

• Many local minima: Different parameter configurations that all achieve low training loss
• Saddle points: Points where the gradient is zero but the Hessian has both positive and negative eigenvalues — neither a minimum nor a maximum
• Plateaus and flat regions: Large areas where the loss changes very slowly
• Sharp ravines: Narrow paths where the loss drops rapidly in one direction but is flat in others

The convex guarantees from Weeks 2–3 (local = global, deterministic convergence rates) no longer hold. Yet deep nets train successfully all the time. How?

Saddles vs Minima#

In high-dimensional random landscapes, saddle points vastly outnumber local minima. This is a statistical fact: for a random Gaussian field in $D$ dimensions, the probability that a critical point is a local minimum decays exponentially with $D$. Most critical points are saddles, with a few negative eigenvalues and many positive ones.

Why does this matter? Gradient-based methods naturally escape strict saddle points. The gradient points away from the saddle along the negative-curvature directions, so GD/SGDStochastic Gradient Descent will eventually slide off — stochastic gradient noise accelerates this escape.

The hierarchy of concern:

1. Saddle points: Common but escapable — gradient methods handle them reasonably well
2. Bad local minima (high loss): Exponentially rare in overparameterised regimes — the real problem is finding them, not avoiding them
3. Good local minima (low loss): Most local minima in overparameterised nets have loss close to the global minimum

This is an empirical observation (backed by growing theory) rather than a theorem: for overparameterised networks, training reliably finds solutions with near-optimal loss.

Overparameterization#

A model is overparameterised when it has more parameters than training examples ($d \gg n$), or more broadly when it has enough capacity to interpolate (fit) the training data perfectly.

In the overparameterised regime:

• The loss landscape becomes more benign — local minima tend to be near-global
• The set of global minimisers becomes a connected manifold rather than isolated points
• GD and SGDStochastic Gradient Descent converge to solutions with particular properties (see implicit bias, Week 10)

The double-descent phenomenon illustrates this: as you increase model capacity past the interpolation threshold, test error first rises (classical overfitting) then falls again. The optimisation problem gets easier in the overparameterised regime — more parameters make the landscape smoother.

Initialization Matters#

The starting point $w_0$ for GD determines which basin of attraction you fall into. Good initialisation:

1. Prevents vanishing/exploding activations at the first forward pass
2. Places parameters in a region with informative gradients (not a flat plateau)
3. Encourages symmetry breaking so different neurons learn different features

Xavier/Glorot initialisation (for tanh/sigmoid activations):

• Sample $W_{ij} \sim \mathcal{U}\left[-\frac{\sqrt{6}}{\sqrt{n_{\text{in}} + n_{\text{out}}}}, \frac{\sqrt{6}}{\sqrt{n_{\text{in}} + n_{\text{out}}}}\right]$
• Keeps the variance of activations and gradients roughly constant across layers

He/Kaiming initialisation (for ReLU activations):

• Sample $W_{ij} \sim \mathcal{N}\left(0, \frac{2}{n_{\text{in}}}\right)$
• Accounts for ReLU zeroing out half the activations (variance scaling by factor 2)

Why scale matters: If initial weights are too small, activations shrink to zero layer by layer (vanishing gradients). If too large, activations explode (exploding gradients). The correct scaling keeps the forward and backward passes numerically balanced.

Normalization as Landscape Shaping#

Normalisation layers reshape the loss landscape to make optimisation easier:

Batch Normalisation (BN): Normalises activations within a mini-batch to zero mean and unit variance. BN:

• Reduces internal covariate shift: the distribution of layer inputs stays stable during training
• Smooths the loss landscape: the Lipschitz constant $L$ of the loss is reduced, allowing larger step sizes
• Acts as a mild regulariser: the mini-batch statistics inject noise

Layer Normalisation (LN): Normalises across features for each example independently (used in transformers). LN works when batch size is small or varies.

BN and LN don't change the global minima of the loss, but they make the path to those minima smoother and faster. With normalisation, training is less sensitive to the learning rate and initialisation — you can use larger $\eta$ and simpler init schemes.

What Transfers from Convex Theory#

The convex toolkit still matters for nonconvex deep learning:

Transfers well:

• Step-size intuition: $\eta = 1/L$ as a starting point (estimate $L$ via power iteration on the Hessian)
• Momentum and its $\sqrt{\kappa}$ acceleration (Week 7) — still works, just without the crisp theory
• Conditioning diagnosis: if some parameters change much faster than others, adaptive methods (Week 8) help
• Gradient clipping (Week 9): still prevents catastrophic divergence
• Early stopping (Week 10): still an effective regulariser

Does not transfer:

• Global optimality guarantees: a converged deep net is not guaranteed to be globally optimal
• Sublinear-to-linear convergence rate theory: convergence curves in deep learning are empirical, not theoretically predicted
• Convex reformulation: the problem is inherently nonconvex — you cannot "make it convex" by a clever change of variables

The pragmatic approach: use convex theory to diagnose and debug training (why is it slow? is it oscillating?), and use nonconvex heuristics (overparameterisation, init, normalisation) to make training possible in the first place.

Browser lab#

Train a tiny 2-layer MLP (numpy, manual backprop) on XOR. Compare convergence with and without He initialisation. Plot loss vs iteration.

import numpy as np
import matplotlib.pyplot as plt

def relu(x): return np.maximum(0, x)
def drelu(x): return (x > 0).astype(float)
def sigmoid(x): return 1 / (1 + np.exp(-x))

X = np.array([[0, 0], [0, 1], [1, 0], [1, 1]]).T
y = np.array([[0, 1, 1, 0]])

def init_weights(d_in, d_hid, d_out, scale_fn):
    W1 = scale_fn(d_in, d_hid)
    b1 = np.zeros((d_hid, 1))
    W2 = scale_fn(d_hid, d_out)
    b2 = np.zeros((d_out, 1))
    return W1, b1, W2, b2

def forward(W1, b1, W2, b2, X):
    z1 = W1 @ X + b1; a1 = relu(z1)
    z2 = W2 @ a1 + b2; a2 = sigmoid(z2)
    return z1, a1, z2, a2

def compute_loss(a2, y):
    return -np.mean(y * np.log(a2 + 1e-12) + (1 - y) * np.log(1 - a2 + 1e-12))

def train(W1, b1, W2, b2, X, y, eta, K):
    losses = []
    for _ in range(K):
        z1, a1, z2, a2 = forward(W1, b1, W2, b2, X)
        losses.append(compute_loss(a2, y))
        dz2 = a2 - y
        dW2 = dz2 @ a1.T / 4; db2 = np.mean(dz2, axis=1, keepdims=True)
        dz1 = (W2.T @ dz2) * drelu(z1)
        dW1 = dz1 @ X.T / 4; db1 = np.mean(dz1, axis=1, keepdims=True)
        W1 -= eta * dW1; b1 -= eta * db1
        W2 -= eta * dW2; b2 -= eta * db2
    return losses

d_in, d_hid, d_out = 2, 4, 1
K, eta = 2000, 0.5

loss_he = train(*init_weights(d_in, d_hid, d_out,
    lambda i, o: np.random.randn(i, o) * np.sqrt(2.0 / i)), X, y, eta, K)

loss_small = train(*init_weights(d_in, d_hid, d_out,
    lambda i, o: np.random.randn(i, o) * 0.01), X, y, eta, K)

plt.figure(figsize=(8, 5))
plt.plot(loss_he, label="He init (√(2/n_in))", linewidth=1.5)
plt.plot(loss_small, label="Small randn × 0.01", linewidth=1.5)
plt.yscale("log"); plt.xlabel("Iteration"); plt.ylabel("Loss")
plt.title("XOR: He Init vs Poor Init")
plt.legend(); plt.grid(True, alpha=0.3); plt.show()