Neural networks

What it is and when you reach for it

A multilayer perceptron is a stack of affine maps separated by elementwise nonlinearities: multiply by a weight matrix, add a bias, bend the result, repeat. Remove the nonlinearities and the whole stack collapses into a single matrix, no more expressive than linear regression; keep them and the composition can approximate any continuous function on a compact set to arbitrary accuracy, which is the universal approximation theorem. Its neighbors define its role. A linear model reaches for it when the decision boundary is not linear and you do not want to hand-craft the features that would make it so. A convolutional network is an MLP with weight sharing and locality baked in for grid data; a transformer alternates attention with MLP blocks. The plain MLP is what you use when the input has no exploitable structure, or as the trainable head bolted onto something that already extracted structure for you. Backpropagation is not specific to MLPs, but the one-hidden-layer MLP is the smallest model where the full mechanism appears, which is why deriving it once by hand pays for itself across every architecture you will ever train.

The math

The forward pass

Take a batch X of shape (N, d_in), rows as samples, and one hidden layer of width h feeding d_out classes. The forward pass is four lines:

Z1 = X W1 + b1        (N, h)       pre-activations
A1 = φ(Z1)            (N, h)       hidden activations, φ = ReLU here
Z2 = A1 W2 + b2       (N, d_out)   logits
L  = CE(softmax(Z2), y)            scalar loss, averaged over the batch

W₁ is (d_in, h), W₂ is (h, d_out), and the biases broadcast across rows. The loss is softmax cross-entropy; the derivation of its clean gradient lives on the softmax page, and we will use the result here: for the averaged loss, the gradient at the logits is (softmax(Z₂) − onehot(y)) / N. That single simplification is the reason the rest of the backward pass stays tidy.

The chain rule through one hidden layer, written out

Backpropagation is nothing but the chain rule applied in reverse topological order, with one bookkeeping convention: for every intermediate tensor T, define dT = ∂L/∂T with exactly the same shape as T. Shapes then become a proof system. Each rule below can be checked by asking: what is the only way to combine the tensors on hand so the output has the shape it must have?

Start at the logits. From the softmax cross-entropy result,

dZ2 = (P − Y) / N                         (N, d_out),  P = softmax(Z2), Y = onehot(y)

Now walk backward through Z₂ = A₁W₂ + b₂. The loss depends on W₂ only through Z₂, and each entry Z₂[n,k] = Σ_j A₁[n,j]·W₂[j,k] + b₂[k]. Differentiating and summing over every path from W₂[j,k] to the loss (that is, over every sample n) gives dW₂[j,k] = Σ_n A₁[n,j]·dZ₂[n,k], which is exactly a matrix product:

dW2 = A1ᵀ dZ2         (h, d_out)    same shape as W2
db2 = Σ_rows dZ2      (d_out,)      bias broadcast forward ⇒ sum backward
dA1 = dZ2 W2ᵀ         (N, h)        gradient flowing to the layer below

The three lines are the whole pattern for a linear layer: the weight gradient is (input)ᵀ(output gradient), the bias gradient sums whatever the broadcast duplicated, and the input gradient is (output gradient)(weight)ᵀ. Next comes the nonlinearity A₁ = φ(Z₁). Because φ acts elementwise, its Jacobian is diagonal and the chain rule reduces to an elementwise product with φ′ evaluated at the pre-activation. For ReLU, φ′(z) is 1 where z > 0 and 0 elsewhere, so the gradient is simply gated:

dZ1 = dA1 ⊙ φ′(Z1) = dA1 ⊙ 1[Z1 > 0]     (N, h)

And the first layer repeats the linear pattern with X in place of A₁:

dW1 = Xᵀ dZ1          (d_in, h)
db1 = Σ_rows dZ1      (h,)

That is the entire algorithm. The forward pass costs two matrix multiplies; the backward costs four (two per layer), so the classic rule of thumb that a training step is roughly three times a forward pass falls straight out of the derivation. Notice also what had to be remembered from the forward pass: X, Z₁, A₁, and the softmax output. Those saved tensors are exactly what frameworks call the autograd tape, and their memory footprint, not compute, is usually what limits batch size.

A worked numeric example

One sample makes the machinery concrete. Let x = [1, 2], target class 0, with W₁ = [[1, −1], [0, 1]], W₂ = [[1, 0], [0, −1]], and zero biases. Forward: z₁ = xW₁ = [1·1+2·0, 1·(−1)+2·1] = [1, 1]; ReLU leaves it alone, a₁ = [1, 1]; z₂ = a₁W₂ = [1, −1]. Softmax of [1, −1] is [0.8808, 0.1192], so the loss is −ln 0.8808 ≈ 0.1269. Backward: dz₂ = p − onehot = [−0.1192, 0.1192]; dW₂ = a₁ᵀdz₂ = [[−0.1192, 0.1192], [−0.1192, 0.1192]]; da₁ = dz₂W₂ᵀ = [−0.1192, −0.1192]; both entries of z₁ are positive so the ReLU gate passes everything, dz₁ = da₁; and dW₁ = xᵀdz₁ = [[−0.1192, −0.1192], [−0.2384, −0.2384]]. Every sign makes sense: the network already prefers the correct class, so the gradients are small, and they push z₂[0] up and z₂[1] down.

Initialization: why scale is everything

Before the first gradient step, the weights are random, and the variance of that randomness decides whether the network is trainable. Two failure modes bracket the problem. Initialize everything to zero (or any constant) and every hidden unit computes the same value forward and receives the same gradient backward, so the units never differentiate: symmetry is never broken and the layer behaves as if it had width one. Initialize with variance too large or too small and the pre-activation variance grows or shrinks geometrically with depth, saturating activations or drowning the signal, and the gradients explode or vanish in the same geometric fashion on the way back.

The fix is to choose the variance so that each layer roughly preserves the variance of what passes through it. For z = Σ_i w_ix_i with independent zero-mean terms, Var(z) = fan_in · Var(w) · Var(x), so keeping Var(z) ≈ Var(x) requires Var(w) = 1/fan_in. Xavier (Glorot) initialization compromises between the forward pass and the backward pass, which cares about fan_out, using Var(w) = 2/(fan_in + fan_out); it was derived for symmetric activations like tanh. ReLU zeroes half of every zero-mean input, cutting the variance in half at each layer, and He initialization compensates with exactly a factor of two: Var(w) = 2/fan_in. The rule of thumb is Xavier for tanh and sigmoid, He for ReLU and its relatives, and the difference is not cosmetic: at depth 20, a factor-of-two variance error per layer compounds to a factor of a million.

Choosing the activation

ReLU, max(0, z), is the default: cheap, non-saturating on the positive side, and its gradient is exactly 0 or 1 so it neither amplifies nor shrinks what flows through active units. Its known pathology is the dead unit: a unit whose pre-activation goes negative for every input receives zero gradient forever, which is why a too-hot learning rate can permanently kill a fraction of a layer. LeakyReLU patches that with a small negative slope. GELU, z·Φ(z) with Φ the Gaussian CDF, is a smooth relaxation that has become the default inside transformers (BERT, GPT); SiLU/Swish, z·σ(z), is its close cousin and shows up in EfficientNet and most modern LLM feed-forward blocks in a gated form (SwiGLU). Tanh survives mostly in gates and small control networks. In practice the choice moves final accuracy by fractions of a point; the initialization matched to the choice moves whether training works.

Implementation, twice

Manual forward and backward, no autograd

The first pair of implementations uses each framework purely as an array library: every gradient below is the hand-derived formula from the math section, transcribed. The invariant worth internalizing is that each dW has the same shape as its W; most transposition bugs die at that check before any numerics run.

import torch
import torch.nn.functional as F

def init_params(d_in, d_h, d_out, seed=0):
    g = torch.Generator().manual_seed(seed)
    # He for the ReLU layer, Xavier-style 1/fan_in for the linear output
    W1 = torch.randn(d_in, d_h, generator=g) * (2.0 / d_in) ** 0.5
    W2 = torch.randn(d_h, d_out, generator=g) * (1.0 / d_h) ** 0.5
    return W1, torch.zeros(d_h), W2, torch.zeros(d_out)

def forward(params, X):
    W1, b1, W2, b2 = params
    Z1 = X @ W1 + b1              # (N, h)
    A1 = Z1.clamp(min=0)          # ReLU
    Z2 = A1 @ W2 + b2             # (N, d_out) logits
    return Z1, A1, Z2             # cache = the autograd tape, by hand

def loss(Z2, y):
    return F.cross_entropy(Z2, y)  # fused log-softmax + NLL

def backward(params, X, Z1, A1, Z2, y):
    """Every line is the derived formula; every dW matches its W's shape."""
    W1, b1, W2, b2 = params
    N = X.shape[0]
    P = torch.softmax(Z2, dim=1)
    dZ2 = (P - F.one_hot(y, Z2.shape[1]).float()) / N   # (N, d_out)
    dW2 = A1.T @ dZ2                                    # (h, d_out)
    db2 = dZ2.sum(0)
    dA1 = dZ2 @ W2.T                                    # (N, h)
    dZ1 = dA1 * (Z1 > 0).float()                        # ReLU gate
    dW1 = X.T @ dZ1                                     # (d_in, h)
    db1 = dZ1.sum(0)
    return dW1, db1, dW2, db2

# gradient check against autograd on a fixed seed
X = torch.randn(8, 4); y = torch.randint(0, 3, (8,))
params = [p.requires_grad_() for p in init_params(4, 16, 3)]
Z1, A1, Z2 = forward(params, X)
loss(Z2, y).backward()
for manual, p in zip(backward(params, X, Z1, A1, Z2, y), params):
    assert torch.allclose(manual, p.grad, atol=1e-6)

import jax
import jax.numpy as jnp

def init_params(key, d_in, d_h, d_out):
    k1, k2 = jax.random.split(key)
    # He for the ReLU layer, Xavier-style 1/fan_in for the linear output
    return dict(
        W1=jax.random.normal(k1, (d_in, d_h)) * jnp.sqrt(2.0 / d_in),
        b1=jnp.zeros(d_h),
        W2=jax.random.normal(k2, (d_h, d_out)) * jnp.sqrt(1.0 / d_h),
        b2=jnp.zeros(d_out))

def forward(params, X):
    Z1 = X @ params["W1"] + params["b1"]      # (N, h)
    A1 = jnp.maximum(Z1, 0.0)                 # ReLU
    Z2 = A1 @ params["W2"] + params["b2"]     # (N, d_out) logits
    return Z1, A1, Z2

def loss(Z2, y):
    logp = jax.nn.log_softmax(Z2)
    return -jnp.take_along_axis(logp, y[:, None], axis=1).mean()

def backward(params, X, Z1, A1, Z2, y):
    """Same pytree shape as params: that is the contract grad honors too."""
    N = X.shape[0]
    P = jax.nn.softmax(Z2, axis=1)
    dZ2 = (P - jax.nn.one_hot(y, Z2.shape[1])) / N      # (N, d_out)
    dA1 = dZ2 @ params["W2"].T                          # (N, h)
    dZ1 = dA1 * (Z1 > 0)                                # ReLU gate
    return dict(W1=X.T @ dZ1, b1=dZ1.sum(0),
                W2=A1.T @ dZ2, b2=dZ2.sum(0))

# gradient check against jax.grad on a fixed key
key = jax.random.PRNGKey(0)
X = jax.random.normal(key, (8, 4))
y = jax.random.randint(key, (8,), 0, 3)
params = init_params(jax.random.PRNGKey(1), 4, 16, 3)
Z1, A1, Z2 = forward(params, X)
auto = jax.grad(lambda p: loss(forward(p, X)[2], y))(params)
manual = backward(params, X, Z1, A1, Z2, y)
for k in params:
    assert jnp.allclose(manual[k], auto[k], atol=1e-5)

Both tabs end with the only test that matters for a hand-written backward: agreement with the framework's own differentiation on a fixed seed, to tolerance. Write this check before trusting a single training step.

The idiomatic version

The second pair is how you would actually write the model. In PyTorch, parameters live inside nn.Module objects and autograd records the tape as the forward runs; calling loss.backward() replays it in reverse and deposits each gradient in p.grad. In JAX, parameters live in an explicit pytree (here a list of per-layer dicts) and jax.grad transforms the loss function into a function returning a gradient pytree of identical structure. Same math, two philosophies: implicit state and recorded tape versus explicit state and function transformation.

import torch
from torch import nn

class MLP(nn.Module):
    def __init__(self, d_in, d_hidden, d_out, depth=2):
        super().__init__()
        layers, d = [], d_in
        for _ in range(depth):
            layers += [nn.Linear(d, d_hidden), nn.ReLU()]
            d = d_hidden
        layers.append(nn.Linear(d, d_out))
        self.net = nn.Sequential(*layers)
        # nn.Linear defaults to Kaiming-uniform; make He-normal explicit
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
                nn.init.zeros_(m.bias)

    def forward(self, x):
        return self.net(x)   # raw logits: CrossEntropyLoss owns the softmax

model = MLP(2, 64, 2)
x, y = torch.randn(32, 2), torch.randint(0, 2, (32,))
loss = nn.functional.cross_entropy(model(x), y)
loss.backward()              # autograd fills p.grad for every parameter

import jax
import jax.numpy as jnp

def init_mlp(key, sizes):
    """sizes = [d_in, h1, ..., d_out]; params is a list-of-dicts pytree."""
    params = []
    for d_in, d_out in zip(sizes[:-1], sizes[1:]):
        key, sub = jax.random.split(key)
        params.append(dict(
            W=jax.random.normal(sub, (d_in, d_out)) * jnp.sqrt(2.0 / d_in),
            b=jnp.zeros(d_out)))
    return params

def apply_mlp(params, x):
    for layer in params[:-1]:
        x = jnp.maximum(x @ layer["W"] + layer["b"], 0.0)
    return x @ params[-1]["W"] + params[-1]["b"]   # raw logits

def loss_fn(params, x, y):
    logp = jax.nn.log_softmax(apply_mlp(params, x))
    return -jnp.take_along_axis(logp, y[:, None], axis=1).mean()

key = jax.random.PRNGKey(0)
params = init_mlp(key, [2, 64, 64, 2])
x = jax.random.normal(key, (32, 2))
y = jax.random.randint(key, (32,), 0, 2)
grads = jax.grad(loss_fn)(params, x, y)   # same pytree structure as params

Using it on a real shape of problem

Two moons is the canonical sanity check for an MLP: two interleaved crescents in the plane that no linear model can separate, easy enough to train full-batch in seconds on a laptop. The loops below generate the data inline (no dependency on scikit-learn), train a two-hidden-layer network of width 64, and print progress. Expect the loss to start near ln 2 ≈ 0.693, the entropy of a coin flip over two classes, drop quickly through the first hundred steps, and settle below 0.05 with accuracy at or near 100% by step 500; exact values vary by machine and seed. If you shrink the hidden width to 2 or 3, the network visibly struggles, which is a nice direct demonstration that capacity is what buys the curved boundary.

import math
import torch
from torch import nn

def two_moons(n, noise=0.08, seed=0):
    g = torch.Generator().manual_seed(seed)
    t = torch.rand(n // 2, generator=g) * math.pi
    upper = torch.stack([t.cos(), t.sin()], dim=1)
    lower = torch.stack([1.0 - t.cos(), 0.5 - t.sin()], dim=1)
    X = torch.cat([upper, lower]) + noise * torch.randn(n, 2, generator=g)
    y = torch.cat([torch.zeros(n // 2), torch.ones(n // 2)]).long()
    return X, y

torch.manual_seed(0)
X, y = two_moons(1024)
model = MLP(2, 64, 2)                       # from the previous block
opt = torch.optim.Adam(model.parameters(), lr=1e-2)

for step in range(501):
    opt.zero_grad()                          # grads accumulate otherwise
    loss = nn.functional.cross_entropy(model(X), y)
    loss.backward()
    opt.step()
    if step % 100 == 0:
        with torch.no_grad():
            acc = (model(X).argmax(1) == y).float().mean()
        print(f"step {step:4d}  loss {loss.item():.4f}  acc {acc:.3f}")
# loss falls from ~0.69 toward ~0.02; accuracy reaches ~1.000

import jax
import jax.numpy as jnp

def two_moons(key, n, noise=0.08):
    k1, k2 = jax.random.split(key)
    t = jax.random.uniform(k1, (n // 2,)) * jnp.pi
    upper = jnp.stack([jnp.cos(t), jnp.sin(t)], axis=1)
    lower = jnp.stack([1.0 - jnp.cos(t), 0.5 - jnp.sin(t)], axis=1)
    X = jnp.concatenate([upper, lower]) + noise * jax.random.normal(k2, (n, 2))
    y = jnp.concatenate([jnp.zeros(n // 2), jnp.ones(n // 2)]).astype(jnp.int32)
    return X, y

X, y = two_moons(jax.random.PRNGKey(0), 1024)
params = init_mlp(jax.random.PRNGKey(1), [2, 64, 64, 2])  # previous block

@jax.jit                                    # whole step compiles to one XLA program
def train_step(params, X, y, lr=0.5):
    loss, grads = jax.value_and_grad(loss_fn)(params, X, y)
    params = jax.tree_util.tree_map(lambda p, g: p - lr * g, params, grads)
    return params, loss

for step in range(501):
    params, loss = train_step(params, X, y)
    if step % 100 == 0:
        acc = (apply_mlp(params, X).argmax(1) == y).mean()
        print(f"step {step:4d}  loss {loss:.4f}  acc {acc:.3f}")
# plain SGD needs a few more steps than Adam but lands in the same place

Applications

Pure MLPs solving whole problems end to end are rarer than they were, but the MLP as a component is everywhere; it is the universal glue of deep learning. The largest single deployment is hiding inside every transformer: the feed-forward block that follows attention in each layer is a two-layer MLP, typically expanding d_model by a factor of four and projecting back, and it holds roughly two thirds of a GPT-style model's parameters. Attention moves information between positions; the MLP is where per-position computation actually happens.

The second big family is heads and adapters. Nearly every transfer-learning recipe ends with a small MLP mapped onto a frozen or fine-tuned backbone: linear-probe and MLP heads on vision encoders, classification heads on BERT-style encoders, reward-model heads on LLMs, the policy and value heads on AlphaGo-lineage systems. Parameter-efficient fine-tuning uses the same shape: adapter layers are small bottleneck MLPs inserted between frozen blocks. Beyond that, MLPs carry whole subfields on their own where inputs are just feature vectors: the two-tower retrieval models behind large-scale recommenders (the YouTube recommendation architecture is the canonical write-up), ranking models over tabular features in ad systems, and coordinate networks such as NeRF, where an MLP maps a 5-D position and viewing direction to color and density and, trained on photos, becomes a 3-D scene.

Against the real libraries

The reference implementations above are honest but you should never hand-roll backprop in production, and it is worth being precise about why. What PyTorch (notes on the codebase at /oss/pytorch) adds over the manual version is not the chain rule, it is everything around it: a differentiation engine with derivatives registered and tested for some two thousand operators, fused and numerically hardened kernels (its cross-entropy never materializes the softmax, exactly the log-sum-exp story from the softmax page), mixed precision, torch.compile graph fusion, and distributed-training plumbing. The nn module layer contributes correct default initializations, train/eval mode handling, and serialization; optimizers like Adam are implemented once, fused, and battle-tested.

On the JAX side the layer libraries are flax, Google's, whose NNX and Linen APIs manage parameter pytrees, initialization, and mutable state such as batch-norm statistics, and equinox, a leaner design where a model simply is a pytree and composes directly with jit, grad, and vmap with no special ceremony. Optimizers come from optax as composable gradient transformations. For the raw pytree style shown above, these libraries add exactly what stops scaling by hand: naming and nesting of hundreds of parameter arrays, per-layer RNG threading, and state that changes outside of gradient descent.

When is the from-scratch version enough? For anything small, fixed, and latency-critical: a two-layer controller exported to a microcontroller, a scoring function embedded in a C++ service, or teaching. The moment you want a second architecture, a schedule, checkpointing, or a GPU, the framework version is both faster and less likely to be subtly wrong. The verification recipe is the one already embedded in the code above and it generalizes: fix a seed, run your manual backward and the library's autograd on the same batch, and assert allclose within about 1e-6 in float32 (looser in reduced precision). Doing backprop by hand once is how you learn what autograd saves, what it costs in memory, and why a shape error two layers deep produces the exact error message it does; doing it twice is a liability.

Traps and misconceptions

Forgetting to zero gradients in PyTorch. backward() accumulates into p.grad rather than overwriting it, a deliberate design for gradient accumulation across micro-batches. Skip opt.zero_grad() and every step applies the running sum of all past gradients; the loss curve goes weird in a way that looks like a bad learning rate. JAX sidesteps this class of bug entirely because gradients are return values, not state.

Putting a softmax before the loss. CrossEntropyLoss and optax.softmax_cross_entropy_with_integer_labels expect raw logits and apply log-softmax internally. Feed them probabilities and you compute log of a softmax of a softmax: training still limps in the right direction, which is what makes the bug so long-lived, but gradients are squashed and the model underperforms for no visible reason.

Zero or constant initialization. It feels like the neutral choice and it is the one initialization guaranteed to fail: identical units receive identical gradients and remain identical forever. Symmetry breaking requires randomness in the weights (zero biases are fine).

Misreading universal approximation. The theorem says a wide-enough single hidden layer can represent any continuous function; it says nothing about whether gradient descent will find that representation, how many units are needed (possibly exponentially many), or generalization. Depth exists because deep composition represents many functions exponentially more compactly than width, not because the theorem demanded it.

Assuming deeper always trains better. A plain MLP stack degrades past a handful of layers as signal and gradient variance drift layer by layer; this is exactly the problem that residual connections and normalization were invented to fix. If a 10-layer plain MLP trains worse than your 3-layer one, that is expected behavior, not a bug in your code.