Step 2 — Wire a Neural Network by Hand

Released with Lecture 2 · ~3 hours · Starts from: your Step 1 notebook, though this construction is self-contained.

In what follows, there are blocks of code provided for you. It’s important that you can read the code. If you are not proficient with Python, then I propose the following workflow. For each block of code, keywords are provided for the Python syntax which is being used, so you can look up syntax and functions you don’t recognise. Beneath code that introduces a new idea, you will find a Python background foldout with a small runnable example. First experiment with the code and print intermediate values until you believe you understand it. Then point AI toward the page and task number, explain precisely what you think each line does, and ask for corrections. In this way you will quickly become a code reader (if not a code writer). — Kate

Goal

Construct a genuine one-hidden-layer neural network without autograd and without training. Four ReLU units will build the nonlinear feature

|x_1|+|x_2|,

and an affine output layer plus softmax will classify points as inside or outside a diamond. You will first write the units separately, then assemble the same computation into matrices, inspect every intermediate value, and change the geometry yourself.

The point of this step is wiring. A neural network diagram is not a metaphor: every arrow is one matrix entry, every node has an affine preactivation and an activation, and the dimensions determine which connections are possible.

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. It is short, and it gives you an answer against which every tensor can be checked.

Tasks

2.1 Write the hidden units separately

Begin with five points chosen to exercise both signs, both classes, and the boundary:

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:

the hidden rows sum to X.abs().sum(dim=1);
the inside probabilities equal torch.sigmoid(2 * gamma * (1 - X.abs().sum(dim=1))).

These checks connect the unit-by-unit construction, the matrix construction, and the closed formula from your paper derivation.

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))
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(label='p(inside | x)')
plt.show()

The black $p=1/2$ contour should be a diamond with vertices at $(\pm1,0)$ and $(0,\pm1)$ . Label the four straight pieces with the corresponding equations, such as $x_1+x_2=1$ in the first quadrant.

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?

This is an experiment about logits, not a training sweep: no parameter is being learned.

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

This is the deliverable that shows you understand the wiring rather than only having copied the original matrices.

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.

If you’re stuck

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.