Step 2 — Wire a Neural Network by Hand

Released with Lecture 2 · Starts from: your Step 1 notebook, though this construction is self-contained · Solution: notebook · read online

Goal

Construct a one-hidden-layer neural network by hand, with four hidden ReLU units. Consider the diamond

|x_1|+|x_2| = 1.

The network will classify points as inside or outside this diamond.

On paper first

Let $\mathbf{x}=(x_1,x_2)^\top\in\mathbb R^2$ .

Verify
$\operatorname{ReLU}(t)+\operatorname{ReLU}(-t)=|t|.$
Use it to design four hidden units whose activations sum to $r=|x_1|+|x_2|$ .
Write those four units in the matrix form
$a_1=W_1\mathbf{x}+b_1, \qquad h=\operatorname{ReLU}(a_1),$
including the dimensions of $W_1$ , $b_1$ , $a_1$ , and $h$ .
We want two output logits, ordered as outside and inside:
$z_{\mathrm{out}}=\gamma(r-1), \qquad z_{\mathrm{in}}=\gamma(1-r), \qquad \gamma>0.$
Find $W_2\in\mathbb R^{2\times4}$ and $b_2\in\mathbb R^2$ such that $z=W_2h+b_2$ .
Use the two-class identity from Lecture 2, Aside 3.1a to show
$p(\mathrm{in}\mid\mathbf{x}) = \sigma\bigl(2\gamma(1-|x_1|-|x_2|)\bigr).$
Deduce the decision boundary and explain why changing $\gamma$ cannot move it.

Do this derivation before writing code.

Tasks

2.1 Write the hidden units separately

Begin with five points:

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

Each row of $X$ is one feature vector. Compute the four hidden activations one at a time:

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

For each of the four columns, say in words which input half-plane makes that unit active. Then trace the row for $(-0.4,-0.2)$ by hand before looking at the printed result.

Python background: tensor slices, torch.stack, and dim

For a two-dimensional array, A[:, 0] means “every row, column 0.” Stacking several one-dimensional arrays with dim=1 makes them columns of a new matrix. NumPy has the same operations, so the browser example below uses NumPy.

import numpy as np
A = np.array([[10, 11],
              [20, 21],
              [30, 31]])
first_column = A[:, 0]
second_column = A[:, 1]
print(first_column)
print(np.stack([first_column, second_column], axis=1))
print(np.stack([first_column, second_column], axis=0))

Compare the last two shapes. In PyTorch the argument is named dim; in NumPy it is named axis.

2.2 Assemble the first layer as a matrix

Now encode exactly the same four units in one weight matrix and one bias vector:

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)

The lecture writes one input as a column and uses $W_1\mathbf{x}$ . Code conventionally stores a batch with one sample per row, so the same multiplication is X @ W1.T. Check the dimensions on paper:

(5\times2)(2\times4)+(4) = (5\times4).

The bias vector is added to every row. This reuse is called broadcasting.

Python background: matrix multiplication, transpose, and broadcasting

The operator @ performs matrix multiplication, while .T transposes a two-dimensional array. Adding a length- $m$ vector to an $n\times m$ matrix adds that vector to every row.

import numpy as np
X = np.array([[1., 2.],
              [3., 4.],
              [5., 6.]])
W = np.array([[10.,  0.],
              [ 0., 10.]])
b = np.array([1., -1.])
print(X @ W.T)
print(X @ W.T + b)
print((X @ W.T + b).shape)

Matrix dimensions are not decoration: the inner dimensions must agree, and the remaining dimensions give the output shape.

2.3 Wire the output layer and softmax

Use $\gamma=4$ first:

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))

Explain every dimension in

(5\times4)(4\times2)+(2)=(5\times2).

Then explain why subtracting the largest logit in each row changes no softmax probability. Your row sums should equal $1$ up to floating-point roundoff.

The output columns are outside and inside, in that order. On the boundary the logits tie, so both probabilities are $1/2$ . If you use argmax, PyTorch breaks that tie by returning the first maximizing index; that software convention is separate from the mathematical classifier.

2.4 Package one forward pass

Turn the two matrix layers into a function. Return the intermediate values, not only the final probabilities: this step is about seeing the computation.

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)

Add two assertions of your own:

verify that the hidden layer adds to $|x_1| + |x_2|$ ;
the inside probabilities for the points each equal $\sigma(2 \gamma (1 - |x_1| - |x_2|))$ .

2.5 Draw the decision surface

Evaluate the network on a fine grid, then plot its inside probability:

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=1/2$ contour should be a diamond with vertices at $(\pm1,0)$ and $(0,\pm1)$ .

Python background: grids, reshape, and contour plots

A grid begins with two one-dimensional coordinate arrays. meshgrid repeats them into matrices of all coordinate pairs. Flattening those matrices and stacking their entries produces the usual “one point per row” batch. After evaluating the batch, reshape restores the grid so a plotting function knows where each value belongs.

import numpy as np
axis = np.linspace(-1, 1, 3)
gx, gy = np.meshgrid(axis, axis, indexing='xy')
points = np.stack([gx.reshape(-1), gy.reshape(-1)], axis=1)
print(gx)
print(gy)
print(points)
print(points.shape)

A filled contour plot colors regions according to a scalar value. A single contour at level $1/2$ draws the classifier’s decision boundary.

2.6 Change confidence without changing geometry

Run the same grid with

\gamma\in\{0.25,1,4,20\}.

Put the four probability plots side by side with the same color scale. Answer:

Which set of points has probability exactly $1/2$ in every plot?
What happens to probabilities away from the boundary as $\gamma$ grows?
What happens as $\gamma$ approaches zero?
Why is changing $\gamma$ equivalent to changing softmax temperature?

2.7 Move and stretch the diamond

Choose a center $(a,b)$ and positive radii $r_1,r_2$ . Modify the first layer so that its hidden activations sum to

\frac{|x_1-a|}{r_1}+\frac{|x_2-b|}{r_2}.

Keep the output layer’s rule “inside when the sum is below $1$ .” Derive your new $W_1$ and $b_1$ on paper, implement them, and plot the result. The boundary should have vertices

(a\pm r_1,b), \qquad (a,b\pm r_2).

Checkpoints

h_units and the matrix-layer activation h agree exactly up to floating-point arithmetic.
The five hidden rows are
$\begin{pmatrix} 0&0&0&0\\ 0.5&0&0.25&0\\ 1&0&0&0\\ 0.8&0&0.5&0\\ 0&0.4&0&0.2 \end{pmatrix}.$
At $\gamma=4$ , the inside probabilities for the first four points are approximately $0.9997$ , $0.8808$ , $0.5000$ , and $0.0832$ . The fifth is approximately $0.9608$ .
Every probability row sums to $1$ .
The $p=1/2$ contour is the diamond $|x_1|+|x_2|=1$ for every positive $\gamma$ .
Your moved-and-stretched boundary has the four vertices predicted on paper.

Troubleshooting

X @ W1 has a shape error: the lecture uses column vectors, while the batch stores samples as rows. Use X @ W1.T.
The hidden sum is $x_1+x_2$ rather than $|x_1|+|x_2|$ : inspect the two negative-coordinate units and their signs.
Inside and outside are reversed: inspect the order of the two rows of $W_2$ and the two entries of $b_2$ .
Probabilities do not sum to $1$ : normalize across dim=1, the class dimension, and retain it with keepdim=True.
The plot is transposed or rotated: use the same indexing=‘xy’ convention in meshgrid as in the supplied code.

Going further (optional)

A square needs another composition. Use $\max(u,v)=\operatorname{ReLU}(u-v)+v$ to build $\max(|x_1|,|x_2|)$ , then classify $\max(|x_1|,|x_2|)<1$ . Draw the extra layer and give every matrix dimension.
Rotate the boundary. Replace the coordinate directions in $W_1$ by two nonparallel direction vectors. Predict the four edge normals before plotting.
Make a dataset. Sample random points in $[-1.5,1.5]^2$ , label them from the exact inequality, and measure classification accuracy away from the boundary. Explain why this measures your implementation, not generalization.
Count the wiring. Count all scalar entries in the two weight matrices and two bias vectors. Then count how many are nonzero and explain what each nonzero group does.

Catch-up

This step is self-contained. If you are joining now, skim Lecture 2, Section 4 and Section 7, then run the supplied blocks in a fresh notebook and do Task 2.7 in full. You do not need the Shakespeare corpus again until Step 3.