Step 2 — Wire a Network by Hand, Then Train One

Released with Lecture 2 · Starts from: Part A is self-contained; Part B needs your Step 1 ids (or solutions/step-01.ipynb) · Solution: notebook · read online

Goal

Two constructions, in order:

1. Part A — wire. Construct a one-hidden-layer neural network by hand, with four hidden ReLU units, that classifies points as inside or outside the diamond

   $$|x_1|+|x_2| = 1.$$

   Nothing is trained; the point is to see every number in a forward pass.

2. Part B — train. Rebuild Step 1’s bigram model, but as a differentiable parametric family trained by gradient descent — with a gradient you derive by hand, no autograd allowed yet. This is the conceptual pivot of the whole course: from estimating a table to optimizing a function. You will prove, and then watch numerically, that the two coincide.

Part A — Wire a diamond classifier by hand

On paper first

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

1. Verify

   $$\operatorname{ReLU}(t)+\operatorname{ReLU}(-t)=|t|.$$

   Use it to design four hidden units whose activations sum to $r=|x_1|+|x_2|$.

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

3. 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$.

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

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.

Look and ask. Before moving on, interrogate h_units rather than just checking it runs.

1. Count the nonzero entries in each row. What is the largest number of the four units that can be active at one point, and why can no point ever activate three?
2. The map $\mathbf{x}\mapsto h$ is nonlinear, yet nothing is lost: explain how to recover $\mathbf{x}$ from $h$. Then say what is lost when the four activations are summed to $r$. Name two very different points the classifier cannot tell apart.

print((h_units > 0).sum(dim=1))          # tensor([0, 2, 1, 2, 2])
print(h_units[:, 0] - h_units[:, 1])     # recovers x1 for every point
print(h_units[:, 2] - h_units[:, 3])     # recovers x2

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

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:

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

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)

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

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.

Look and ask.

1. The points $(0.7,0)$ and $(0.5,0.5)$ are at essentially the same Euclidean distance from the origin ($0.700$ against $0.707$). Predict their inside probabilities, then run them through the network. What notion of distance is this network measuring?
2. Take the five probabilities you just printed at $\gamma=4$ and predict, before running, what they become at $\gamma=1$ and at $\gamma=20$. Which entries move, which cannot, and which way?

pair = torch.tensor([[0.7, 0.0], [0.5, 0.5]])
r = pair.abs().sum(dim=1)                     # tensor([0.7000, 1.0000])
print(torch.sigmoid(2 * gamma * (1 - r)))     # inside: 0.9168 and 0.5000
# (or, once Task 2.4 is done: diamond_net(pair)[3][:, 1])

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

Look and ask. The picture has more in it than the black line.

1. Before you plot, predict where the $p=0.9$ and $p=0.1$ contours lie at $\gamma=4$ (solve $\sigma(2\gamma(1-r))=0.9$ for $r$). Then add levels=[0.1, 0.5, 0.9] to the plt.contour call and check.
2. The colour bands run parallel to the diamond’s edges instead of forming concentric circles. Why? What happens to the bands at the four corners?
3. Give all four hidden units the same bias, b1 = torch.full((4,), -0.3), and predict the shape of the $p=1/2$ contour before plotting it. Then try +0.3. Explain both pictures.

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)

Part B — From counts to parameters

You have now wired a network whose weights were chosen. The rest of the course is about weights that are learned, and this part is the first training run: the same bigram model as Step 1, rebuilt as softmax of a linear map and trained by gradient descent. To be clear about what is and is not here: the model is multinomial logistic regression, with no hidden layer and no embedding, so it is not yet the neural $n$-gram of Lecture 2, Example 4.7. Gradient descent is here in full; what is missing is a gradient through a hidden layer, and that is Step 3.

On paper first

This derivation is the mathematical core of the step; the code is 30 lines.

1. The model. One-hot encode the current character: $e_a \in \mathbb{R}^V$. Set logits $z = W^{\top} e_a = W_{a,:}$ for a parameter matrix $W \in \mathbb{R}^{V\times V}$, and model

   $$p_\theta(b \mid a) = \operatorname{softmax}(z)_b = \frac{e^{z_b}}{\sum_c e^{z_c}}.$$

   Convince yourself this is the fully general row-stochastic matrix, reparametrized: softmax maps $\mathbb{R}^V$ onto the interior of the simplex, with kernel the constants (so $W$ is identifiable only up to a constant per row). This is exactly Lecture 2’s multinomial logistic regression with one-hot features.

2. The gradient. For one observed pair $(a,b)$ the loss is $\ell = -\log p_\theta(b\mid a) = -z_b + \log\sum_c e^{z_c}$. Show:

   $$\frac{\partial \ell}{\partial z_j} = p_j - \mathbf 1_{j=b}, \qquad\text{i.e.}\qquad \nabla_z \ell = p - y$$

   where $p = \operatorname{softmax}(z)$ and $y = e_b$ is the one-hot target — this is Proposition 3.4. Probabilities minus target — memorize this; it is the last gradient anyone ever computes by hand, and in Step 3 it becomes the seed of backpropagation.

3. Sanity structure. Note $\nabla_z \ell$ has components summing to zero. Why must it, given the kernel observation in 1?

4. Convexity (this model only!). $\ell$ is convex in $z$ (log-sum-exp is convex; linear terms don’t spoil it), and the total loss is convex in $W$. So for this one week gradient descent is guaranteed to find the global optimum — which by Lecture 1’s MLE computation is the count model. From Step 3 onward, convexity is gone forever.

2.6 Data as tensors

From Step 1’s ids, build training pairs:

xs = torch.tensor(ids[:-1])   # list slicing; current char, shape (n,)
ys = torch.tensor(ids[1:])    # next char,    shape (n,)
n = len(xs)                   # 1,115,393

Look and ask. A million training pairs is a lot of tensor. Is it more information than Step 1’s $65\times65$ count matrix $N$?

1. Write the training loss $\mathcal L(W)$ as a formula in $N$ alone, and explain why the million-row tensor cannot contain anything the counts do not.
2. Then say why we build it anyway. (What will be different about the input in Step 3?)

2.7 Forward pass

Implement the loss as a function of $W$ (a (V, V) float tensor, initialized torch.randn(V, V) * 0.01):

def loss_fn(W):                                     # function definition (def)
    logits = W[xs]                                  # indexing by a tensor; (n, V): row a of W is z for input a
    logits = logits - logits.max(dim=1, keepdim=True).values   # stability: softmax is shift-invariant
    p = logits.exp()
    p = p / p.sum(dim=1, keepdim=True)              # (n, V) softmax rows
    return -p[torch.arange(n), ys].log().mean(), p  # method chaining; returns a tuple

def divide(a, b):
    return a // b, a % b       # the comma builds a tuple: two results at once
q, r = divide(17, 5)           # unpack into two names
print(q, r)
both = divide(17, 5)           # or keep the pair as a single tuple
print(both)
print(both[0])                 # tuples index and slice like lists
_, r = divide(17, 5)           # _ = a name for "I don't need this one"
print(r)

import numpy as np
W = np.array([[10., 11.],
              [20., 21.],
              [30., 31.]])
idx = np.array([2, 0, 0, 1])
print(W[idx])                  # one row per index — shape (4, 2)
print(W[idx].shape)

The subtraction of the row max is your first encounter with the difference between mathematics and floating point: prove (one line) that it changes nothing, then try removing it with * 10 larger initial weights and watch the infs.

Look and ask. Before training, evaluate the loss at three weight matrices you choose by hand, predicting each number before you run it.

1. $W=0$.
2. $W=\log P$, with $P$ Step 1’s add-one transition matrix.
3. $W=\log P$ plus a different constant on every row, say + 3 * torch.randn(V, 1).

Then: what is the loss at your random initialization, and why is it so close to the first number?

print(loss_fn(torch.zeros(V, V))[0])                          # 4.1744
print(loss_fn(P.log())[0])                                    # 2.4549
print(loss_fn(P.log() + 3 * torch.randn(V, 1))[0])            # 2.4549

2.8 Backward pass — by hand

Using your derivation ($\nabla_z \ell = p - y$, averaged over the batch, routed back to the rows of $W$):

def grad_fn(W):
    _, p = loss_fn(W)                      # tuple unpacking; _ = throwaway name
    dlogits = p.clone()
    dlogits[torch.arange(n), ys] -= 1.0    # p - y
    dlogits /= n                           # in-place division (/=)
    dW = torch.zeros_like(W)
    dW.index_add_(0, xs, dlogits)          # trailing _ = in-place torch method; rows of W accumulate over their occurrences
    return dW

a = [1, 2, 3]
b = a              # b is another name for the SAME list, not a copy
b += [4]           # in-place: modifies the one shared list
print(a)           # a sees the change
b = b + [5]        # rebinding: builds a NEW list, b points at it
print(a)           # a is unchanged this time
print(b)

for t in range(len(xs)):
    dW[xs[t]] += dlogits[t]

Validate before training. Check against a central finite difference: pick a few random entries $(i,j)$ and compare $\bigl(\mathcal{L}(W + h\,E_{ij}) - \mathcal{L}(W - h\,E_{ij})\bigr)/2h$ to grad_fn(W)[i, j]. Adopt this habit permanently: never trust a gradient you haven’t finite-differenced.

⚠️ Do this in float64 — cast W with .double() first. In float32 the check fails on a perfectly correct gradient, and it is worth understanding why before you waste an afternoon on it. A central difference carries two competing errors: truncation $O(h^2)$ from the Taylor remainder, and roundoff $O(\varepsilon |\mathcal L| / h|g|)$ from cancelling two nearly equal numbers. Here $\mathcal L \approx 2.5$ while an individual $\partial\mathcal L/\partial W_{ij} \approx 2\times 10^{-5}$ (the loss is a mean over $n$ examples, so each of the 4,225 entries gets a small share), and the numerator $\mathcal L(W+hE)-\mathcal L(W-hE) \approx 2hg \approx 4\times10^{-8}$ sits below float32’s resolution of $\varepsilon|\mathcal L| \approx 3\times10^{-7}$. The signal is entirely lost in the subtraction. Measured relative error, same code and same $h$, only the dtype changing:

dtype	$h=10^{-3}$	$h=10^{-1}$
float32	5.9 ❌	$8.9\times10^{-2}$
float64	$3.3\times10^{-7}$ ✅	$3.2\times10^{-3}$

Note the columns move in opposite directions: raising $h$ helps float32 (less cancellation) and hurts float64 (more truncation) — the classic U-shaped error curve, with optimum near $h \sim \varepsilon^{1/3}$. In float64 at $h=10^{-3}$ you get the ~6 significant figures the check is supposed to give.

Once in float64, agreement to ~6 significant figures (relative error $\lesssim 10^{-6}$) or your gradient is wrong.

Look and ask. Compute dW = grad_fn(W) at the initialization and study it before you take a single step with it.

1. Print dW.sum(dim=1). Predict the answer from the sanity-structure item of the derivation. Then deduce something about W.sum(dim=1) during training, and check it after 300 steps.
2. Rank the rows of dW by norm (dW.norm(dim=1).argsort()). Predict which characters are at the top and at the bottom before looking, and compare with how often each character occurs. Why is the row for e only about $150$ times larger than the row for Q, when e occurs $410$ times more often?

dW = grad_fn(W)
print(dW.sum(dim=1).abs().max())            # ~ 4e-5: zero, to float32 roundoff
norms = dW.norm(dim=1)
order = norms.argsort(descending=True)
print([vocab[i] for i in order[:5]])        # [' ', 'e', 't', 'h', 'o']
print([vocab[i] for i in order[-5:]])       # ['Z', 'X', '3', '&', '$']
counts = N.sum(dim=1).float()
print(torch.corrcoef(torch.stack([norms, counts]))[0, 1])   # 0.94

2.9 Gradient descent

W = torch.randn(V, V) * 0.01
for step in range(300):             # for loop over range
    L, _ = loss_fn(W)
    W -= 50.0 * grad_fn(W)          # yes, learning rate 50 — see checkpoint discussion
    if step % 20 == 0:              # % = modulo (remainder)
        print(step, L.item())       # .item(): one-element tensor -> Python number

for step in range(5):          # 0, 1, 2, 3, 4 — stops before 5
    print(step)
print(list(range(2, 20, 5)))   # optional start and stride
print(17 % 5)                  # remainder of 17 divided by 5
total = 0
for step in range(100):
    if step % 20 == 0:         # true for step = 0, 20, 40, 60, 80
        total += 1
print(total)

Then experiment with the learning rate: 0.1, 1, 50, 500. Record what you see. (With the full dataset as one batch the loss is smooth and convex, so this is the cleanest look at step-size behaviour you’ll ever get.)

Measured, loss after just 30 steps from the same initialization:

Two lessons in one table. The tiny rates are not wrong — they are converging, just impractically slowly from $\ln 65 = 4.174$. And lr 500 has gone above where it started: it overshoots the valley each step. Why does this well-behaved convex problem tolerate a learning rate of 50 when Step 3’s network will need $\approx 0.1$? Think about the scale of $\nabla_W$: it is a mean over a million examples spread across 4,225 entries, hence tiny (of order $10^{-5}$), and the step size has to compensate. Learning rates are not transferable constants; they are inverse to the curvature and gradient scale of the specific problem.

Look and ask. The table says $50$ beats $0.1$, $1$, and $500$. It does not say $50$ is best.

1. Sweep the learning rate over $50, 100, 150, 200, 250, 300, 400, 500$ from the same initialization and record the loss after one step and after thirty. Which rate wins each contest? Why are they different?
2. Plot the gap between your loss and Step 1’s $2.4549$ against the step number on a log scale. Is it a straight line? What would a straight line have meant?

2.10 The punchline

Compare your trained model to Step 1’s count model:

• final training loss vs the count model’s cross-entropy from Step 1;
• softmax(W, dim=1) rows vs P’s rows (plot a few, or (softmax_W - P).abs().max());
• samples from the trained model, using your Step 1 sampler.

You have two proofs that they must agree in the limit. But look carefully at what 300 steps actually delivers — the numbers are in the checkpoints, and the honest answer is “close, but visibly not there yet.” That gap is the point of the following question, which is really a numerical-analysis question rather than a machine-learning one:

Gradient descent on a convex problem converges at a rate governed by the condition number of the Hessian. What is the Hessian here, and why is it badly conditioned? (Hint: the row of $W$ for the character Z receives gradient contributions only from the handful of times Z occurs. Rows whose characters are common get large, informative gradients; rows for rare characters get almost nothing. Now look at the per-row agreement in the checkpoints and see whether it tracks that prediction.)

This is the first appearance of a theme that runs through Lecture 6: a convergence guarantee and a convergence rate are very different things, and essentially all of the practical craft of training lives in the gap.

Finally: which of Step 1’s ingredients plays the role of the $+1$ smoothing here? (Hint: what does initializing $W \approx 0$ mean in probability space, and does 300 steps fully erase it?)

Look and ask. Put numbers on both questions.

1. Every 50 steps, record $\max_b\lvert\operatorname{softmax}(W)_{ab}-P_{ab}\rvert$ for the rows $a\in\{$space, e, t, Q, Z, X$\}$ and plot the six curves together. Predict the ordering before you look.
2. After 300 steps, find the never-seen pair (a cell with $N_{ab}=0$) to which the model assigns the most probability. Which pair is it, what probability does it get, and why that pair? Compare with what add-one smoothing would give it.

Checkpoints

Part A:

• 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$.

• Look-and-ask answers: at most two hidden units are ever active (active counts for the five points: $0,2,1,2,2$); $(0.7,0)$ and $(0.5,0.5)$ score $0.9168$ and $0.5000$ inside; at $\gamma=4$ the $p=0.9$ and $p=0.1$ contours are the diamonds $r=0.725$ and $r=1.275$; a shared hidden bias of $-0.3$ turns the boundary into an octagon with vertices $(\pm1.3,\pm0.3)$ and $(\pm0.3,\pm1.3)$, and $+0.3$ leaves nothing inside (largest inside probability $0.168$).

Part B:

• Finite-difference check in float64 passes: relative error $\approx 3\times10^{-7}$ at $h=10^{-3}$. (In float32 it “fails” at $\approx 5.9$ — see the warning above. If you see that, your gradient is probably fine and your dtype is not.)

• Loss decreases monotonically (full-batch, sane learning rate), approaching Step 1’s 2.4549 from above but not reaching it in 300 steps. Measured trajectory at lr 50:

  step	0	50	100	150	200	250	300
  loss	4.1731	2.6530	2.5690	2.5352	2.5168	2.5052	2.4973

  Still 0.042 nats above the count model, and closing slowly — the tail of this curve is the ill-conditioning discussed above.

• Agreement between $\operatorname{softmax}(W)$ and $P$ after 300 steps is strongly row-dependent, exactly as the conditioning argument predicts:

  quantity	value
  $\max_{ij}\lvert\operatorname{softmax}(W)-P\rvert$ (all rows)	0.542
  mean $\lvert\cdot\rvert$ over all entries	0.0052
  $\max\lvert\cdot\rvert$ over the 53 rows with $>1000$ counts	0.070

  So the typical entry is already accurate to three decimals, the well-populated rows to about one, and the single worst entry — in the row for a rare character — is still off by half a unit of probability. (Measured: the worst row belongs to Q, which occurs 231 times against e’s 94,611 — rare enough that 300 steps barely move it.) Identify that row yourself and the answer should feel inevitable.

• Look-and-ask answers: the loss is $\ln 65=4.1744$ at $W=0$, $2.4549$ at $W=\log P$, and unchanged by per-row constants; gradient rows sum to zero and W.sum(dim=1) is conserved (float64 drift $4\times10^{-12}$); gradient row norms correlate $0.94$ with character counts, with space, e, t largest and &, $ smallest; over 30 steps the best learning rate is about $100$, and $300$ stalls; the most-favoured never-seen pair after 300 steps is &→; at $0.0157\approx1/65$.

Troubleshooting

Part A:

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

Part B:

• Loss stuck at $\ln 65 \approx 4.17$: your gradient is zero or you’re not actually updating W (in-place vs reassignment).
• Loss nan: learning rate too large, or you removed the max-subtraction.
• index_add_ unfamiliar: do it with a loop over a 10,000-pair subset first; then read its docs. It is exactly “accumulate $\nabla_z\ell$ into row $x_t$.”

Going further (optional)

Part A:

• Change confidence without changing geometry. Run the grid of Task 2.5 with $\gamma\in\{0.25,1,4,20\}$ and put the four probability plots side by side with the same color scale. Which set of points has probability exactly $1/2$ in every plot? What happens to probabilities away from the boundary as $\gamma$ grows, and as $\gamma$ approaches zero? Why is changing $\gamma$ equivalent to changing softmax temperature?

• 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},$$

  keeping 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)$ and $(a,b\pm r_2)$.

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

Part B:

• Add an MLP: $z = W_2\,\tanh(W_1 e_a)$ with a width-100 hidden layer. Derive the gradient by the chain rule (two more lines of the same style) and train it. It can’t beat the bigram table (why not? — the input is still only one character), which makes it the perfect controlled introduction to hidden layers. This is also a preview of Step 3’s pain: by the time you have three layers, hand gradients stop scaling — hence autodiff.
• Regularize: add $\lambda \lVert W\rVert_F^2$ to the loss, sweep $\lambda$, and show (paper + plot) it interpolates toward the uniform model, playing exactly the role of Laplace smoothing.

Catch-up

Joining now? Part A is self-contained: skim Lecture 2, Section 4 and Section 7, then run the supplied blocks in a fresh notebook. Part B depends only on Lecture 2, Section 3 and the shape of Step 1’s data: run solutions/step-01.ipynb (5 min), skim its prose, then do Tasks 2.6–2.10 in full. You do not need anything else from the Shakespeare corpus until Step 3.