Linear Algebra in a Tiny Neural Network

A neural network can look like a string of matrix operations, yet each step changes what the numbers mean. Following one small network from input to output reveals the whole route: matrix shapes determine which maps can connect, biases translate those maps, rank limits independent directions, and an activation changes what the composition can express.

By the end, you will be able to trace each shape and calculate a two-layer network by hand or in NumPy, then explain where its bias, rank, and activation affect the result.

This lesson brings together three earlier tools: solving linear systems without explicitly inverting a matrix, using kernel and image to describe lost and reachable directions, and using bias to translate an affine map. Chapter 2 of Mathematics for Machine Learning supplies the linear algebra behind the synthesis. The network, numbers, code, visual, and assessments below are original.

Follow the dimensions through two affine layers

Here, “two-layer network” means two learned affine maps with a hidden activation between them. We will keep the output as two scores; no probability or class decision is applied. Let the input be a column vector $x\in\mathbb {R}^2$ . Each affine map combines a matrix with a bias vector:

A two-layer network and its shapes

\begin{aligned} z &= W_1x+b_1, & W_1&\in\mathbb{R}^{2\times 2}, & b_1&\in\mathbb{R}^2, & z&\in\mathbb{R}^2,\\ h &= \sigma(z), & h&\in\mathbb{R}^2,\\ y &= W_2h+b_2, & W_2&\in\mathbb{R}^{2\times 2}, & b_2&\in\mathbb{R}^2, & y&\in\mathbb{R}^2. \end{aligned}

The hidden vector $z$ is the first affine map’s preactivation. The function $\sigma$ acts on each of its coordinates; in our example it will be ReLU, $\sigma(t)=\max(0,t)$ . The activation preserves shape, so $h$ is still a two-coordinate vector and can be passed to the second matrix.

The shared dimension between adjacent maps must match: $W_1$ has two columns for the two input features, and two rows for its hidden outputs. Then $W_2$ accepts those two hidden values and returns two output scores. The bias $b_1$ is added once to each hidden coordinate, and $b_2$ once to each output. A matrix with incompatible dimensions cannot perform the stated mapping; this is the same shape check that NumPy’s matrix product applies.

Trace a two-input, two-hidden-unit networkMove either input coordinate or switch the hidden activation. The initial state uses x = (1, 1) and ReLU; all values below are rendered before JavaScript runs.

Input · 2 features
Affine map · 2 hidden values
Hidden activation
Affine map · 2 scores
Output · 2 scores

z = W₁x + b₁

h = σ(z)

y = W₂h + b₂

Change the input and hidden rule

Use the arrow keys or slider to change each input from −2 to 2 in steps of 0.25.

Input x₁1.00Input x₂1.00

Hidden ruleReLUIdentity (no activation)

Preactivation z: (-0.50, 0.50)
Hidden output h: (0.00, 0.50)
Scores y: (1.50, -0.50)
Active hidden units: 1 / 2
Local Jacobian rank: 1

The starting values in the figure use these fixed, hand-chosen parameters:

W_1=\begin{bmatrix}1&-1\\1&1\end{bmatrix},\quad b_1=\begin{bmatrix}-0.5\\-1.5\end{bmatrix},\quad W_2=\begin{bmatrix}1&2\\-1&1\end{bmatrix},\quad b_2=\begin{bmatrix}0.5\\-1\end{bmatrix}.

The example is small enough to follow entry by entry. Its values are chosen to show one hidden unit turning off under ReLU; they are not trained parameters or a claim about a typical model.

Compute one prediction, keeping every intermediate

Take $x=(1,1)^\mathsf{T}$ . The first matrix combines the two input coordinates in different ways. Adding the bias shifts each result:

Forward pass with a ReLU hidden step

\begin{aligned} z &= W_1x+b_1 = \begin{bmatrix}1&-1\\1&1\end{bmatrix} \begin{bmatrix}1\\1\end{bmatrix} +\begin{bmatrix}-0.5\\-1.5\end{bmatrix} = \begin{bmatrix}0\\2\end{bmatrix} +\begin{bmatrix}-0.5\\-1.5\end{bmatrix} = \begin{bmatrix}-0.5\\0.5\end{bmatrix},\\[4pt] h &= \operatorname{ReLU}(z) = \begin{bmatrix}\max(0,-0.5)\\\max(0,0.5)\end{bmatrix} = \begin{bmatrix}0\\0.5\end{bmatrix},\\[4pt] y &= W_2h+b_2 = \begin{bmatrix}1&2\\-1&1\end{bmatrix} \begin{bmatrix}0\\0.5\end{bmatrix} +\begin{bmatrix}0.5\\-1\end{bmatrix} = \begin{bmatrix}1\\0.5\end{bmatrix} +\begin{bmatrix}0.5\\-1\end{bmatrix} = \begin{bmatrix}1.5\\-0.5\end{bmatrix}. \end{aligned}

The shapes at each step are $(2\times2)(2,)+(2,)=(2,)$ , then the activation keeps shape $(2,)$ , then $(2\times2)(2,)+(2,)=(2,)$ . The resulting $y$ is two scores. A separate output rule would be needed to turn them into a label or a probability distribution.

NumPy’s @ operator performs matrix multiplication. The + b operations below use NumPy’s elementwise broadcasting to add a length-two bias to a length-two result. This small implementation prints the actual intermediate vectors:

import numpy as np

W1 = np.array([[1.0, -1.0], [1.0, 1.0]])  # (2, 2)
b1 = np.array([-0.5, -1.5])               # (2,)
W2 = np.array([[1.0, 2.0], [-1.0, 1.0]])  # (2, 2)
b2 = np.array([0.5, -1.0])                # (2,)
x = np.array([1.0, 1.0])                  # (2,)

z = W1 @ x + b1                           # (2,) = (2, 2) @ (2,) + (2,)
h = np.maximum(z, 0.0)                    # (2,), elementwise ReLU
y = W2 @ h + b2                           # (2,) = (2, 2) @ (2,) + (2,)

print('z:', z)  # [-0.5  0.5]
print('h:', h)  # [0.  0.5]
print('y:', y)  # [ 1.5 -0.5]

The same width logic applies to a batch. If $B$ examples are stored as rows, their array has shape $(B,2)$ ; multiply on the right by $W_1^\mathsf{T}$ to produce $(B,2)$ hidden values, and then by $W_2^\mathsf{T}$ to produce $(B,2)$ scores. PyTorch’s Linear layer uses that row-oriented convention: its weight shape is $(\text{out features},\text{in features})$ and it applies $y=xA^\mathsf{T}+b$ . A Linear(2, 2) therefore stores a $(2,2)$ weight and a length-two bias. PyTorch’s ReLU is applied elementwise and preserves the input shape, matching the step used above (Linear docs, ReLU docs, NumPy matmul docs).

Without an activation, two layers are still one affine map

The bias in an affine layer is a translation. If the hidden step is the identity, the two layers compose by substitution:

y=W_2(W_1x+b_1)+b_2 =(W_2W_1)x+(W_2b_1+b_2).

So stacking more affine maps without a nonlinear activation cannot turn the overall input-output rule into a nonlinear one; the matrices multiply into a single matrix, and the biases combine into one offset. For our chosen values, $W_2W_1=\begin{bmatrix}3&1\\0&2\end{bmatrix}$ and $W_2b_1+b_2=(-3,-2)^\mathsf{T}$ . The identity version therefore maps $(1,1)^\mathsf{T}$ to $(1,0)^\mathsf{T}$ .

The rank of a matrix counts its independent output directions. In this identity example both $W_1$ and $W_2$ have rank two, and their product also has rank two. In a network whose hidden width is $r$ , the linear part of an affine composition cannot have rank greater than $r$ , since it factors through only $r$ coordinates. If either matrix has lower rank, the composition’s rank can be lower still. These are upper bounds: a narrow layer is a possible bottleneck, but width alone does not guarantee a particular rank.

ReLU makes the map piecewise affine

ReLU changes a hidden coordinate to zero when its preactivation is negative. For any region where the signs of all preactivations stay fixed, the set of active units is fixed too. Let $D$ be a diagonal matrix with a $1$ for each active ReLU unit and a $0$ for each inactive one. Within that region, the local linear part of the two-layer network is

J=W_2DW_1.

For the worked input, $z=(-0.5,0.5)^\mathsf{T}$ , so $D=\begin{bmatrix}0&0\\0&1\end{bmatrix}$ . Then

J= \begin{bmatrix}1&2\\-1&1\end{bmatrix} \begin{bmatrix}0&0\\0&1\end{bmatrix} \begin{bmatrix}1&-1\\1&1\end{bmatrix} =\begin{bmatrix}2&2\\1&1\end{bmatrix}, \qquad \operatorname{rank}(J)=1.

Near this input, both output scores change with the same input combination $x_1+x_2$ ; the difference direction $(1,-1)$ has no effect while the first unit remains inactive. The matrix $J$ describes local changes inside this activation region, not one global matrix for the entire network. On a ReLU boundary where a preactivation equals zero, the function has a kink and no single region’s Jacobian describes both sides.

The hidden width and weights also constrain what the network can express. In this particular example, the two preactivation-zero lines partition the input plane into regions with different active-unit patterns; each region has an affine rule, and crossing a line can change that rule. More generally, adding a nonlinear activation makes nonlinear functions possible, but does not make every choice of weights expressive. A hidden unit may remain off for the inputs of interest, or the matrices may collapse directions. The matrix shapes, their ranks, the bias offsets, and the activation pattern all contribute to the behavior.

That is why the output values alone are a thin debugging signal. Trace shapes to catch incompatible connections, compare $z$ and $h$ to see which features were gated, and inspect the layer ranks or local Jacobian to identify lost directions. When a separate fitting step asks you to solve $Ax\approx b$ , Lesson 025 explains why a numerical solver avoids the hazards of forming $A^{-1}$ . A network forward pass itself uses matrix products; it does not require inverting its learned weights.

Once the maps and their directions are clear, the next question is how large a vector is. Lesson 039 introduces norms to measure vector magnitude, compare embedding scales, and describe parameter penalties.

Check your understanding

Question 1

For the network in this lesson, state the shapes of x, W₁, b₁, z, h, W₂, b₂, and y. Why can W₂ accept h?