Deep Networks & Optimization

A single hidden layer network is a universal approximator — given enough neurons, it can approximate any continuous function on a compact set (the Universal Approximation Theorem). So why use depth at all?

The answer is efficiency. Depth buys exponential representational power per parameter. Consider computing the parity function on n bits: a shallow network needs exponentially many neurons (Ω(2^n)), while a deep network with O(n) layers needs only O(n) neurons total.

More intuitively, deep networks learn hierarchies of features. In vision: pixels → edges → textures → object parts → objects. Each layer combines the features of the previous layer into more abstract, compositional features. This compositional structure matches the compositional structure of the world, which is why depth provides such dramatic leverage on real data.

Formally, consider functions that can be expressed as compositions: f(x) = f_L(f_{L-1}(...f_1(x))). If each f_i is a 'simple' function (e.g., a linear map followed by a pointwise nonlinearity), then composing them builds complexity combinatorially. A network with L layers and width w can represent O(w^L) distinct linear regions in its input space — exponential in depth.

But depth is not free. More layers means more matrix multiplications in the forward pass and more chain-rule applications in the backward pass. This is where the pathologies begin.

import numpy as np

# Compare shallow vs deep networks for parity (conceptual)
# Parity of n bits: XOR of all bits

# Shallow: needs exponential neurons
# Deep: 2-neuron-per-layer pipeline works

def deep_xor(bits):
    """Conceptual: log2(n) layers, 2 neurons each."""
    x = bits.copy()
    while len(x) > 1:
        # Each layer XORs pairs
        x = [a ^ b for a, b in zip(x[0::2], x[1::2])]
    return x[0]

bits = [1, 0, 1, 1, 0, 1, 0, 0]  # 8 bits
print(deep_xor(bits))  # 0  (even parity)
# A shallow net needs 2^8 = 256 hidden neurons for same task

Training a neural network means computing the gradient of the loss with respect to every parameter via backpropagation, which is just the chain rule applied recursively. For a network with L layers, the gradient of the loss L with respect to parameters in layer 1 is:

∂L/∂W₁ = (∂L/∂aL) · (∂aL/∂a_{L-1}) · ... · (∂a₂/∂a₁) · (∂a₁/∂W₁)

Each factor (∂a_{l+1}/∂a_l) is the Jacobian of one layer's activations with respect to the previous layer's activations. For a layer with weight matrix W and activation function σ, this Jacobian has spectral norm approximately |W| · σ'.

The core problem: if the average singular value of each layer's Jacobian is < 1, the product of L such matrices → 0 exponentially fast as L grows. Gradients in early layers become numerically zero, and those layers stop learning entirely.

For sigmoid activations (σ(z) = 1/(1+e^{-z})), the derivative σ'(z) ≤ 0.25 everywhere, and at saturation (|z| >> 0) it approaches 0. Combined with typical random weight initialization (Gaussian with std ~1), the product of L factors each ≤ 0.25 shrinks to 10^{-7} after just 14 layers.

This was discovered empirically by Hochreiter (1991) and studied theoretically by Bengio et al. (1994). It explains why before ~2012, 'deep' networks (> 4-5 layers) almost never worked in practice.

import numpy as np

def sigmoid(z):
    return 1 / (1 + np.exp(-z))

def sigmoid_grad(z):
    s = sigmoid(z)
    return s * (1 - s)  # max value = 0.25

# Simulate gradient norm through L sigmoid layers
# random weights, random inputs
np.random.seed(42)
L = 20
n = 100  # neurons per layer

gradient_norm = 1.0
for layer in range(L):
    W = np.random.randn(n, n) * 0.1  # small init
    z = np.random.randn(n)
    # Jacobian spectral norm ≈ ||W|| * mean(sigma')
    jacob_norm = np.linalg.norm(W, ord=2) * np.mean(np.abs(sigmoid_grad(z)))
    gradient_norm *= jacob_norm
    if layer % 5 == 4:
        print(f"After layer {layer+1}: gradient norm ≈ {gradient_norm:.2e}")
# After layer 5:  gradient norm ≈ 1.2e-07
# After layer 10: gradient norm ≈ 1.4e-16  ← numerically zero!

The flip side of vanishing gradients: if the average Jacobian spectral norm > 1, the gradient product grows exponentially with depth. This manifests as:

• Loss oscillates wildly or diverges to NaN
• Parameter updates are enormous, overshooting any reasonable minimum
• The gradient norm spike can be sudden — training looks fine for 1000 steps then diverges on step 1001

The standard diagnostic is to plot the gradient norm (∥∇L∥) over training. A healthy run should show a rough decay toward a small value. Exploding gradients show as sudden spikes.

Gradient clipping is the standard cure: if ∥g∥ > threshold τ, replace g with τ · g/∥g∥. This is global (clips the entire gradient vector, not coordinate-wise), which preserves direction while bounding magnitude. Clip value is a hyperparameter, typically 1.0–5.0.

Exploding gradients are especially severe in recurrent networks (RNNs), where the same weight matrix W is applied repeatedly (L times for a sequence of length L), so the product W^L either → 0 or ∞ depending on the spectral radius ρ(W). This is why LSTMs and Transformers largely replaced vanilla RNNs — they route gradients around this product via gating (LSTM) or attention (Transformer).

import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(256, 256), nn.Tanh(),
    nn.Linear(256, 256), nn.Tanh(),
    nn.Linear(256, 10)
)

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

def train_step(x, y):
    optimizer.zero_grad()
    loss = nn.CrossEntropyLoss()(model(x), y)
    loss.backward()
    
    # Gradient clipping — BEFORE optimizer.step()
    grad_norm = nn.utils.clip_grad_norm_(
        model.parameters(),
        max_norm=1.0  # clip to unit norm
    )
    print(f"Gradient norm (pre-clip): {grad_norm:.4f}")
    
    optimizer.step()
    return loss.item()

The key insight: at the start of training, we want signals to flow through the network without vanishing or exploding — not just gradients, but also activations. This is a variance preservation condition.

Forward pass analysis: if layer l has n_in inputs, each weight w ~ N(0, σ²), and activations are independent with unit variance, then the variance of the output z = Σ w_i x_i is n_in · σ². To keep variance of z equal to 1, we need σ² = 1/n_in.

Backward pass analysis: by a symmetric argument, to preserve gradient variance, we need σ² = 1/n_out.

Xavier (Glorot) initialization (2010): compromise for tanh/sigmoid activations: σ² = 2 / (n_in + n_out) [or equivalently, Uniform(-√(6/(n_in+n_out)), +√(6/(n_in+n_out)))]

Xavier assumes the activation function has derivative ≈ 1 near 0 (valid for tanh). It fails for ReLU because ReLU zeroes out half of all activations.

He initialization (2015, Kaiming He et al.): for ReLU activations: σ² = 2 / n_in

The factor of 2 compensates for the half of activations zeroed by ReLU. This is derived rigorously by computing E[ReLU(z)²] = Var(z)/2 when z is zero-mean Gaussian.

Rule of thumb: use He init with ReLU/LeakyReLU; use Xavier init with tanh/sigmoid/linear.

import torch
import torch.nn as nn
import math

# PyTorch defaults to Kaiming Uniform for Linear layers,
# but you can be explicit:

def init_weights(module):
    if isinstance(module, nn.Linear):
        # He initialization for ReLU networks
        nn.init.kaiming_normal_(
            module.weight,
            mode='fan_in',         # use n_in for variance
            nonlinearity='relu'
        )
        nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Conv2d):
        # He init for conv layers
        nn.init.kaiming_normal_(
            module.weight,
            mode='fan_out',        # fan_out for conv
            nonlinearity='relu'
        )

# For tanh networks: Xavier
def init_xavier(module):
    if isinstance(module, nn.Linear):
        nn.init.xavier_normal_(module.weight)
        nn.init.zeros_(module.bias)

model = nn.Sequential(nn.Linear(784, 256), nn.ReLU(),
                      nn.Linear(256, 256), nn.ReLU(),
                      nn.Linear(256, 10))
model.apply(init_weights)

# Verify: check activation variance at init
x = torch.randn(1000, 784)
with torch.no_grad():
    h1 = torch.relu(model[0](x))
    h2 = torch.relu(model[2](h1))
    print(f"Input var: {x.var():.3f}")
    print(f"h1 var:    {h1.var():.3f}")
    print(f"h2 var:    {h2.var():.3f}")

Even with careful initialization, as training progresses, the distribution of each layer's inputs shifts as the parameters of preceding layers change. Ioffe and Szegedy (2015) called this Internal Covariate Shift. Batch Normalization (BatchNorm) addresses it by normalizing the pre-activation at each layer on every mini-batch.

Algorithm for a mini-batch B = {x₁, ..., xₘ}:

1. μ_B = (1/m) Σ xᵢ (batch mean)
2. σ²_B = (1/m) Σ (xᵢ - μ_B)² (batch variance)
3. x̂ᵢ = (xᵢ - μ_B) / √(σ²_B + ε) (normalize; ε ≈ 1e-5 for numerical stability)
4. yᵢ = γ x̂ᵢ + β (scale and shift by learned parameters γ, β)

The learned γ (scale) and β (shift) parameters allow the network to undo the normalization if that's what the loss requires — BatchNorm is not forcing any particular distribution; it is giving the network a better-conditioned starting point.

At inference: we use running statistics (exponential moving averages of μ and σ²) accumulated during training, not per-batch statistics (since we may evaluate on a single example).

Why it works (multiple mechanisms):

• Reduces sensitivity to initialization (helps gradients flow)
• Has a regularization effect (stochastic noise from batch statistics)
• Allows higher learning rates (the landscape becomes more bowl-shaped)
• Reduces training time dramatically — often 14× speedup

Placement: conventionally after the linear/conv layer, before the activation: Linear → BN → ReLU. Some architectures (e.g., ResNets in later variants) apply it before the residual add. LayerNorm (used in Transformers) normalizes across the feature dimension rather than the batch dimension.

import torch
import torch.nn as nn

# Manual BatchNorm implementation to build intuition
class ManualBatchNorm1d(nn.Module):
    def __init__(self, num_features, eps=1e-5, momentum=0.1):
        super().__init__()
        self.eps = eps
        self.momentum = momentum
        # Learned scale and shift
        self.gamma = nn.Parameter(torch.ones(num_features))
        self.beta  = nn.Parameter(torch.zeros(num_features))
        # Running stats for inference
        self.register_buffer('running_mean', torch.zeros(num_features))
        self.register_buffer('running_var',  torch.ones(num_features))

    def forward(self, x):
        if self.training:
            mean = x.mean(dim=0)           # (F,)
            var  = x.var(dim=0, unbiased=False)  # (F,)
            # Update running stats (EMA)
            self.running_mean = (1 - self.momentum) * self.running_mean \
                                + self.momentum * mean.detach()
            self.running_var  = (1 - self.momentum) * self.running_var  \
                                + self.momentum * var.detach()
        else:
            mean = self.running_mean
            var  = self.running_var

        x_hat = (x - mean) / torch.sqrt(var + self.eps)
        return self.gamma * x_hat + self.beta

# Usage in a network
net = nn.Sequential(
    nn.Linear(784, 256),
    ManualBatchNorm1d(256),  # or nn.BatchNorm1d(256)
    nn.ReLU(),
    nn.Linear(256, 10)
)

Dropout (Srivastava et al., 2014) is a form of regularization that prevents co-adaptation of neurons — the phenomenon where neurons collectively learn to memorize training data in a way that fails to generalize.

Algorithm: during each forward pass of training, each neuron (activation) is independently set to zero with probability p (the dropout rate, typically 0.1–0.5). At test time, no dropout is applied, but activations are scaled by (1-p) to maintain the same expected value, OR equivalently, during training activations are scaled by 1/(1-p) (inverted dropout — the standard in practice).

Inverted dropout (what PyTorch uses):

• Train: activate with prob (1-p), scale surviving activations by 1/(1-p)
• Test: use all activations, no scaling needed

Why it works:

1. Ensemble interpretation: each forward pass corresponds to a different 'thinned' subnetwork. Training with dropout is approximately averaging over 2^n different networks (for n neurons). At test time, the full network approximates this ensemble.
2. Feature robustness: a neuron cannot rely on any single co-adapted partner being present — it must learn features that are individually useful.
3. Reduction of complex co-adaptations: similar in spirit to bagging, but within a single network.

Where to apply: typically after fully-connected layers, not (usually) after conv layers. After each hidden layer before the activation, or after the activation — both are common. Don't apply before the input layer.

Modern note: BatchNorm and Dropout interact poorly in some settings — Dropout introduces noise in the running statistics of BN. In Transformers, Dropout is used extensively; BN is replaced by LayerNorm.

import torch
import torch.nn as nn

# Manual inverted dropout
class InvertedDropout(nn.Module):
    def __init__(self, p=0.5):
        super().__init__()
        self.p = p

    def forward(self, x):
        if not self.training or self.p == 0:
            return x
        # Bernoulli mask: keep each neuron w.p. (1-p)
        keep_prob = 1 - self.p
        mask = torch.bernoulli(torch.full_like(x, keep_prob))
        # Inverted: scale UP by 1/keep_prob during training
        # so no change needed at test time
        return x * mask / keep_prob

# Use in a classifier
classifier = nn.Sequential(
    nn.Linear(784, 512), nn.ReLU(),
    nn.Dropout(p=0.5),         # or InvertedDropout(0.5)
    nn.Linear(512, 256), nn.ReLU(),
    nn.Dropout(p=0.3),
    nn.Linear(256, 10)
)

# IMPORTANT: must switch modes
classifier.train()  # dropout active
classifier.eval()   # dropout off, uses full network

Stochastic Gradient Descent (SGD) is the foundation of all deep learning optimization. Instead of computing the exact gradient over the entire dataset (which is O(n) per step), SGD computes a noisy gradient estimate using a mini-batch of size B:

w ← w - η · ∇L(w; mini-batch)

where η is the learning rate. The noise is not a bug — it's a feature. It acts as implicit regularization and helps escape sharp (poorly generalizing) minima.

Problems with vanilla SGD:

1. Oscillation: in elongated ravines (common in neural network loss landscapes), SGD oscillates across the narrow dimension while making slow progress along the long dimension.
2. Local minima / saddle points: gradient is nearly zero; SGD stalls.

Momentum (Polyak, 1964) introduces a velocity term: v ← μ·v - η·∇L(w)
w ← w + v

where μ ∈ [0.9, 0.99] is the momentum coefficient. Momentum accumulates gradients like a ball rolling down a hill — it accelerates in directions of consistent gradient and dampens oscillations across the narrow ravine dimension.

Intuition: if gradients consistently point the same direction (like in the long valley direction), velocity builds up. If gradients alternate sign (oscillating across the narrow direction), they cancel out in the velocity.

Nesterov Accelerated Gradient (NAG): a smarter version that computes the gradient at the lookahead position (w + μ·v) rather than the current position. This gives better theoretical convergence and is often slightly better in practice.

Typical hyperparameters: lr = 0.01–0.1 for SGD+momentum in vision; momentum = 0.9.

import numpy as np

def sgd_with_momentum(grad_fn, w_init, lr=0.01, mu=0.9,
                      n_steps=1000):
    """
    grad_fn: function that returns noisy gradient at w
    w_init:  initial parameters
    """
    w = w_init.copy()
    v = np.zeros_like(w)   # velocity, initialized to zero

    trajectory = [w.copy()]
    for t in range(n_steps):
        g = grad_fn(w)           # stochastic gradient
        v = mu * v - lr * g      # accumulate velocity
        w = w + v                # update params
        trajectory.append(w.copy())

    return w, np.array(trajectory)

# Example: minimize f(x,y) = x^2 + 10*y^2 (elongated bowl)
def grad(w, noise=0.01):
    return np.array([2*w[0], 20*w[1]]) + np.random.randn(2)*noise

w0 = np.array([3.0, 3.0])
w_sgd, traj_sgd = sgd_with_momentum(
    lambda w: grad(w), w0, lr=0.01, mu=0.0  # SGD: mu=0
)
w_mom, traj_mom = sgd_with_momentum(
    lambda w: grad(w), w0, lr=0.01, mu=0.9  # Momentum
)
print(f"SGD final:      {w_sgd}")    # slow convergence
print(f"Momentum final: {w_mom}")   # faster

Adam (Kingma and Ba, 2015) is the default optimizer for most deep learning research. It combines two ideas:

1. Momentum (first moment): exponential moving average of the gradient
2. Adaptive learning rates (second moment): divide by EMA of the squared gradient, so that parameters with historically large gradients get smaller effective learning rates

Algorithm (for parameter θ, gradient g_t at step t):

m_t = β₁ · m_{t-1} + (1-β₁) · g_t       # 1st moment (momentum)
v_t = β₂ · v_{t-1} + (1-β₂) · g_t²      # 2nd moment (RMS)
m̂_t = m_t / (1 - β₁ᵗ)                   # bias correction
v̂_t = v_t / (1 - β₂ᵗ)                   # bias correction
θ ← θ - η · m̂_t / (√v̂_t + ε)

Default hyperparameters: β₁=0.9, β₂=0.999, ε=1e-8, η=1e-3.

Bias correction: at initialization, m₀=v₀=0 (biased toward zero). The correction m̂_t = m_t/(1-β₁ᵗ) inflates the estimate in early steps, counteracting the cold start. By step ~1000, (1-0.9^1000) ≈ 1, so the correction becomes negligible.

Why v_t matters: consider two parameters. Parameter A has consistent large gradients (high v). Parameter B has sparse, small gradients (low v). Adam gives A a smaller effective learning rate (learning rate shrinks by √v), and B a larger effective learning rate. This is especially powerful for sparse gradients (e.g., embeddings, where most tokens aren't updated each step).

Limitations of Adam:

• Does not always generalize as well as SGD+momentum on image classification (Wilson et al., 2017)
• Second-moment estimate v can dominate, effectively causing very small updates late in training (a form of 'gradient signal death')
• AdamW (Loshchilov & Hutter, 2019) fixes the weight decay implementation by decoupling it from the gradient update — use AdamW, not Adam, whenever weight decay is needed

import numpy as np

def adam(grad_fn, theta_init, lr=1e-3,
         beta1=0.9, beta2=0.999, eps=1e-8,
         n_steps=1000):
    theta = theta_init.copy()
    m = np.zeros_like(theta)   # 1st moment
    v = np.zeros_like(theta)   # 2nd moment

    for t in range(1, n_steps + 1):
        g = grad_fn(theta)
        # Update biased moments
        m = beta1 * m + (1 - beta1) * g
        v = beta2 * v + (1 - beta2) * (g ** 2)
        # Bias-corrected moments
        m_hat = m / (1 - beta1 ** t)
        v_hat = v / (1 - beta2 ** t)
        # Parameter update
        theta -= lr * m_hat / (np.sqrt(v_hat) + eps)

    return theta

# In PyTorch:
import torch.optim as optim
# Standard Adam (weight decay baked into gradient — suboptimal)
# opt = optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)

# AdamW — correct weight decay (decoupled)
# opt = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)