PythonMastery
advanced 25 min read · lesson 1 of 9 in AI & Deep Learning

Neural Networks: Architectures and Theory

1 · The lesson

read

You already know that a neural network is layers of weighted sums followed by non-linearities, trained by gradient descent. That is the overview. This lesson goes deeper: the linear-algebra notation that papers use, why a deep network with tanh activations can refuse to train, why the Adam optimiser is used almost everywhere, and why a 152-layer ResNet is even possible.

Everything here is grounded in one fact: training a deep network is an optimisation problem in millions of parameters, with a non-convex loss surface, fitted by a noisy first-order method. Every technique in this lesson (activation choice, initialisation, normalisation, residuals, optimisers, schedules) exists because that fact causes real problems in training.


1. The MLP in Formal Notation

A multi-layer perceptron (MLP) is the basic architecture. Layer $\ell$ takes input $\mathbf{a}^{(\ell-1)} \in \mathbb{R}^{n_{\ell-1}}$ and produces:

$$
\mathbf{z}^{(\ell)} = W^{(\ell)} \mathbf{a}^{(\ell-1)} + \mathbf{b}^{(\ell)}, \qquad \mathbf{a}^{(\ell)} = \sigma\!\left(\mathbf{z}^{(\ell)}\right)
$$

  • $W^{(\ell)} \in \mathbb{R}^{n_\ell \times n_{\ell-1}}$: the weight matrix
  • $\mathbf{b}^{(\ell)} \in \mathbb{R}^{n_\ell}$: the bias
  • $\sigma$: the activation function, applied element-wise
  • $\mathbf{z}^{(\ell)}$: the pre-activation (the value before $\sigma$)
  • $\mathbf{a}^{(\ell)}$: the activation (the value after $\sigma$)

The full forward pass for an $L$-layer network is the function composition:

$$
f(\mathbf{x}) = \sigma_L \circ W_L \circ \sigma_{L-1} \circ W_{L-1} \circ \cdots \circ \sigma_1 \circ W_1 (\mathbf{x})
$$

Without the $\sigma$'s, the whole thing collapses to a single linear map $W_L W_{L-1} \cdots W_1 \mathbf{x}$, no matter how many layers you stack. The non-linearity is what lets the network represent complex functions. A network of any depth with only linear activations is exactly equivalent to one linear layer.

In pure NumPy code, it looks like this:

python
import numpy as np

def forward(x, params):
    a = x
    for W, b, sigma in params:
        z = W @ a + b
        a = sigma(z)
    return a

That is all. Everything else in this lesson is about making that loop trainable.


2. Activation Functions: Picking the Right Non-Linearity

The activation function determines the gradient signal that flows back through the network. If you choose badly, your gradients vanish, explode, or stay at zero forever.

Sigmoid: common in the past, avoided today

$$
\sigma(z) = \frac{1}{1+e^{-z}}, \qquad \sigma'(z) = \sigma(z)(1-\sigma(z))
$$

The derivative peaks at $0.25$ when $z = 0$ and decays to zero on both sides. Stack ten sigmoid layers and the gradient is multiplied by at most $0.25^{10} \approx 10^{-6}$. This is the vanishing gradient problem. Sigmoid is still used in the output layer of a binary classifier (it maps $\mathbb{R} \to (0,1)$), but it is no longer used in hidden layers.

Tanh: a zero-centred version of sigmoid

$$
\tanh(z) = \frac{e^z - e^{-z}}{e^z + e^{-z}}, \qquad \tanh'(z) = 1 - \tanh^2(z)
$$

Its output $(-1, 1)$ is centred on zero, which makes optimisation slightly easier than with sigmoid. It still saturates on both sides. It is used in older RNNs and for the cell/candidate values inside LSTMs and GRUs (their gates use sigmoid).

ReLU: the standard choice since 2012

$$
\text{ReLU}(z) = \max(0, z), \qquad \text{ReLU}'(z) = \begin{cases}1 & z > 0 \\ 0 & z \le 0\end{cases}
$$

It is cheap to compute, it doesn't saturate on the positive side, and its gradient is exactly $1$ for active units, so gradients pass through many layers without shrinking. ReLU made much deeper networks practical to train.

The dying-ReLU problem: if an unlucky update pushes a unit to $z < 0$ for every input, its gradient is $0$, so no future update can change it. Dead units stay dead. With a high learning rate, a large part of your network can die in the first epoch and never recover.

Leaky ReLU: keep the dead units alive

$$
\text{LReLU}(z) = \begin{cases}z & z > 0 \\ \alpha z & z \le 0\end{cases}, \qquad \alpha = 0.01 \text{ (typical)}
$$

A small slope for negative inputs means these units still get a gradient. This is a cheap way to prevent the dying-ReLU problem.

GELU: the transformer default

$$
\text{GELU}(z) = z \cdot \Phi(z)
$$

where $\Phi$ is the standard normal CDF. It is smooth and non-monotonic, and it scales each input by a factor that depends on its value instead of cutting it off sharply at zero. It is used in BERT, GPT and ViT, and in many other modern transformers. It is slightly more expensive than ReLU, but in practice it usually works better in transformer language models.

Swish / SiLU: used in EfficientNet, Llama

$$
\text{Swish}(z) = z \cdot \sigma(\beta z)
$$

It is smooth, it multiplies the input by its own sigmoid ("self-gated"), and it performs similarly to GELU. nn.SiLU() in PyTorch, keras.activations.swish in Keras.

When to use what: a summary

ActivationUse when
ReLUDefault for CNNs and most MLPs. Cheap, fast, reliable.
Leaky ReLU / PReLUYou're seeing dead units, or training very deep without BN.
GELUTransformers. Now the default for LLMs and ViT.
Swish/SiLUEfficientNet-style CNNs, modern Llama-family LLMs.
TanhLSTM/GRU cell and candidate values (the gates themselves are sigmoid). Almost nowhere else.
SigmoidOutput layer of binary classifier. Never hidden.
SoftmaxOutput layer of multi-class classifier.

3. Weight Initialisation: Why Random Isn't Enough

If every weight starts at the same value, every hidden unit computes the same thing. This is called symmetry. Gradient descent updates them identically. You effectively have a one-unit-per-layer network. So weights must be random.

But how you choose the random values matters a lot. Consider a layer with $n_{\text{in}}$ inputs. The pre-activation variance is:

$$
\operatorname{Var}(z) = n_{\text{in}} \cdot \operatorname{Var}(W) \cdot \operatorname{Var}(a)
$$

If $\operatorname{Var}(W)$ is too large, $\operatorname{Var}(z)$ grows at every layer, so the pre-activations and the gradients explode. If $\operatorname{Var}(W)$ is too small, signal vanishes layer by layer.

Goal: pick $\operatorname{Var}(W)$ so that signal variance is preserved layer to layer.

Xavier / Glorot initialisation (2010)

For symmetric activations (tanh, sigmoid):

$$
W \sim \mathcal{N}\!\left(0, \frac{2}{n_{\text{in}} + n_{\text{out}}}\right) \quad \text{or} \quad W \sim \mathcal{U}\!\left(-\sqrt{\tfrac{6}{n_{\text{in}}+n_{\text{out}}}}, \sqrt{\tfrac{6}{n_{\text{in}}+n_{\text{out}}}}\right)
$$

It keeps the variance roughly constant in both directions: forward (activations) and backward (gradients).

He initialisation (2015)

For ReLU and its variants. ReLU sets about half of its inputs to zero, so you need double the variance to compensate:

$$
W \sim \mathcal{N}\!\left(0, \frac{2}{n_{\text{in}}}\right)
$$

Always use He init when your activations are ReLU/LReLU/GELU/Swish. Keras: kernel_initializer="he_normal". PyTorch: nn.init.kaiming_normal_.

python
# Keras: pick the initialisation that matches the activation
from tensorflow.keras import layers, initializers

layers.Dense(128, activation="relu",
             kernel_initializer=initializers.HeNormal())

layers.Dense(64,  activation="tanh",
             kernel_initializer=initializers.GlorotNormal())

If the initialisation doesn't match the activation, the loss curve can look wrong from epoch 1: flat, NaN, or stuck. If they match, the first few epochs usually converge smoothly.


4. Backpropagation: A Worked Example

Let's do one full backward pass by hand. Two-layer network, scalar input, scalar output, MSE loss. No batching, no broadcasting, only the chain rule.

Forward:
$$
z_1 = w_1 x + b_1, \quad a_1 = \text{ReLU}(z_1), \quad z_2 = w_2 a_1 + b_2, \quad \hat{y} = z_2
$$
$$
\mathcal{L} = \tfrac{1}{2}(\hat{y} - y)^2
$$

Concrete numbers: $x = 2$, $y = 1$, $w_1 = 0.5$, $b_1 = 0$, $w_2 = -1$, $b_2 = 0.3$.

  • $z_1 = 0.5 \cdot 2 + 0 = 1.0$
  • $a_1 = \text{ReLU}(1.0) = 1.0$
  • $z_2 = -1 \cdot 1.0 + 0.3 = -0.7$
  • $\hat{y} = -0.7$
  • $\mathcal{L} = \tfrac{1}{2}(-0.7 - 1)^2 = \tfrac{1}{2}(2.89) = 1.445$

Backward: work from the output back, multiplying derivatives (chain rule):

$$
\frac{\partial \mathcal{L}}{\partial \hat{y}} = \hat{y} - y = -1.7
$$

$$
\frac{\partial \mathcal{L}}{\partial z_2} = \frac{\partial \mathcal{L}}{\partial \hat{y}} \cdot \frac{\partial \hat{y}}{\partial z_2} = -1.7 \cdot 1 = -1.7
$$

$$
\frac{\partial \mathcal{L}}{\partial w_2} = \frac{\partial \mathcal{L}}{\partial z_2} \cdot a_1 = -1.7 \cdot 1.0 = -1.7
$$

$$
\frac{\partial \mathcal{L}}{\partial b_2} = \frac{\partial \mathcal{L}}{\partial z_2} = -1.7
$$

$$
\frac{\partial \mathcal{L}}{\partial a_1} = \frac{\partial \mathcal{L}}{\partial z_2} \cdot w_2 = -1.7 \cdot -1 = 1.7
$$

$$
\frac{\partial \mathcal{L}}{\partial z_1} = \frac{\partial \mathcal{L}}{\partial a_1} \cdot \text{ReLU}'(z_1) = 1.7 \cdot 1 = 1.7 \quad (\text{since } z_1 > 0)
$$

$$
\frac{\partial \mathcal{L}}{\partial w_1} = \frac{\partial \mathcal{L}}{\partial z_1} \cdot x = 1.7 \cdot 2 = 3.4
$$

$$
\frac{\partial \mathcal{L}}{\partial b_1} = \frac{\partial \mathcal{L}}{\partial z_1} = 1.7
$$

Gradient descent step with learning rate $\eta = 0.1$:

  • $w_1 \leftarrow 0.5 - 0.1 \cdot 3.4 = 0.16$
  • $w_2 \leftarrow -1 - 0.1 \cdot (-1.7) = -0.83$
  • $b_1 \leftarrow 0 - 0.1 \cdot 1.7 = -0.17$
  • $b_2 \leftarrow 0.3 - 0.1 \cdot (-1.7) = 0.47$

That is one training step. Vector-valued inputs, batched data, and matrix weights follow this exact pattern: every gradient is a product of partial derivatives along the path from $\mathcal{L}$ back to the parameter. Autograd in TensorFlow/PyTorch builds the computation graph and does this multiplication for you, but it is the same chain rule on each edge.

If $z_1$ had been negative, $\text{ReLU}'(z_1) = 0$ would make every gradient before that node zero. That is why dying ReLU is so harmful.


5. Loss Functions: More Than MSE

The loss defines what "correct" means. If you choose the wrong loss, the optimiser can converge to something useless without any error.

LossFormulaWhen
MSE$\frac{1}{n}\sum (y-\hat{y})^2$Regression, Gaussian noise, no outliers
MAE / L1$\frac{1}{n}\sum \lvert y-\hat{y} \rvert$Regression with outliers (less sensitive to them)
Huberquadratic near zero, linear far awayRegression that needs both: a robust default
Binary cross-entropy$-[y \log\hat{y} + (1-y)\log(1-\hat{y})]$Binary classification with sigmoid output
Categorical cross-entropy$-\sum_k y_k \log \hat{y}_k$Multi-class with softmax output (one-hot $y$)
Sparse cat. cross-entropysame, but $y$ is integer labelSame as above, cheaper when classes are huge
Focal loss$-(1-\hat{y})^\gamma \log\hat{y}$Heavy class imbalance (gives less weight to easy examples)
Contrastive / tripletbased on distances in embedding spaceMetric learning (face recognition, embeddings)

The loss changes the shape of the problem. MSE on sigmoid or softmax outputs gives very small gradients for confident wrong predictions, so it corrects them slowly. Cross-entropy expects probabilities between 0 and 1, so it doesn't make sense for regression targets. Choose the loss for the task before you do anything else.

A note on focal loss: in object detection, the vast majority of anchor boxes are background. With plain cross-entropy, these easy negatives make up most of the loss. Focal loss adds a $(1-\hat{y})^\gamma$ factor that reduces the loss for confident, correct examples, so the network focuses on the hard ones.


6. Optimisers: The Update Rules

Vanilla SGD updates parameters with $\theta \leftarrow \theta - \eta \nabla_\theta \mathcal{L}$. It works, but it is slow, and it can get stuck in narrow valleys and at saddle points. Every modern optimiser is SGD with extra techniques added.

SGD with momentum

Momentum accumulates an exponentially weighted velocity:

$$
v_t = \mu v_{t-1} + \nabla_\theta \mathcal{L}, \qquad \theta \leftarrow \theta - \eta v_t
$$

$\mu = 0.9$ is the usual choice. Momentum speeds up movement in consistent directions and reduces oscillation in noisy ones.

Nesterov momentum

Look ahead: compute the gradient at $\theta - \eta \mu v_{t-1}$ rather than $\theta$. It converges better on convex problems, but in deep learning the difference is small.

RMSprop

Per-parameter learning rate, scaled by recent gradient magnitude:

$$
s_t = \beta s_{t-1} + (1-\beta)\, g_t^2, \qquad \theta \leftarrow \theta - \frac{\eta}{\sqrt{s_t}+\epsilon}\, g_t
$$

Where $g_t = \nabla_\theta \mathcal{L}$. Parameters with consistently large gradients get a smaller effective rate. It was the most common choice for RNNs before Adam.

Adam: adaptive moment estimation

Momentum + RMSprop combined:

$$
m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t \quad (\text{first moment})
$$
$$
v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2 \quad (\text{second moment})
$$
$$
\hat{m}_t = \frac{m_t}{1-\beta_1^t}, \quad \hat{v}_t = \frac{v_t}{1-\beta_2^t} \quad (\text{bias correction})
$$
$$
\theta \leftarrow \theta - \frac{\eta}{\sqrt{\hat{v}_t}+\epsilon}\, \hat{m}_t
$$

Defaults: $\beta_1=0.9$, $\beta_2=0.999$, $\epsilon=10^{-8}$, $\eta=10^{-3}$. Adam works well on almost every problem with its default settings, so it is the right starting point unless you have a reason to choose something else.

AdamW: decoupled weight decay

Original Adam applies L2 regularisation by adding $\lambda \theta$ to the gradient, then dividing by $\sqrt{\hat{v}}$. This means the weight decay gets scaled by the per-parameter learning rate. Loshchilov & Hutter showed that this weakens the intended regularisation effect. AdamW decouples:

$$
\theta \leftarrow \theta - \eta \left(\frac{\hat{m}_t}{\sqrt{\hat{v}_t}+\epsilon} + \lambda \theta\right)
$$

The $\lambda \theta$ term is applied separately, not coupled to the adaptive scaling. AdamW is the standard optimiser for transformers and for most modern training. Keras: keras.optimizers.AdamW(weight_decay=0.01).


7. Learning Rate Schedules

A single fixed learning rate is rarely right for the whole of training. You want a high rate early, to move away from saddle points quickly, and a low rate later, to settle into a minimum.

  • Step decay: reduce $\eta$ by a factor every $k$ epochs. Simple, but you must choose the steps by hand.
  • Exponential decay: $\eta_t = \eta_0 \gamma^t$. Smoother than step decay.
  • Cosine annealing: $\eta_t = \eta_{\min} + \tfrac{1}{2}(\eta_{\max} - \eta_{\min})(1 + \cos(\pi t / T))$. A default for modern training. The rate decreases smoothly to almost zero.
  • Warmup: increase $\eta$ linearly from $0$ to $\eta_0$ over the first $w$ steps. This is very important for transformers: large models trained with adaptive optimisers often fail without it, because the variance estimates in the early steps are unreliable.
  • OneCycle (Smith, 2017): warmup, then cosine decay, with momentum changed in the opposite direction. It uses high learning rates, and it often works very well for CNNs.
python
import tensorflow as tf

schedule = tf.keras.optimizers.schedules.CosineDecay(
    initial_learning_rate=1e-3,
    decay_steps=10_000,
    alpha=0.01,            # final learning rate = alpha * initial rate
)
opt = tf.keras.optimizers.AdamW(learning_rate=schedule, weight_decay=0.01)

The schedule lives inside the optimiser. Pass learning_rate=schedule and Keras handles the per-step update.


8. Batch Normalisation: What It Does

For each mini-batch and each feature, BatchNorm normalises:

$$
\hat{z} = \frac{z - \mu_{\text{batch}}}{\sqrt{\sigma^2_{\text{batch}} + \epsilon}}, \qquad \text{out} = \gamma \hat{z} + \beta
$$

$\gamma, \beta$ are learned per-feature scale and shift. At test time, $\mu$ and $\sigma$ come from running averages tracked during training.

The original (2015) paper said BN works by reducing "internal covariate shift". Many researchers now doubt that explanation. A 2018 paper by Santurkar et al. argued that BN works by smoothing the loss surface, which makes the gradients more stable. Whatever the exact reason, the observed effect is: BN lets you use higher learning rates, train deeper networks, and acts as a mild regulariser.

An important trap: BN behaves differently during training and during evaluation. Keras handles this for you: predict(), evaluate() and model(x) run BN in inference mode (trainable isn't the switch). In PyTorch, if you forget .eval() at inference, BN uses the statistics of each new batch and keeps updating the running averages with whatever data you feed in. This breaks the model without any error message.


9. Layer / Group / Instance Normalisation

BatchNorm breaks down when the batch is small (batch=1 in RL, batch=2 in 3D vision). The alternatives normalise over different axes:

NormNormalises over
BatchNormEach feature across the batch dimension
LayerNormAll features within a single sample
InstanceNormEach channel within a single sample, over its spatial positions (used in style transfer)
GroupNormEach group of channels within a single sample (between LayerNorm and InstanceNorm)

LayerNorm is the usual choice for transformers, because it doesn't depend on batch size at all. Almost every transformer block uses LayerNorm (or a close variant such as RMSNorm). GroupNorm is the usual choice for computer vision with small batches. It is common in detection and segmentation, where a batch size of 2 is normal.


10. Residual Connections: Why Very Deep Networks Work

A residual block computes $y = F(x) + x$ rather than $y = F(x)$. The skip connection lets the gradient flow directly through:

$$
\frac{\partial \mathcal{L}}{\partial x} = \frac{\partial \mathcal{L}}{\partial y} \cdot \left(\frac{\partial F}{\partial x} + I\right)
$$

The $+I$ term means the gradient always has a "shortcut" path: even if $\partial F/\partial x$ vanishes, the gradient still flows back through the identity. ResNet's main contribution was not depth itself, but an architecture that made training deep networks stable. Before ResNet, plain networks with more than about 20 layers often got worse as they got deeper, and the cause was the optimisation, not overfitting. With ResNet, networks with more than 1000 layers could be trained.

Most modern architectures use residual connections: transformers, U-Nets, EfficientNet, ConvNeXt. If you build something deeper than about 10 layers, add residual connections.


11. The Bias-Variance Trade-off, DL Edition

Classical statistical learning: too few parameters → high bias (underfitting), too many → high variance (overfitting). The optimal model is in the middle.

Deep learning does not follow this rule. Modern networks are often heavily over-parameterised: they have far more parameters than training examples, and they still generalise. The classical theory predicts severe overfitting, but it usually doesn't happen.

Double descent (Belkin et al., 2019): plot test error against model size. As you increase model size, the error first falls and then rises (the classical U-curve). It peaks at the interpolation threshold (where the model can exactly fit the training set), and then falls again as the model gets even bigger. Most modern deep learning models are on the right side of this curve.

One explanation is implicit regularisation: in the over-parameterised case, SGD tends to find "simple" solutions. Weight decay and other techniques add explicit regularisation. The classical bias-variance picture is not wrong, but it is incomplete.

In practice: don't be afraid to use a model that "should" overfit. With proper regularisation, a bigger model often generalises better. Check this on your validation set.


Common Mistakes

1. Forgetting normalisation between layers

A 12-layer MLP with ReLU and no BatchNorm/LayerNorm will train, but slowly and with high variance across seeds. Add a normalisation layer after every Dense/Conv layer (or use a block design that includes normalisation). It usually makes training faster and more stable, and it adds some regularisation.

2. Same LR for all layers when fine-tuning

When you fine-tune a pretrained model on a new task, the new head needs a normal LR, but the backbone needs a very small LR (its pretrained features are mostly correct already). Use layer-wise LR decay, or use a small LR (1e-5) for the backbone and 1e-3 for the head only. One LR for everything either destroys the pretrained features or trains the head too little.

3. Treating losses as drop-in replacements

Changing from MSE to cross-entropy is more than a formula change: the loss surface is different. Cross-entropy with logits is much steeper near wrong predictions, so a learning rate that worked for MSE may be too high for CE. Always tune the LR again when you change the loss.

4. Init mismatch

Dense(128, activation="relu") with default Glorot init will train, but slower than Dense(128, activation="relu", kernel_initializer="he_normal"). With deep networks, the difference matters. Sometimes it decides whether the model converges at all.

5. Ignoring the warmup

Train a transformer or large vision model with constant LR from step 0 and you'll see early loss spikes or NaNs. Adam's variance estimates are unreliable for the first few hundred steps. A 1000-step linear warmup costs almost nothing and prevents this problem.


🎯 Your Turn: A 2-Layer MLP Forward Pass in Pure NumPy

Implement a forward pass for a 2-layer classifier: input → Dense(hidden, ReLU) → Dense(num_classes, softmax). No Keras and no autograd: only NumPy and the equations from Section 1.

Requirements:

  • Accept a batch X of shape (batch, in_features).
  • Weights are passed as a tuple (W1, b1, W2, b2).
- W1 shape: (hidden, in_features), b1 shape: (hidden,) - W2 shape: (num_classes, hidden), b2 shape: (num_classes,)
  • Apply ReLU on the hidden layer pre-activations.
  • Apply numerically stable softmax on the output logits.
  • Return a probability matrix of shape (batch, num_classes), rows summing to 1.
python
import numpy as np

def forward(X, params):
    W1, b1, W2, b2 = params
    # TODO 1: compute hidden pre-activations  z1 = X @ W1.T + b1
    # TODO 2: apply ReLU                        a1 = max(0, z1)
    # TODO 3: compute output logits             z2 = a1 @ W2.T + b2
    # TODO 4: numerically stable softmax along axis=1
    ...

# Sanity check
np.random.seed(0)
X  = np.random.randn(4, 3)
W1 = np.random.randn(8, 3) * np.sqrt(2/3)   # He init
b1 = np.zeros(8)
W2 = np.random.randn(2, 8) * np.sqrt(2/8)
b2 = np.zeros(2)

probs = forward(X, (W1, b1, W2, b2))
print(probs.shape)              # (4, 2)
print(probs.sum(axis=1))        # [1. 1. 1. 1.]
Hint 1: Broadcasting the bias X @ W1.T has shape (batch, hidden). Adding b1 of shape (hidden,) broadcasts correctly: NumPy treats b1 as (1, hidden). Don't reshape or transpose b1. Add it directly.
Hint 2: Numerically stable softmax A naive softmax exp(z) / exp(z).sum() overflows for large logits (e.g. exp(1000) = inf). Subtract the per-row max first: z - z.max(axis=1, keepdims=True). The result is mathematically the same (softmax doesn't change when you subtract a constant), and it is numerically safe. The keepdims=True matters for broadcasting.
Show full solution
python
import numpy as np

def forward(X, params):
    """Forward pass: Dense(ReLU), then Dense(softmax)."""
    W1, b1, W2, b2 = params

    # Hidden layer
    z1 = X @ W1.T + b1              # (batch, hidden)
    a1 = np.maximum(0, z1)          # ReLU, element-wise

    # Output layer
    z2 = a1 @ W2.T + b2             # (batch, num_classes)

    # Numerically stable softmax (subtract per-row max)
    z2_shift = z2 - z2.max(axis=1, keepdims=True)
    exp_z2   = np.exp(z2_shift)
    probs    = exp_z2 / exp_z2.sum(axis=1, keepdims=True)

    return probs


np.random.seed(0)
X  = np.random.randn(4, 3)
W1 = np.random.randn(8, 3) * np.sqrt(2/3)
b1 = np.zeros(8)
W2 = np.random.randn(2, 8) * np.sqrt(2/8)
b2 = np.zeros(2)

probs = forward(X, (W1, b1, W2, b2))
print(probs.shape)                  # (4, 2)
print(probs.sum(axis=1))            # [1. 1. 1. 1.]   (each row sums to 1)
print(np.allclose(probs.sum(axis=1), 1.0))   # True

Two design choices worth noting:

  • W1.T convention: storing weights as (out, in) matches PyTorch's nn.Linear layout, where each row is one output unit's weight vector (Keras stores Dense.kernel the other way, as (in, out)). Because of that convention, you need the transpose in the matmul. Some texts use (in, out) and skip the transpose. Pick one and use it consistently.
  • keepdims=True in the softmax max and sum: without it, the reduction removes the axis, and broadcasting lines up the wrong dimensions without any error. This bug is hard to see, so always pass keepdims=True when you will broadcast the result back.

To make this trainable, you'd add a backward pass: cross-entropy gradient at the output ($\hat{y} - y$ for one-hot labels, because softmax with cross-entropy has this simple gradient), backprop through W2, through ReLU's $\{0,1\}$ mask, through W1. Each step is a transpose-matmul plus an element-wise operation. A complete NumPy MLP trainer is about 40 lines, and it is a useful exercise to write once.


What You Learned

  • The MLP is layered $\mathbf{a}^{(\ell)} = \sigma(W^{(\ell)} \mathbf{a}^{(\ell-1)} + \mathbf{b}^{(\ell)})$. Without a non-linearity, extra layers add nothing.
  • Activations: ReLU is the default. Use GELU/Swish for transformers and modern CNNs, and sigmoid/softmax only at the output. The dying ReLU problem is real, and Leaky ReLU is a cheap fix.
  • Initialisation must match the activation: He for ReLU-family, Xavier/Glorot for tanh/sigmoid. Mismatch causes vanishing/exploding signal.
  • Backprop is the chain rule applied along the computation graph, with one multiplication per edge. Autograd automates it, but the math doesn't change.
  • Loss functions define the shape of the problem. Cross-entropy for classification, Huber for robust regression, focal for severe class imbalance.
  • Adam/AdamW is the default optimiser. SGD+momentum still wins on vision sometimes. Warmup + cosine decay is the modern schedule.
  • BatchNorm probably works by smoothing the loss surface (not by reducing "covariate shift", as first claimed). LayerNorm for transformers, GroupNorm for small-batch CV.
  • Residual connections are what made deep networks trainable. Use them in networks deeper than about 10 layers. Modern designs almost always include them.
  • Double descent means over-parameterisation can help generalisation. The classical bias-variance picture is incomplete for deep learning.

Next: TensorFlow & Keras: Production Engine: the framework that turns these building blocks into trained, deployed and scaled models.

Found this useful? Share it with a friend who’s learning Python.