Step 2 — Wire a Neural Network by Hand

Solution notebook. A one-hidden-layer ReLU network, built unit by unit and then as two matrix layers, that classifies points of the plane as inside or outside the diamond $|x_1|+|x_2|=1$ — exactly, with weights written down on paper rather than learned. No corpus, no training, no autograd: this step is about seeing every number in a forward pass. It runs in a few seconds.

The construction is Lecture 2, Example 4.6 and Theorem 7.1; the derivation below is the "On paper first" section of the task, written out.

On paper first

1. Four units whose activations sum to $r=|x_1|+|x_2|$. For any real $t$, exactly one of $t$ and $-t$ is positive (or both are zero), so $\operatorname{ReLU}(t)+\operatorname{ReLU}(-t)=\max(t,0)+\max(-t,0)=|t|$. Apply it to each coordinate: the four hidden units

$$ h=\bigl(\operatorname{ReLU}(x_1),\ \operatorname{ReLU}(-x_1),\ \operatorname{ReLU}(x_2),\ \operatorname{ReLU}(-x_2)\bigr)^\top $$

have $h_1+h_2=|x_1|$ and $h_3+h_4=|x_2|$, so $\mathbf 1^\top h=r$.

2. As a matrix layer. Each unit is $\operatorname{ReLU}$ of an affine function of $\mathbf x$, so $a_1=W_1\mathbf x+b_1$, $h=\operatorname{ReLU}(a_1)$ with

$$ W_1=\begin{pmatrix}1&0\\-1&0\\0&1\\0&-1\end{pmatrix}\in\mathbb R^{4\times2}, \qquad b_1=\mathbf 0\in\mathbb R^4, \qquad a_1,h\in\mathbb R^4 . $$

3. The output layer. We want $z_{\mathrm{out}}=\gamma(r-1)$ and $z_{\mathrm{in}}=\gamma(1-r)$. Since $r=\mathbf 1^\top h$,

$$ W_2=\gamma\begin{pmatrix}1&1&1&1\\-1&-1&-1&-1\end{pmatrix}\in\mathbb R^{2\times4}, \qquad b_2=\gamma\begin{pmatrix}-1\\1\end{pmatrix}\in\mathbb R^2, \qquad z=W_2h+b_2 . $$

4. The probability and the boundary. For two classes, softmax reduces to a sigmoid of the logit difference (Lecture 2, Aside 3.1a): $p(\mathrm{in}\mid\mathbf x)=\sigma(z_{\mathrm{in}}-z_{\mathrm{out}}) =\sigma\bigl(2\gamma(1-|x_1|-|x_2|)\bigr)$. Since $\sigma$ is increasing with $\sigma(0)=\tfrac12$, the classifier says inside exactly when $|x_1|+|x_2|<1$: the boundary is the diamond with vertices $(\pm1,0)$, $(0,\pm1)$. The sign of $1-r$ does not depend on $\gamma$, so $\gamma$ rescales the logits — how confident the network is — without moving the boundary at all.

2.1 Write the hidden units separately

Five test points, one per row. Each column of h_units is one hidden unit evaluated on all five points at once.

import torch

X = torch.tensor([
    [ 0.0,  0.0],   # inside
    [ 0.5,  0.25],  # inside
    [ 1.0,  0.0],   # boundary
    [ 0.8,  0.5],   # outside
    [-0.4, -0.2],   # inside, negative coordinates
])
print(X.shape)      # (5, 2): five points, two features per point

h_units = torch.stack([
    torch.relu( X[:, 0]),
    torch.relu(-X[:, 0]),
    torch.relu( X[:, 1]),
    torch.relu(-X[:, 1]),
], dim=1)

print(h_units)
print(h_units.shape)                 # (5, 4)
r_units = h_units.sum(dim=1)
print(r_units)                       # |x1| + |x2| for each point

The four columns are active on the four half-planes $x_1>0$, $x_1<0$, $x_2>0$, $x_2<0$ respectively — each unit measures how far the point sits into "its" half-plane and is silent elsewhere. For $(-0.4,-0.2)$: unit 1 sees $-0.4$ and outputs $0$; unit 2 sees $+0.4$ and outputs $0.4$; unit 3 sees $-0.2$, outputs $0$; unit 4 outputs $0.2$. Row: $(0,\,0.4,\,0,\,0.2)$, sum $0.6=|{-0.4}|+|{-0.2}|$. The five rows are the checkpoint matrix.

2.2 Assemble the first layer as a matrix

The same four units as one weight matrix and one bias vector. The lecture writes one input as a column and computes $W_1\mathbf x$; code stores a batch with one sample per row, so the same product is X @ W1.T: $(5\times2)(2\times4)+(4)=(5\times4)$, the bias added to every row by broadcasting.

W1 = torch.tensor([
    [ 1.0,  0.0],
    [-1.0,  0.0],
    [ 0.0,  1.0],
    [ 0.0, -1.0],
])
b1 = torch.zeros(4)

a1 = X @ W1.T + b1
h = torch.relu(a1)

print(a1.shape, h.shape)             # both (5, 4)
assert torch.allclose(h, h_units)
print('matrix layer == four separate units ✓')

2.3 Wire the output layer and softmax

$(5\times4)(4\times2)+(2)=(5\times2)$: two logits per point, columns ordered outside, inside. Subtracting each row's largest logit before exponentiating changes nothing — softmax is invariant under adding a constant to every logit in a row, since the constant factors out of numerator and denominator alike — but it keeps exp from overflowing when $\gamma$ is large.

gamma = 4.0
W2 = gamma * torch.tensor([
    [ 1.0,  1.0,  1.0,  1.0],   # outside logit
    [-1.0, -1.0, -1.0, -1.0],   # inside logit
])
b2 = gamma * torch.tensor([-1.0, 1.0])

logits = h @ W2.T + b2            # (5, 2)
shifted = logits - logits.max(dim=1, keepdim=True).values
weights = shifted.exp()
probs = weights / weights.sum(dim=1, keepdim=True)

print(logits)
print(probs)
print(probs.sum(dim=1))

Inside probabilities at $\gamma=4$: 0.9997, 0.8808, 0.5000, 0.0832, 0.9608 — i.e. $\sigma(8)$, $\sigma(2)$, $\sigma(0)$, $\sigma(-2.4)$, $\sigma(3.2)$, since $2\gamma(1-r)=8(1-r)$ and $r=0,\,0.75,\,1,\,1.3,\,0.6$. On the boundary point $(1,0)$ the logits tie and both probabilities are exactly $\tfrac12$; argmax would return index 0 (outside) there, a tie-breaking convention of the software, not a property of the classifier. Every row sums to 1.

2.4 Package one forward pass

The two layers as a function that returns every intermediate value. Then the two assertions the task asks for: the hidden layer sums to $|x_1|+|x_2|$, and the inside probability is $\sigma(2\gamma(1-|x_1|-|x_2|))$ — Theorem 7.1, checked numerically.

def diamond_net(X, gamma=4.0):
    W1 = torch.tensor([
        [ 1.0,  0.0],
        [-1.0,  0.0],
        [ 0.0,  1.0],
        [ 0.0, -1.0],
    ], dtype=X.dtype)
    b1 = torch.zeros(4, dtype=X.dtype)

    W2 = gamma * torch.tensor([
        [ 1.0,  1.0,  1.0,  1.0],
        [-1.0, -1.0, -1.0, -1.0],
    ], dtype=X.dtype)
    b2 = gamma * torch.tensor([-1.0, 1.0], dtype=X.dtype)

    a1 = X @ W1.T + b1
    h = torch.relu(a1)
    logits = h @ W2.T + b2

    shifted = logits - logits.max(dim=1, keepdim=True).values
    weights = shifted.exp()
    probs = weights / weights.sum(dim=1, keepdim=True)
    return a1, h, logits, probs

a1, h, logits, probs = diamond_net(X)
assert a1.shape == (5, 4)
assert h.shape == (5, 4)
assert logits.shape == (5, 2)
assert probs.shape == (5, 2)

# the two assertions of our own
r = X.abs().sum(dim=1)                                   # |x1| + |x2|
assert torch.allclose(h.sum(dim=1), r)
assert torch.allclose(probs[:, 1], torch.sigmoid(2 * 4.0 * (1 - r)))
print('hidden sum = |x1|+|x2| ✓   p(in) = sigmoid(2γ(1-|x1|-|x2|)) ✓')

2.5 Draw the decision surface

Evaluate the network on a $201\times201$ grid and plot the inside probability. meshgrid turns two coordinate axes into all coordinate pairs; flattening and stacking makes the usual one-point-per-row batch; reshape puts the answers back on the grid for the plotting routine.

import matplotlib.pyplot as plt

axis = torch.linspace(-1.6, 1.6, 201)
gx, gy = torch.meshgrid(axis, axis, indexing='xy')
grid = torch.stack([gx.reshape(-1), gy.reshape(-1)], dim=1)

_, _, _, grid_probs = diamond_net(grid, gamma=4.0)
p_inside = grid_probs[:, 1].reshape(gx.shape)

plt.figure(figsize=(6, 5))
contour_plot = plt.contourf(gx.numpy(), gy.numpy(), p_inside.numpy(),
                            levels=30, cmap='Purples')
plt.contour(gx.numpy(), gy.numpy(), p_inside.numpy(),
            levels=[0.5], colors='black', linewidths=2)
plt.scatter(X[:, 0], X[:, 1], c='red', s=30)
plt.xlabel('x1')
plt.ylabel('x2')
plt.axis('equal')
plt.colorbar(contour_plot, label='p(inside | x)')
plt.show()

The black $p=\tfrac12$ contour is the diamond with vertices $(\pm1,0)$, $(0,\pm1)$. A single logistic-regression unit can only draw one line in this plane; four hidden ReLUs measure distances from the two axes, and their sum bends one line into four.

2.6 Change confidence without changing geometry

The same grid at $\gamma\in\{0.25,1,4,20\}$, one colour scale for all four (vmin=0, vmax=1, so the panels are comparable).

gammas = [0.25, 1.0, 4.0, 20.0]
fig, axes = plt.subplots(1, 4, figsize=(16, 4), sharex=True, sharey=True)
for ax, g in zip(axes, gammas):
    _, _, _, gp = diamond_net(grid, gamma=g)
    pi = gp[:, 1].reshape(gx.shape).numpy()
    im = ax.contourf(gx.numpy(), gy.numpy(), pi, levels=torch.linspace(0, 1, 31).numpy(),
                     cmap='Purples', vmin=0, vmax=1)
    ax.contour(gx.numpy(), gy.numpy(), pi, levels=[0.5], colors='black', linewidths=2)
    ax.set_title(f'γ = {g}')
    ax.set_aspect('equal')
fig.colorbar(im, ax=axes, label='p(inside | x)', shrink=0.8)
plt.show()

# the boundary does not move: p = 1/2 on the diamond's vertices for every γ
vertices = torch.tensor([[1., 0.], [0., 1.], [-1., 0.], [0., -1.]])
for g in gammas:
    _, _, _, vp = diamond_net(vertices, gamma=g)
    assert torch.allclose(vp[:, 1], torch.full((4,), 0.5))
print('p(inside) = 1/2 on all four vertices, for every γ ✓')

1. The diamond itself — every point with $|x_1|+|x_2|=1$ — has probability exactly $\tfrac12$ in all four plots, because there $z_{\mathrm{in}}=z_{\mathrm{out}}=0$ whatever $\gamma$ is.
2. As $\gamma$ grows, probabilities away from the boundary saturate toward $0$ and $1$: the purple band of uncertainty narrows, and at $\gamma=20$ the plot is essentially a two-colour picture of the diamond's indicator function.
3. As $\gamma\to0$, both logits go to $0$ and every point tends to $p=\tfrac12$: the network is still correct (the sign of $z_{\mathrm{in}}-z_{\mathrm{out}}$ never changes) but maximally unconfident, and the whole plane fades to the same mid-purple.
4. Softmax at temperature $T$ is $\operatorname{softmax}(z/T)$. Here every logit is proportional to $\gamma$, so $\operatorname{softmax}(z(\gamma))=\operatorname{softmax}(z(1)/T)$ with $T=1/\gamma$: changing $\gamma$ is changing the temperature. Temperature rescales a distribution's sharpness and leaves its argmax — the geometry — untouched. (Lecture 7 returns to this when we sample.)

2.7 Move and stretch the diamond

For a centre $(a,b)$ and radii $r_1,r_2$ we want the hidden activations to sum to $\dfrac{|x_1-a|}{r_1}+\dfrac{|x_2-b|}{r_2}$. Since $|x_1-a|/r_1=\operatorname{ReLU}\!\bigl(\tfrac{x_1-a}{r_1}\bigr)+\operatorname{ReLU}\!\bigl(\tfrac{a-x_1}{r_1}\bigr)$, the four units are affine in $\mathbf x$ with

$$ W_1=\begin{pmatrix}1/r_1&0\\-1/r_1&0\\0&1/r_2\\0&-1/r_2\end{pmatrix}, \qquad b_1=\begin{pmatrix}-a/r_1\\ a/r_1\\ -b/r_2\\ b/r_2\end{pmatrix}, $$

and the output layer is unchanged: inside when the sum is below $1$. The boundary $\frac{|x_1-a|}{r_1}+\frac{|x_2-b|}{r_2}=1$ has vertices $(a\pm r_1,b)$ and $(a,b\pm r_2)$.

def stretched_diamond_net(X, a=0.5, b=-0.25, r1=1.2, r2=0.6, gamma=4.0):
    W1 = torch.tensor([
        [ 1/r1,   0.0],
        [-1/r1,   0.0],
        [  0.0,  1/r2],
        [  0.0, -1/r2],
    ], dtype=X.dtype)
    b1 = torch.tensor([-a/r1, a/r1, -b/r2, b/r2], dtype=X.dtype)
    W2 = gamma * torch.tensor([[1., 1., 1., 1.], [-1., -1., -1., -1.]], dtype=X.dtype)
    b2 = gamma * torch.tensor([-1.0, 1.0], dtype=X.dtype)

    h = torch.relu(X @ W1.T + b1)
    logits = h @ W2.T + b2
    return torch.softmax(logits, dim=1)        # same as the by-hand softmax above

a, b, r1, r2 = 0.5, -0.25, 1.2, 0.6
axis2 = torch.linspace(-1.5, 2.5, 201)
gx2, gy2 = torch.meshgrid(axis2, axis2, indexing='xy')
grid2 = torch.stack([gx2.reshape(-1), gy2.reshape(-1)], dim=1)
p2 = stretched_diamond_net(grid2, a, b, r1, r2)[:, 1].reshape(gx2.shape)

plt.figure(figsize=(6, 5))
plt.contourf(gx2.numpy(), gy2.numpy(), p2.numpy(), levels=30, cmap='Purples')
plt.contour(gx2.numpy(), gy2.numpy(), p2.numpy(), levels=[0.5], colors='black', linewidths=2)
predicted = torch.tensor([[a + r1, b], [a - r1, b], [a, b + r2], [a, b - r2]])
plt.scatter(predicted[:, 0], predicted[:, 1], c='red', s=40, zorder=3, label='predicted vertices')
plt.legend(); plt.axis('equal'); plt.xlabel('x1'); plt.ylabel('x2')
plt.show()

assert torch.allclose(stretched_diamond_net(predicted, a, b, r1, r2)[:, 1], torch.full((4,), 0.5))
print('p(inside) = 1/2 exactly at the four predicted vertices ✓')

Going further — a square needs another composition

$\max(u,v)=\operatorname{ReLU}(u-v)+v$, so $\max(|x_1|,|x_2|)$ needs one more ReLU layer after the absolute values are formed — a composition, not wider first layer. With $h$ the four units above, $u=h_1+h_2$ and $v=h_3+h_4$:

$$ a_2=W_{2}'\,h=\begin{pmatrix}1&1&-1&-1\\0&0&1&1\end{pmatrix}h =\begin{pmatrix}u-v\\ v\end{pmatrix}, \qquad h_2=\operatorname{ReLU}(a_2)=\begin{pmatrix}\operatorname{ReLU}(u-v)\\ v\end{pmatrix} $$

(the second unit passes $v\ge0$ through unchanged), and $\mathbf 1^\top h_2=\max(u,v)$. The output layer is the same rule as before on this sum. Dimensions: $W_1\in\mathbb R^{4\times2}$, $W_2'\in\mathbb R^{2\times4}$, $W_3\in\mathbb R^{2\times2}$.

def square_net(X, gamma=4.0):
    h = torch.relu(X @ W1.T + b1)                                 # (N, 4): |x1|, |x2| split into halves
    W2p = torch.tensor([[1., 1., -1., -1.], [0., 0., 1., 1.]])    # (2, 4)
    h2 = torch.relu(h @ W2p.T)                                    # (N, 2): (ReLU(u - v), v)
    W3 = gamma * torch.tensor([[1., 1.], [-1., -1.]])             # (2, 2)
    b3 = gamma * torch.tensor([-1.0, 1.0])
    return torch.softmax(h2 @ W3.T + b3, dim=1)

m = torch.maximum(X[:, 0].abs(), X[:, 1].abs())
assert torch.allclose(square_net(X)[:, 1], torch.sigmoid(2 * 4.0 * (1 - m)))

p_sq = square_net(grid)[:, 1].reshape(gx.shape)
plt.figure(figsize=(5, 5))
plt.contourf(gx.numpy(), gy.numpy(), p_sq.numpy(), levels=30, cmap='Purples')
plt.contour(gx.numpy(), gy.numpy(), p_sq.numpy(), levels=[0.5], colors='black', linewidths=2)
plt.axis('equal'); plt.title('max(|x1|, |x2|) < 1: three layers, one square')
plt.show()

Count the wiring. The diamond network has $4\cdot2+4+2\cdot4+2=22$ scalar parameters, of which $4+0+8+2=14$ are nonzero: four in $W_1$ (one per unit, choosing an axis and a sign), eight in $W_2$ (every hidden unit feeds both logits, with opposite signs), two in $b_2$ (the threshold $r=1$). Every one of them has a job you can name — which is the point of building a network by hand before letting gradient descent fill the numbers in.

→ Continue with Step 3: differentiate this very network — first with an autograd engine you write yourself, then with PyTorch's — and then train the first model of the course whose weights nobody wrote down.