Module 1: The smallest network that can fail

The task: two moons

We need a task that is easy enough for a tiny network but impossible for a linear model; otherwise depth would be decoration, and this workshop is about what depth does to gradients. The classic choice is two moons: two interleaved crescents of points, one per class. The script generates it from scratch (two half-circles plus Gaussian noise), so there is no dataset dependency:

# Upper moon: half circle of radius 1 centered at the origin.
t = rng.uniform(0.0, np.pi, size=n_half)
upper = np.stack([np.cos(t), np.sin(t)], axis=1)

# Lower moon: mirrored half circle, shifted to interleave with the upper one.
t = rng.uniform(0.0, np.pi, size=n_half)
lower = np.stack([1.0 - np.cos(t), 0.5 - np.sin(t)], axis=1)

# Stack both moons into one dataset and jitter every point with Gaussian noise.
X = np.concatenate([upper, lower]).astype(np.float32)
X += rng.normal(0.0, NOISE, size=X.shape).astype(np.float32)

Section 1's pass of the experiment script then shuffles, holds out a test set, and logs the dataset's vital signs; these log(...) lines write the capture below, at captures/results.log lines 8–10:

# Shuffle, then hold out the last 200 points as the never-trained-on test set.
perm = rng.permutation(N_SAMPLES)
X, y = X[perm], y[perm]
X_train, y_train = torch.from_numpy(X[:N_TRAIN]), torch.from_numpy(y[:N_TRAIN])
X_test, y_test = torch.from_numpy(X[N_TRAIN:]), torch.from_numpy(y[N_TRAIN:])

# Record the dataset's vital signs in the citable log.
log(f"  samples: {N_SAMPLES} ({n_half} per class), noise sigma = {NOISE}")
log(f"  split: {N_TRAIN} train / {N_SAMPLES - N_TRAIN} test")
log(f"  input feature std: x1 = {X[:,0].std():.3f}, x2 = {X[:,1].std():.3f}")

samples: 1000 (500 per class), noise sigma = 0.1
split: 800 train / 200 test
input feature std: x1 = 0.890, x2 = 0.511

Let's go deeper: 1,000 points, perfectly balanced at 500 per class, so 50% accuracy is exactly the coin-flip baseline; 800 points to train on and 200 held out, so every accuracy we quote is on data the network never saw; and the two input features have standard deviations 0.89 and 0.51 (order one). That last line looks like trivia now, but it is a load-bearing measurement: Module 2 tracks how this order-one signal shrinks as it moves through the layers, and the shrinking starts from here.

Because the moons interleave, any straight-line boundary misclassifies a big chunk of one moon's tip. A network must curve the boundary through the gap, a genuinely nonlinear job.

The network: 2 in, four hidden layers of 4, 1 out

Follow the diagram top to bottom: the two input features enter at the input box; each hidden box is a Linear layer mapping into 4 neurons followed by an element-wise sigmoid (the function from Module 0); the output box is a final Linear down to a single logit; and the bottom edge, labeled p = sigmoid(logit), converts that logit to a probability which the cross-entropy loss box scores against the true label y.

The constructor, from experiment.py:

class TinyNet(nn.Module):
    """One network class for every variant in the workshop.

    plain / batch norm / LayerNorm come from `norm`; ReLU from `act`;
    residual from `residual`; the depth-8 nets from `depth`. Sharing one
    code path is itself a control: variants differ ONLY in these options.
    """

    def __init__(self, norm="none", act="sigmoid", residual=False, depth=DEPTH):
        super().__init__()
        # Set the manual seed so that different configurations (norm/act/...)
        # start with the exact same weight values, for a fair comparison.
        torch.manual_seed(SEED)

        # Dimensions: input (2 features) -> depth hidden layers of WIDTH -> output.
        dims = [2] + [WIDTH] * depth

        # One Linear layer per hidden level.
        self.hidden = nn.ModuleList(
            nn.Linear(dims[i], dims[i + 1]) for i in range(depth))

        # The normalizer that sits between each Linear and its activation.
        if norm == "batch":
            self.norms = nn.ModuleList(nn.BatchNorm1d(WIDTH) for _ in range(depth))
        elif norm == "layer":
            self.norms = nn.ModuleList(nn.LayerNorm(WIDTH) for _ in range(depth))
        else:
            # Identity() is a do-nothing placeholder: the plain net runs the
            # exact same forward() code path as the normalized variants.
            self.norms = nn.ModuleList(nn.Identity() for _ in range(depth))

        # Final layer projects to a single logit; the sigmoid lives in the loss.
        self.out = nn.Linear(WIDTH, 1)

        # Activation applied at every hidden layer.
        self.act = torch.sigmoid if act == "sigmoid" else torch.relu
        self.residual = residual
        self.rec = None   # intermediates from the last recorded forward pass

In the plain network the norms list holds do-nothing nn.Identity() modules; Module 4 swaps them for BatchNorm1d and changes nothing else. With DEPTH = 4 and WIDTH = 4, Section 3's pass of the experiment script builds the two main networks and counts parameters:

plain = TinyNet(norm="none")
bn = TinyNet(norm="batch")
n_params = sum(p.numel() for p in plain.parameters())
log(f"  plain net parameters: {n_params}")

and, skipping past the checksum lines you will meet below, inspects one middle layer's starting weights:

# The starting weight scale: one of the two sub-unity factors in Module 2's
# per-layer hop product.
w1 = plain.hidden[1].weight
log(f"  layer-2 weight stats: std = {w1.std().item():.4f}, "
    f"max|w| = {w1.abs().max().item():.4f}  (PyTorch default init, fan_in = {WIDTH})")

When Section 3's pass runs (it is the part of the script that constructs the networks and inspects their starting state), these are the statistics obtained (captures/results.log lines 25 and 29):

plain net parameters: 77
layer-2 weight stats: std = 0.3114, max|w| = 0.4777  (PyTorch default init, fan_in = 4)

Let's go deeper: the whole network is 77 numbers; you could print them on an index card. And the starting weights, drawn by PyTorch's default nn.Linear initialization, are small: for a layer whose inputs are 4 wide (fan_in = 4), the weights come out with standard deviation 0.31 and none exceeds 0.48 in magnitude. Small starting weights are standard practice and usually sensible, but hold onto that 0.31: in Module 2 it becomes one of the two culprits.

Why deep and tiny, rather than wide and shallow?

A single hidden layer of, say, 16 neurons would learn two moons easily, but it would tell us nothing, because vanishing gradients is a depth phenomenon: as Module 2 will show, the backward signal picks up one multiplicative factor per layer, so the damage compounds geometrically with depth and is invisible in shallow nets. We want the smallest network where the compounding is plainly visible, so we spend our tiny parameter budget on depth: four layers deep, four neurons wide.

And to be clear, 77 parameters is plenty for two moons; capacity is not the problem. The proof arrives in Modules 4 and 5: this exact architecture, with the same 77-ish parameter budget, reaches 100% test accuracy the moment gradients are allowed to flow. When the plain version fails, it fails because of trainability, not size.

The control: same starting weights, always

The comparisons ahead ("plain vs batch norm") are only meaningful if both networks start from the same place. Look back at the constructor: torch.manual_seed(SEED) runs inside __init__, immediately before the Linear layers are created. Every TinyNet, whatever its options, therefore draws exactly the same starting weights in exactly the same order. The script verifies this by summing the absolute values of all hidden weights in both nets:

def weight_checksum(net):
    """Sum of |w| over all hidden-layer parameters: a cheap fingerprint used
    in S3 to PROVE two variants start from identical weights."""
    return sum(p.abs().sum().item() for lin in net.hidden for p in lin.parameters())

Running the check on both networks writes three lines to the log (lines 26–28):

hidden-weight checksum, plain = 18.770746
hidden-weight checksum, bn    = 18.770746
identical starting weights: True

The two checksums agree to the last printed digit: whatever differences we measure between these two networks from here on, starting weights are not the explanation. One caveat for honesty: the batch-norm variant does add its own parameters (a scale and shift per neuron; Module 4 covers them), but they are initialized to the deterministic values 1 and 0, so they draw no random numbers and the shared weights stay identical.