Neural network playground

This page runs a real neural network in your browser. It is a multilayer perceptron (MLP): a stack of layers, where every neuron takes a weighted sum of its inputs, adds a bias and applies an activation function. The task is binary classification in the plane. Every point $(x_1, x_2)$ is blue or orange, and the network learns to predict the colour from the position.

Everything is computed on this page, in plain JavaScript: the forward pass, backpropagation, SGD and Adam. Nothing is sent to a server. We checked the gradients against numerical derivatives (see Checking the gradients below).

What you see in the figure:

• Output (the large square). The background colour at each point is the probability that the network gives to the blue class there: blue means "blue", orange means "orange", pale means "not sure". The line where the colour changes is the decision boundary.
• Network (in the middle). The first column holds the input features. Then come the hidden layers and the output neuron. Each small square shows what that neuron outputs at every point of the plane (for ReLU and linear neurons, the colours are scaled so that the largest absolute value gets the full colour; the colour bar under the large square shows the scale). Each line is one weight: blue for a positive weight, orange for a negative one, and thicker for a larger $|w|$.
• Loss (below the network). The training loss and the test loss after every epoch (one pass through the training data).

The playground #

Interactive figure: loading…

The controls, from top to bottom:

• Step trains one epoch; Reset starts again from new random weights.
• Learning rate $\eta$, batch size $B$ and L2 regularization $\lambda$ can change while the network trains. Changing the activation, a layer or a feature restarts training.
• Epochs per frame only makes the animation faster; it does not change the result.
• Noise moves the points randomly; training data is the share of the 400 points used for training. The rest is the test set: the network never trains on it, so its loss tells us how well the network generalizes to new data.

Try this #

Before each experiment, guess the answer to the question. Then run it and compare.

1. A linear model and feature engineering. Press 1. There is no hidden layer, so the model is logistic regression on $x_1, x_2$. Press Play.

   Question: can a straight line separate the blue disk from the orange ring?

   Observe: the test loss stays near $\ln 2 \approx 0.69$, the loss of a coin flip. Now click the features $x_1^2$ and $x_2^2$ to switch them on. Within about 100 epochs the test accuracy reaches 100%, still with no hidden layer. Also try the XOR data with only the feature $x_1x_2$.

2. Hidden neurons learn features. Press 2: one hidden layer with one ReLU neuron. Press Play, then add neurons one at a time with the + above the layer.

   Question: how many ReLU neurons do you need before the boundary closes around the blue disk?

   Observe: with 1 neuron the boundary is a straight line (about 63% test accuracy), with 2 an open wedge (about 80%), and with 3 a closed polygon around the disk (about 98%). Hover over the hidden neurons: each one is zero on one side of a line and grows linearly on the other side.

3. Depth versus width. Press 3: the spiral, with one hidden layer of 8 ReLU neurons. Press Play and wait for 1000 epochs.

   Question: is one layer of 8 neurons enough to follow the arms of the spiral?

   Observe: in our runs the test accuracy stayed below about 80%. Now press the + next to "1 hidden layer" twice, to get three layers of 8 neurons. The test accuracy usually reaches 95–100%. If the loss jumps around, lower the learning rate to 0.03. Then go back to one layer and switch on $\sin x_1$ and $\sin x_2$: with these features, the same single layer reaches about 90%. Good features can do part of the work of depth.

4. Overfitting and regularization. Press 4: a network with 249 parameters, but only 40 noisy training points (10% of the data). The rings are the test points. Press Play.

   Question: the training loss goes almost to zero. Does the test loss follow it?

   Observe: the test loss is lowest after only 10 to 30 epochs. Then it rises (to about 0.8 after 1000 epochs), while the boundary bends around single training points. Now set L2 regularization to 0.03 and press Reset: the test loss stays near 0.2, and the boundary stays close to a straight line. Try also a smaller network, or 90% training data.

5. Learning rate and batch size. Press 5: the circle data with learning rate 10. Press Play.

   Question: why does a larger step not give faster learning?

   Observe: the loss jumps between large values and does not settle. Try 1 (fast, with spikes), 0.1 (smooth) and 0.001 (almost no progress in 300 epochs). Then, at learning rate 0.1, compare batch size 1 with batch size 30: batch size 1 makes 200 noisy steps per epoch, batch size 30 makes 7 smoother ones.

6. Why we need nonlinear activations. Press 6: four hidden layers of 8 neurons with the linear activation $\sigma(z) = z$. Press Play.

   Question: 249 parameters and four layers: can this network draw a circle?

   Observe: no. The boundary is a straight line, and the test loss stays near $\ln 2$. Switch the activation to tanh or ReLU, and the same network solves the task in less than 100 epochs.

The math behind the playground #

The model #

The input of the network is the vector of the chosen features, $h^{(0)} = \phi(x) \in \R^k$, for example $\phi(x) = (x_1, x_2, x_1^2, x_2^2)$. Each feature is divided by its largest absolute value on the square $[-6, 6]^2$, so every input lies in $[-1, 1]$. This feature scaling keeps $x_1^2$ (up to 36) from dominating $x_1$ (up to 6).

Each hidden layer $l = 1, \dots, L$ computes an affine map followed by the activation $\sigma$, applied to every entry:

The output neuron is affine only. Its value $z$ is the logit, and the sigmoid turns it into a probability:

The network predicts blue when $\hat p > 0.5$, that is, when $z > 0$. So the decision boundary is the set of points where $z = 0$. The activations in the playground are ReLU, $\sigma(z) = \max(z, 0)$; tanh; the sigmoid; and the linear (identity) function. The weights start random: $W_{ij} \sim \mathcal N(0, 2/n_{\text{in}})$ in the hidden layers of a ReLU network (He initialization), and $\mathcal N(0, 1/n_{\text{in}})$ for the other activations and for the output layer. Here $n_{\text{in}}$ is the number of inputs of the layer. The hidden biases start at 0.1 (so fewer ReLU neurons start "dead", with zero output everywhere), and the output bias starts at 0.

The loss #

For one example with label $y \in \{0, 1\}$ (1 = blue), the loss is the binary cross-entropy

The second form gives the same number, but it never takes the logarithm of 0, so it is safe for large $|z|$. One training step uses a mini-batch $\mathcal B$ of $B$ examples and minimizes

The second term is the L2 penalty (also called weight decay). It acts on the weights only, not on the biases. The curves in the figure show the mean cross-entropy without the penalty.

Backpropagation #

Backpropagation is the chain rule, applied from the output back to the input. Write $\delta = \partial \ell / \partial z$ for the error signal of a neuron. For the output neuron, the derivative of cross-entropy through the sigmoid is very simple:

Each hidden layer gets its error signal from the layer above it. Here $\odot$ is the entrywise product:

with $\sigma'(z) = \mathbb 1[z > 0]$ for ReLU, $1 - \tanh^2 z$ for tanh, and $s(1 - s)$ for the sigmoid $s$. The gradients of the weights are outer products, averaged over the batch, plus the gradient of the penalty:

The update: SGD and Adam #

Stochastic gradient descent (SGD) moves every parameter $\theta$ a small step against its gradient $g = \partial J / \partial \theta$:

Adam keeps running averages of the gradient and of its square, for every parameter separately ($\beta_1 = 0.9$, $\beta_2 = 0.999$, $\epsilon = 10^{-8}$, step $t = 1, 2, \dots$):

The divisions by $1 - \beta^t$ correct the bias of the averages, which start at zero. Here the L2 term is part of $g$. This is "Adam with L2", which is not the same as AdamW (decoupled weight decay). An epoch is one pass through the training set in mini-batches, in a new random order each time.

Checking the gradients #

A wrong gradient often still trains, only worse, so bugs in backpropagation are easy to miss. We compared every entry of the gradient with a central difference:

We tested all four activations, networks with 0 to 4 hidden layers and up to 7 inputs, and $\lambda$ from 0 to 0.1. The largest relative error was about $10^{-7}$, and the largest absolute error about $10^{-10}$. This is the size we expect from rounding and from the $O(h^2)$ error of the formula. For ReLU, the test skips the rare case where $\theta \pm h$ crosses a kink of $\max(z, 0)$.

Three ideas to take away #

Feature engineering #

A model with no hidden layer is logistic regression: $z = w^\top \phi(x) + b$. Its decision boundary $z = 0$ is a straight line in the space of features, but it does not have to be a line in the plane. With the features $x_1^2$ and $x_2^2$, the boundary is

an ellipse centred at the origin when $w_1$ and $w_2$ have the same sign. This is exactly the shape the circle data needs. For XOR, the class is the sign of $x_1 x_2$, so the single feature $x_1 x_2$ is enough. Choosing such inputs by hand is feature engineering. It works very well when we know the structure of the problem. The hidden layers of a network do the same job automatically: they learn features from the data. This costs more data and more computation, but it also works when we do not know which features to use. Kernel methods, such as SVMs, are a third way: they use a fixed feature map that can be very large.

Depth versus width #

Look at a ReLU neuron in the first hidden layer. It is zero on one side of the line $w^\top x + b = 0$ and grows linearly on the other side. One hidden layer adds such ramps, so its output is piecewise linear, with kinks only along the lines of its neurons. With $n$ neurons and two inputs, these $n$ lines cut the plane into at most $1 + n + \binom{n}{2}$ pieces. Every boundary that one layer can draw is built from these pieces.

A second layer applies ReLUs to sums of these ramps. This folds the pieces again, and each fold can create new kinks in many places at once. The number of linear pieces that a deep ReLU network can make grows exponentially with the depth. For a fixed depth, it grows only polynomially with the width (Montúfar et al., 2014). The spiral needs a boundary with many turns, so depth helps, as experiment 3 shows. Two warnings: this is about what a network can represent, not about what training finds; and deeper networks are harder to train (the jumps in the loss in experiment 3). See MLPs and universal approximation for the tent-map example.

Overfitting and regularization #

A model overfits when it fits the noise of the training data: the training loss keeps falling, but the test loss rises. The difference between test loss and training loss is the generalization gap. Experiment 4 makes overfitting easy: a network with 249 parameters and only 40 noisy points. The two Gaussian blobs overlap, so no boundary can be right for every point. The two blobs have the same spread, so the best boundary is a straight line.

Why does the test loss grow so much, while the test accuracy drops much less? With no regularization, the network separates the training points perfectly. Then cross-entropy can still go down by making $|z|$ larger, so the weights keep growing and the predictions become very confident. A test point on the wrong side with a large $|z|$ costs about $|z|$ in loss. So the test loss rises even when few test predictions change.

There are four standard remedies, and you can try each one in the playground:

• More data: move the training data slider up.
• A smaller model: remove layers or neurons.
• L2 regularization: the penalty $\frac{\lambda}{2}\|W\|^2$ keeps the weights small. Small weights give small slopes, so the boundary stays smooth and the predictions stay less confident.
• Early stopping: stop when the test loss is lowest. (In practice, use a separate validation set for this decision, and keep the test set for the final report only.)

Learn more on this site #

• MLPs and universal approximation: layers, the collapse of linear networks, and why depth helps.
• Training a network: losses, mini-batches and the training loop: cross-entropy from logits, mini-batches and the limit on the learning rate.
• Backpropagation on computational graphs and backpropagation through losses and layers.
• Numerical differentiation and gradient checking: why the central difference is accurate, and how to choose $h$.
• Activations and initialization: ReLU, tanh, the sigmoid, and He initialization.
• Regularization and generalization: weight decay, early stopping and double descent.
• Loss surfaces and gradient descent and stochastic gradient descent and batch size. Adam and AdamW are in the Optimization track.
• ML Fundamentals: logistic regression, the bias–variance trade-off and kernels.

This page is inspired by TensorFlow Playground by Daniel Smilkov and Shan Carter. The code here is written from scratch.