Network Architecture Design

Learning Objectives

By the end of this section, you will be able to:

1. Identify the four design choices that define an MLP: width, depth, activation function, and weight initialization
2. Explain the effect of width on learning capacity and show that more neurons means more "bends" in the function the network can represent
3. Explain the effect of depth and why deeper networks can represent exponentially more complex functions per parameter
4. Compare activation functions (ReLU, Leaky ReLU, GELU, Tanh) and know when to use each
5. Apply He and Xavier initialization and explain why proper initialization prevents signals from vanishing or exploding
6. Build flexible MLPs in PyTorch using nn.Sequential and nn.init

The Architecture Question

In Chapters 7–9, we learned how neural networks learn: forward propagation computes predictions, backpropagation computes gradients, and optimizers update weights. But we always used the same architecture: $4 \to 3 \to 4$ with ReLU and random initialization.

Why 3 hidden neurons? Why not 2, or 8, or 100? Why one hidden layer instead of three? Why ReLU instead of tanh? These choices were arbitrary — and that's a problem, because architecture determines what a network can and cannot learn.

The Central Question of This Chapter: Given a task (like diagonal flip), how do we choose the right architecture? What are the design knobs, and what happens when we turn them?

Think of it like building a house. Knowing how to lay bricks (forward pass), mix mortar (backprop), and plan the schedule (optimizer) is essential — but it does not tell you how many rooms to build or how tall to make the ceilings. Architecture is the blueprint that determines what the finished structure can do.

What Defines an MLP?

A Multi-Layer Perceptron (MLP) — also called a feedforward neural network or fully connected network — has four architectural choices:

Each choice has consequences. Width controls how many features a single layer can detect. Depth controls how many levels of abstraction the network builds. Activation controls the shape of the non-linearity. Initialization controls whether training starts in a healthy state.

Let us examine each one, starting with the most intuitive: width.

Width: How Many Neurons Per Layer?

Consider our diagonal flip task. The network must learn to swap positions 1 and 2 of a 4-element vector while keeping positions 0 and 3 unchanged. How many hidden neurons does it need?

The Effect of Width

Each hidden neuron with ReLU activation contributes one "bend" — a point where the function changes slope. A neuron computes $h_j = \max(0, \mathbf{w}_j^T \mathbf{x} + b_j)$. This is a hinge function: linear on one side, zero on the other. The boundary between the two sides is a hyperplane defined by $\mathbf{w}_j^T \mathbf{x} + b_j = 0$.

With $n$ hidden neurons in a single layer, the network can partition the input space into at most $n + 1$ linear regions. Within each region, the network computes a different linear function. The more regions, the more complex the overall function.

The pattern is clear: more neurons → lower loss. With only 2 neurons, the network can create 3 linear regions — not enough to capture the diagonal flip mapping for all 16 images. With 16 neurons, it has 17 regions and achieves nearly zero loss.

Width as Feature Detectors

Each hidden neuron acts as a feature detector. It learns to respond to a specific pattern in the input. For example, one neuron might detect "position 1 is 1 AND position 2 is 0", while another detects the reverse. The output layer then combines these features to produce the final mapping.

With too few neurons, the network lacks enough feature detectors to distinguish all the patterns it needs. It is forced to make compromises — mapping several distinct inputs to similar outputs. This is underfitting: the model is too simple for the task.

Depth: How Many Hidden Layers?

Width is not the only way to add capacity. We can also add depth — more hidden layers. And depth has a remarkable property: it adds capacity exponentially rather than linearly.

Why Depth Matters

Consider a single hidden layer with $n$ ReLU neurons. It creates at most $n + 1$ linear regions. Now add a second hidden layer with $n$ neurons. Each neuron in the second layer can "fold" the first layer's regions, effectively multiplying the number of regions:

• 1 hidden layer with $n$ neurons: up to $n + 1$ regions
• 2 hidden layers with $n$ neurons each: up to $O(n^2)$ regions
• $L$ hidden layers with $n$ neurons each: up to $O(n^L)$ regions

This is the exponential advantage of depth. A network with 4 neurons per layer and 3 hidden layers has only $4 \times 4 + 4 + 4 \times 4 + 4 + 4 \times 4 + 4 = 60$ parameters, but can theoretically represent $O(4^3) = 64$ linear regions. A single-layer network would need $63$ neurons (and 507 parameters!) to match that number of regions.

Depth vs Width: A Practical Comparison

For small problems like diagonal flip, width often wins because the optimization is easier. For complex problems like image recognition or language understanding, depth wins because real-world data has hierarchical structure — edges combine into shapes, shapes combine into objects, objects combine into scenes. Each layer captures one level of abstraction.

The depth principle: If your task has natural hierarchy (low-level features composing into high-level concepts), add depth. If your task is a complex mapping without obvious hierarchy, add width. Most real-world tasks have hierarchy, which is why modern networks are deep.

Interactive: Architecture Explorer

Use the sliders below to adjust the number of neurons in each layer. Watch how the network structure changes — more neurons mean more connections (weights), and each connection is a learnable parameter. Click "Animate Forward Pass" to see data flow through the layers.

Try these experiments:

1. Set all hidden layers to 1 neuron — notice the extreme bottleneck. Information must squeeze through a single number.
2. Set a hidden layer to 8 neurons — the web of connections grows dramatically. Each connection is a weight the optimizer must tune.
3. Compare 4→8→4 versus 4→4→4→4 — similar parameter counts, but very different structures.

Interactive: Width vs Depth

This visualization shows a real neural network learning a curved target function (dashed line). Adjust width and depth, then click Train to watch the cyan prediction line converge toward the target. The network diagram on the left updates to show the current architecture.

Try these experiments:

1. Width = 2, Depth = 1: Train and watch the network struggle. With only 2 neurons, it can make at most 3 linear segments — not enough to match the wavy target.
2. Width = 8, Depth = 1: Now 9 segments are possible. The fit is much better, but look at the parameter count.
3. Width = 4, Depth = 3: Fewer parameters than Width=8/Depth=1, but potentially more linear regions due to the exponential depth effect. Does it learn better?
4. Width = 16, Depth = 1: Brute force with many neurons. Fast convergence but many parameters.

Activation Functions: Beyond ReLU

We have used ReLU exclusively so far. But there are several activation functions, each with distinct properties. The choice of activation is the third architectural decision.

The Activation Function Zoo

ReLU: The Workhorse

ReLU ($\max(0, x)$) is the default activation for hidden layers. It is fast to compute (just a comparison), produces sparse activations (many neurons output exactly 0), and has a simple gradient (1 if active, 0 if not).

The main weakness is the dying ReLU problem. If a neuron's pre-activation is always negative (for all training inputs), its gradient is always 0, and it can never recover. The neuron is permanently "dead." This happens more with large learning rates or poor initialization.

Leaky ReLU: Fixing Dead Neurons

Leaky ReLU replaces the flat zero region with a small negative slope: $f(x) = \max(\alpha x, x)$ where $\alpha = 0.01$. Now even negative inputs produce a small gradient ($\alpha$), so no neuron can completely die. The output is $x$ when positive, $0.01x$ when negative.

GELU: The Transformer Standard

GELU (Gaussian Error Linear Unit) is used in GPT, BERT, and most modern transformers. It is defined as $\text{GELU}(x) = x \cdot \Phi(x)$ where $\Phi(x)$ is the standard Gaussian CDF. Unlike ReLU's hard cutoff at 0, GELU provides a smooth transition. Small negative values are slightly attenuated rather than completely zeroed out. This smoothness helps gradient-based optimization, especially in very deep networks.

Tanh and Sigmoid: The Classics

Before ReLU, tanh and sigmoid were standard. Tanh ($\tanh(x) = \frac{e^x - e^{-x}}{e^x + e^{-x}}$) outputs values in $(-1, 1)$ and is zero-centered, which helps optimization. Sigmoid ($\sigma(x) = \frac{1}{1 + e^{-x}}$) outputs values in $(0, 1)$ and is mainly used for output layers where we need probabilities.

The main problem with both is vanishing gradients: for large $|x|$, the derivative approaches 0. In deep networks, this means early layers receive almost no gradient signal and learn extremely slowly. ReLU solved this by having a constant gradient of 1 for positive inputs.

Weight Initialization: Why It Matters

The fourth architectural choice is how to set the initial weight values. This might seem trivial — they are just starting values that will be trained away — but initialization has a dramatic impact on whether training succeeds at all.

The Variance Problem

Consider a layer with $n$ inputs: $z = \sum_{i=1}^{n} w_i x_i$. If inputs have variance $\text{Var}(x) = 1$ and weights have variance $\text{Var}(w) = \sigma^2$, then by the properties of independent random variables:

$\text{Var}(z) = n \cdot \sigma^2 \cdot \text{Var}(x) = n\sigma^2$

The output variance is $n$ times the weight variance. For a layer with $n = 512$ inputs and naive $\sigma = 1$ initialization, the output variance is 512 — the signal explodes. After a few layers, activations overflow to infinity. Conversely, if $\sigma = 0.01$, the output variance is $512 \times 0.0001 = 0.0512$ — the signal vanishes toward zero.

Xavier Initialization (for Tanh/Sigmoid)

Xavier Glorot (2010) proposed: set $\sigma = \sqrt{\frac{2}{n_{\text{in}} + n_{\text{out}}}}$. This keeps the variance of both the forward pass and the backward pass approximately 1. It is designed for linear or tanh activations that preserve their input variance.

He Initialization (for ReLU)

Kaiming He (2015) observed that ReLU kills approximately half the activations (all negative values become 0), which halves the variance. To compensate, He init uses $\sigma = \sqrt{\frac{2}{n_{\text{in}}}}$ — doubling the variance compared to Xavier. This is the standard for ReLU networks and what we have been using in our code.

Look at the epoch 99 losses. He initialization reaches the lowest loss (0.0244) — it is specifically designed for ReLU. Xavier is close (0.0305) but slightly worse because it underestimates the variance needed for ReLU. Small random initialization (0.0984) converges 4× slower because the signal vanishes through layers. Large random (0.0658) starts with high loss because the initial predictions are wildly wrong, wasting early training steps recovering.

Rule of thumb: Use He initialization with ReLU (the default in PyTorch). Use Xavier initialization with tanh or sigmoid. Never use plain random initialization — the scale matters enormously.

NumPy: Comparing Widths from Scratch

Let us put the theory into practice. The following code trains three architectures — 4→2→4, 4→8→4, and 4→16→4 — on our diagonal flip task and compares their final losses. This is the same forward/backward code from Chapters 7–8, now wrapped in a function so we can vary the width.

1import numpy as np
2
3# ── Dataset: 2×2 binary images → diagonal flip ──
4# Flattened: [a,b,c,d] → [a,c,b,d] (transpose)
5X = np.array([[i,j,k,l]
6              for i in [0,1] for j in [0,1]
7              for k in [0,1] for l in [0,1]], dtype=float)
8Y = np.array([[x[0],x[2],x[1],x[3]] for x in X])
9
10def relu(x):
11    return np.maximum(0, x)
12
13# ── Train a 4→width→4 network ──
14def train(width, epochs=200, lr=0.05):
15    np.random.seed(42)
16    W1 = np.random.randn(4, width) * np.sqrt(2/4)
17    b1 = np.zeros(width)
18    W2 = np.random.randn(width, 4) * np.sqrt(2/width)
19    b2 = np.zeros(4)
20
21    for epoch in range(epochs):
22        loss = 0.0
23        for i in range(16):
24            h = relu(X[i] @ W1 + b1)
25            y_hat = h @ W2 + b2
26            loss += 0.5 * np.mean((y_hat - Y[i])**2)
27
28            dout = (y_hat - Y[i]) / 4
29            dW2 = np.outer(h, dout)
30            db2 = dout
31            dh = dout @ W2.T * (X[i] @ W1 + b1 > 0)
32            dW1 = np.outer(X[i], dh)
33            db1 = dh
34
35            W1 -= lr * dW1; b1 -= lr * db1
36            W2 -= lr * dW2; b2 -= lr * db2
37
38        if epoch % 50 == 0 or epoch == 199:
39            print(f"  Epoch {epoch:3d}: loss = {loss/16:.4f}")
40
41    params = 4*width + width + width*4 + 4
42    return loss/16, params
43
44# ── Compare three widths ──
45for w in [2, 8, 16]:
46    print(f"\nWidth = {w} (4→{w}→4):")
47    final_loss, params = train(w)
48    print(f"  Parameters: {params}, Final loss: {final_loss:.6f}")

The results tell a clear story:

• Width 2 (22 params): Loss plateaus at 0.065 — the network simply cannot represent the diagonal flip with only 3 linear regions.
• Width 8 (76 params): Loss reaches 0.024 — 2.7× better. With 9 linear regions, the network has enough capacity to approximate the mapping.
• Width 16 (148 params): Loss reaches 0.0005 — nearly perfect. Each of the 17 linear regions precisely captures part of the input-output mapping.

PyTorch: Flexible MLP Builder

In practice, we do not write the forward and backward passes by hand. PyTorch's $\texttt{nn.Sequential}$ lets us define any MLP architecture in a few lines, and autograd handles all the gradient computation. The following code introduces a factory function pattern that builds any MLP from a list of layer sizes.

1import torch
2import torch.nn as nn
3
4# ── Build any MLP architecture with nn.Sequential ──
5def make_mlp(layer_sizes):
6    layers = []
7    for i in range(len(layer_sizes) - 1):
8        layers.append(nn.Linear(layer_sizes[i], layer_sizes[i+1]))
9        if i < len(layer_sizes) - 2:
10            layers.append(nn.ReLU())
11    model = nn.Sequential(*layers)
12
13    # Apply He initialization
14    for m in model.modules():
15        if isinstance(m, nn.Linear):
16            nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
17            nn.init.zeros_(m.bias)
18
19    return model
20
21# ── Dataset ──
22X = torch.tensor([[i,j,k,l]
23                   for i in [0,1] for j in [0,1]
24                   for k in [0,1] for l in [0,1]], dtype=torch.float32)
25Y = torch.stack([torch.tensor([x[0],x[2],x[1],x[3]]) for x in X])
26
27# ── Compare architectures ──
28configs = {
29    "4→2→4":       [4, 2, 4],
30    "4→8→4":       [4, 8, 4],
31    "4→16→4":      [4, 16, 4],
32    "4→4→4→4":     [4, 4, 4, 4],
33    "4→8→8→4":     [4, 8, 8, 4],
34}
35
36torch.manual_seed(42)
37for name, sizes in configs.items():
38    model = make_mlp(sizes)
39    params = sum(p.numel() for p in model.parameters())
40    opt = torch.optim.Adam(model.parameters(), lr=0.01)
41
42    for epoch in range(200):
43        pred = model(X)
44        loss = nn.functional.mse_loss(pred, Y)
45        opt.zero_grad()
46        loss.backward()
47        opt.step()
48
49    print(f"{name:14s}  params={params:4d}  loss={loss.item():.6f}")

The $\texttt{make\_mlp}$ function is a pattern you will use throughout your deep learning career. It separates architecture definition (the list of sizes) from architecture construction (the loop that builds layers). This makes experimentation trivial: just change the list.

Architecture Design Heuristics

There is no formula for the optimal architecture. But decades of practice have produced reliable heuristics:

Starting Points

1. Start with one hidden layer. A single hidden layer can approximate any continuous function (the Universal Approximation Theorem, which we prove in Section 3 of this chapter). Start simple and add depth only if needed.
2. Set width to 2–4× the input dimension. For our 4-input task, that suggests 8–16 hidden neurons — which matches our experimental results. For a 784-input MNIST task, start with 256–512 hidden neurons.
3. Use a pyramid or funnel shape. In multi-layer networks, gradually decrease width: 784→512→256→128→10. This forces the network to compress information into progressively more abstract representations.
4. Equal-width layers work too. Research by Hanin and Rolnick (2019) showed that equal-width networks (same number of neurons in every hidden layer) perform surprisingly well. The funnel shape is traditional but not always superior.

Practical Decision Tree

Connection to Modern Systems

Every concept in this section scales directly to the architectures powering today's AI systems:

Transformers Use MLPs

Every transformer block contains a feed-forward network (FFN) — which is just a 2-layer MLP! In GPT-4 and LLaMA, each FFN takes the form $d_{\text{model}} \to 4d_{\text{model}} \to d_{\text{model}}$. For a model with $d_{\text{model}} = 4096$, that is a $4096 \to 16384 \to 4096$ MLP with over 67 million parameters per layer. The width factor of 4× matches our heuristic.

GELU in Practice

Modern transformers (BERT, GPT, LLaMA) all use GELU activation in their FFN layers, not ReLU. The smooth gradient of GELU helps training stability in networks with 100+ layers. Some variants like LLaMA use SwiGLU, a gated variant that combines two linear transformations: $\text{SwiGLU}(x) = (xW_1) \cdot \sigma(xW_2)$ where $\sigma$ is the sigmoid function. This gives each neuron a learnable gate that controls information flow.

Initialization at Scale

Proper initialization becomes even more critical at scale. GPT-3 (175B parameters, 96 layers) uses a modified initialization where the output projection of each residual block is scaled by $1/\sqrt{2N}$ where $N$ is the number of layers. Without this, the residual connections cause variance to grow with depth. The principle is the same as He initialization — keep the signal variance stable — but adapted for the specific architecture.

Width and Memory

The width of the FFN layers is the largest memory consumer in transformers. During inference with KV-cache, the attention layers store cached key/value pairs (proportional to sequence length), but the FFN layers store the full weight matrices (proportional to $d_{\text{model}}^2$). This is why techniques like MoE (Mixture of Experts) in models like Mixtral-8x7B activate only a subset of the FFN width for each token — getting the capacity of a wide network with the compute cost of a narrow one.

Summary

MLP architecture design comes down to four choices, each with measurable consequences:

1. Width (neurons per layer) controls how many features the network can detect. More width = more linear regions = more complex functions. But more parameters risk overfitting.
2. Depth (number of hidden layers) adds capacity exponentially. $L$ layers of width $n$ can create $O(n^L)$ linear regions. But deeper networks are harder to optimize (vanishing gradients).
3. Activation function determines the shape of the non-linearity. ReLU is the default, GELU for transformers, Leaky ReLU if dead neurons are a problem.
4. Initialization determines whether the signal survives through layers. He init for ReLU ($\sigma = \sqrt{2/n_{\text{in}}}$), Xavier for tanh/sigmoid ($\sigma = \sqrt{2/(n_{\text{in}} + n_{\text{out}})}$).

We demonstrated these concepts on the diagonal flip task: width 2 could not learn (loss 0.065), width 8 learned reasonably (loss 0.024), and width 16 achieved near-perfect accuracy (loss 0.0005). In PyTorch, the $\texttt{make\_mlp}$ factory pattern makes architecture experimentation trivial.

In the next section, we will put this knowledge to work: building complete MLPs in PyTorch, training them on real data, and systematically comparing architectures using proper train/validation splits.