Neural networks — the forward pass with matrices

A network with input x (d numbers), one hidden layer of h neurons and k outputs:

z₁ = W₁x + b₁ (W₁ is h×d, b₁ is h) a₁ = ReLU(z₁) z₂ = W₂a₁ + b₂ (W₂ is k×h, b₂ is k) ŷ = softmax(z₂) (for classification: probabilities that sum to 1)

Parameters: h·d + h + k·h + k. For d = 784, h = 128, k = 10: 100 480 + 1 290 = 101 770.

ReLU(z) = max(0, z). Softmax(z)ᵢ = e^{zᵢ} / Σⱼ e^{zⱼ}.

import numpy as np

def relu(z): return np.maximum(0, z)
def softmax(z):
    e = np.exp(z - z.max())
    return e / e.sum()

rng = np.random.default_rng(0)
W1, b1 = rng.normal(0, 0.1, (128, 784)), np.zeros(128)
W2, b2 = rng.normal(0, 0.1, (10, 128)),  np.zeros(10)

def forward(x):
    a1 = relu(W1 @ x + b1)
    return softmax(W2 @ a1 + b2)

x = rng.random(784)
print(forward(x).sum())   # 1.0

@ is matrix multiplication. With a batch of N examples as an N×784 matrix it becomes X @ W1.T + b1 — the same thing, all the examples at once.