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, then train. 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.$$

   Then erase those weights and use back propagation and gradient descent to learn the same classifier from labelled examples. The comparison separates what the architecture can represent from what optimization actually finds.

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, then train, a diamond classifier

The Neural Function Playground shows every neuron as a function of $(x_1,x_2)$ while you change its weights and bias. You will design the network there first, then transcribe your construction into PyTorch. Do not begin by searching for finished matrices: the point is to discover what each unit has to do.

2.1 Classify by hand in the playground

Open the playground, click Reset, and use the targets across the top as exercises. Zero all parameters before each new attempt, then change weights and biases by dragging, with the arrow keys, or by double-clicking for exact entry. The target is only an overlay; nothing is trained for you.

Try as many of these as time permits:

1. Half-plane (warm-up). Delete the hidden layer. Can one identity output unit classify the target?
2. Diamond. Click Reset again to restore the default architecture: four ReLU hidden units, one identity output, and threshold $\tau=0$. Zero its parameters. First try to classify part of the target region. Then improve it until the solid decision boundary follows all four sides of the dashed diamond.
3. XOR. The score counts the four marked examples, so partial success is visible. Can two ReLU units classify all four?
4. Three sectors. This requires three softmax outputs. Think of each output logit as favouring one direction.

Look and ask. A hidden unit with weight vector $(a,b)$ and bias $c$ has preactivation $ax_1+bx_2+c$. Where is its ReLU active? What geometric feature in its heatmap changes when you alter the weight vector, and what changes when you alter the bias?

2.2 Explain the diamond you found

Let $\mathbf{x}=(x_1,x_2)^\top$. Your diamond network should produce a signed scalar score $s$: positive inside, negative outside, and zero on the boundary. Explain your construction on paper before translating it.

1. Verify

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

   Use the identity to explain the four hidden heatmaps. Which input half-plane makes each unit active?

2. Write the playground’s two layers as

   $$z_1=W_1\mathbf{x}+b_1, \qquad a_1=\operatorname{ReLU}(z_1), \qquad s=W_2a_1+b_2.$$

   Give every matrix and vector dimension. Show that your particular values produce

   $$s=1-|x_1|-|x_2|.$$

3. How many hidden units can be active at one point? Show that the hidden layer still determines the input, but the scalar score does not. Give two very different points with the same score.

2.3 Translate the playground into one forward pass

Begin a notebook with PyTorch:

import torch

Now package the whole forward pass. The playground prints an executable W1 = torch.tensor(...) or W2 = torch.tensor(...) beneath each bundle of wires, and it prints the corresponding two lines of layer computation beneath each destination column. Copy those pieces into the scaffold below. The biases are the b values printed above the neuron tiles; transcribe them into b1 and b2 in the same top-to-bottom order as the matrix rows. We use the playground’s names here: z1 is the hidden layer before ReLU, and a1 is its activation after ReLU. The latter is the vector often called $h$ in formulas for a neural network.

Note that you will have to enter b1 = torch.tensor([...]) and b2 = torch.tensor([...]).

def diamond_net(X):
    # Copy the playground's two layer-1 computation lines.

    # The playground calls this layer's preactivation z2 and activation a2.
    # Copy W2 and the two computation lines, then enter its bias.

    return a2

The lecture writes one input as a column and uses $W_1\mathbf{x}$. The playground’s copied code stores a batch with one sample per row, so it uses X @ W1.T. For a batch of $n$ points, the dimensions are

then

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)

Test diamond_net on the two points $(0.7,0)$ and $(0.5,0.5)$. Predict both outputs before running the code. How do the signs of the outputs determine the classifications?

2.4 Draw the decision surface

Evaluate the network on a fine grid, then plot its signed output:

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)

surface = diamond_net(grid).reshape(gx.shape)

plt.figure(figsize=(6, 5))
contour_plot = plt.contourf(gx.numpy(), gy.numpy(), surface.numpy(),
                            levels=30, cmap='Purples')
plt.contour(gx.numpy(), gy.numpy(), surface.numpy(),
            levels=[0.0], colors='black', linewidths=2)
plt.xlabel('x1')
plt.ylabel('x2')
plt.axis('equal')
plt.colorbar(contour_plot, label='network output a2')
plt.show()

The black $a_2=0$ 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 $a_2=0.3$ and $a_2=-0.3$ contours lie. Then add levels=[-0.3, 0.0, 0.3] 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 $a_2=0$ 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)

2.5 Learn the same classifier from data

You chose the preceding weights by reasoning about the target. Now erase that construction and ask gradient descent to discover a classifier from labelled examples.

A. Set up a trainable diamond network

1. Save a share link for your hand-built diamond if you want to compare with it later. Then click Reset and select the Diamond target.
2. In Visualizers, open Training and close Output and Inspector. Leave the other cards closed for the moment.
3. Change the final activation from identity to sigmoid. Its output $z$ is now a number between $0$ and $1$, interpreted as the predicted probability that the point is inside. The target $y$ supplied to the loss function is $1$ inside and $0$ outside.

First click Zero all parameters. Before training, predict what will happen. All four hidden neurons now have the same weights and bias. Will back propagation give them different gradients, or will they remain identical after every update? Why does this prevent them from learning the four different sides of the diamond? Try one update if you are unsure.

Now click Randomize parameters. Random initialization breaks the symmetry between the hidden neurons.

B. Write the loss

For one example, negative log likelihood is the negative logarithm of the probability assigned to the correct class:

The playground loss editor has no if statement. Find one expression for the probability assigned to the correct class that works for both $y=0$ and $y=1$. Hint: use abs, and remember that $y$ is exactly $0$ or $1$. Replace the default squared-error loss with the corresponding negative log loss. Click away from the editor or press Ctrl/⌘ Enter to check that it is a valid scalar loss.

Test your expression mentally with $y=0$ and $y=1$, for both a good prediction and a bad one. Which case should have the larger loss?

C. Inspect one gradient step

1. Open Back propagation.
2. Click a point in the Training window. The click adds that labelled point to the batch and selects it in the composition tree.
3. Read the computation from input to loss. Red numbers are forward values; blue numbers are derivatives of the loss with respect to that node. Locate the blue derivatives at several weights and biases.
4. Before pressing Train, decide which direction each displayed parameter should move. If a parameter is $\theta$ and the blue number is $\partial\ell/\partial\theta$, write the update using the learning rate $\eta$. Then press Train once and compare the new parameter value with your prediction.
5. Try an inside point and an outside point. Which parts of the network get no gradient when a ReLU is inactive?

D. Train continuously

Open Batch. Generate a random batch, train once, and confirm that its average loss usually decreases on that batch. Then choose a larger number of data and click Train continuously. Each step replaces the batch with new random points before updating the parameters.

This experiment may need many samples and some adjustment of the learning rate. If the loss jumps or the picture becomes unstable, lower $\eta$. If almost nothing changes, raise it. Stop training periodically, reopen Output, and compare the learned decision boundary with the target and with your hand-built network.

• Did training recover the same boundary? Did it recover the same actual weights, or only an equivalent function? Remember that hidden neurons can be permuted, rescaled, or share work.
• What happens with two or three hidden neurons? What changes with more than four?
• Add or remove a hidden layer, randomize again, and retrain. Which architectures learn reliably? Which are difficult to optimize even if they appear capable of drawing the target?

E. Repeat with the other targets

1. XOR. Select the XOR target, use one sigmoid output, randomize, and reuse the binary negative log loss. Can continuous training classify all four marked points? Try different hidden widths and depths.
2. Three sectors. Select the target, use exactly three outputs, and set the final activation to softmax. Now $y$ is a one-hot vector and $z$ is a vector of three probabilities. Construct a scalar negative log loss from the probability assigned to the correct class. You may use abs, sum, and mean; the scalar 1 can be broadcast as a vector of ones. Train, tune the batch size and learning rate, and compare the learned sector boundaries with the target.

return -log(abs(1 - y - z))

return -log(sum(y * z))

return -log((sum(abs(1 - y - z)) - 1) / 2)

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

Why do we subtract the maximum in the second line? Prove that this has no effect on the return values, mathematically. What purpose does it serve computationally?

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:

• You attempted one or more playground targets by hand and completed the diamond before copying its parameters into PyTorch.
• For a sigmoid output, return -log(abs(1 - y - z)) is binary negative log likelihood for both labels. Gradient descent uses $\theta\leftarrow\theta-\eta\,\partial\ell/\partial\theta$.
• With every parameter zero, the hidden and output-weight gradients vanish and only the output bias can move. Random initialization breaks the hidden-unit symmetry.
• You trained on random batches, compared the learned and hand-built diamond classifiers, and repeated the experiment for XOR and the three-sector softmax target.
• For the canonical diamond, $W_1$ has shape $(4,2)$, $b_1$ shape $(4)$, $W_2$ shape $(1,4)$, and $b_2$ shape $(1)$; its scalar output is $s=1-|x_1|-|x_2|$.
• The outputs at $(0.7,0)$ and $(0.5,0.5)$ are $0.3$ and $0$.
• The $a_2=0$ contour is the diamond $|x_1|+|x_2|=1$.
• Look-and-ask answers: $(0.7,0)$ and $(0.5,0.5)$ have outputs $0.3$ and $0$; the $a_2=0.3$ and $a_2=-0.3$ contours are the diamonds $r=0.7$ and $r=1.3$; 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 output $-0.2$).

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:

• Training changes only the output bias: do not initialize every parameter at zero. Randomize the parameters so hidden units can receive different gradients and specialize.
• Loss is invalid for the three-sector target: the loss must reduce the three-component vector to one scalar; use sum or mean. Also confirm that the output has three units and uses softmax.
• Continuous training is unstable or produces a non-finite loss: lower the learning rate. If training barely moves, try a larger batch and then cautiously raise the learning rate.
• 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: the playground score must be positive inside, so check the signs of the single row of $W_2$ and of $b_2$.
• 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:

• 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 complete the playground diamond and translate it with the scaffold in Task 2.3. 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.