← mino.mobi

descent

the fractal edge of neural-network trainability · sohl-dickstein 2024

↑ geometry pack · a Mandelbrot cousin, drawn by gradient descent

One step of gradient descent, θ → θ − η·∇L(θ), is an iterated map — and whether training a tiny neural net converges or explodes, as a function of its two learning rates (η₁ for the first layer, η₂ for the second), is a fractal. It's the exact Mandelbrot escape-time algorithm, with "one training step" in place of z² + c. Pan (drag) and zoom (scroll) the boundary; click any point to watch that run's loss converge or blow up. Or flip the model at each point to a tiny recurrent language model — then every dark pixel is a nanomodel that learned to talk, and clicking it shows what it says.

what lives at each point
training steps · T (max escape)
400
hidden units · H
2
colour map
resolution (CPU cost)
420
probed run · loss vs step
Dark interior = trainable (loss converges, like the Mandelbrot set's body); bright filaments = the boundary, where a hair's change in learning rate flips training from success to blow-up. The structure is genuinely self-similar — keep zooming. Runs on a WebGPU compute shader where available (one thread per pixel) and falls back to a progressive CPU renderer otherwise — the "renderer" readout below says which.

Training is an iterated map

Full-batch gradient descent updates every parameter as θ ← θ − η·∇L(θ). Apply it again and again and you have a dynamical system: a fixed map iterated from a fixed start, parameterised by your hyperparameters. For some learning rates the iteration settles into a minimum — training converges. For others the loss grows without bound — training diverges. The question "which hyperparameters train?" is the question "which starting conditions stay bounded?", the same question the Mandelbrot set answers for z → z² + c.

The boundary is fractal

Jascha Sohl-Dickstein showed in 2024 that the boundary between trainable and untrainable hyperparameters is a fractal, robustly, across architectures. Here the plane is the two learning rates — η₁ for the input layer's weights, η₂ for the output layer's — of a tiny net x → tanh(W₁x + b₁) → W₂·a + b₂ trained on three points. Each pixel runs that net's training for up to T steps from a fixed initialisation; if the loss crosses a blow-up threshold we colour by how fast it escaped (the escape time), exactly as a Mandelbrot renderer does. Survive all T steps and the pixel is "trainable" — the dark interior.

Nobody designed the filigree. It is what the edge of stability looks like.

We measured the boundary directly: as the render resolution doubles, the boundary's length grows about 3.3× rather than 2×, putting its fractal dimension near 1.7 — and fresh islands of trainability keep appearing at every zoom level. A smooth curve does neither.

Why it happens

The gradient-descent map is nonlinear (the tanh sees to that). Near the edge of stability — learning rates close to the critical 2/curvature — the linearised step has eigenvalues hugging the unit circle, the dynamics turn chaotic, and small hyperparameter changes are amplified without limit. That sensitive dependence, folded over and over by the nonlinearity, is the engine of self-similarity, just as it is for the quadratic map behind Mandelbrot. Click a pixel on the explore tab to watch a single run's loss curve hug the boundary, converge, or detonate.

A field of nanomodels

Flip the "what lives at each point" toggle and the curve-fitter becomes a genuine — if absurdly tiny — language model: a recurrent net h_t = tanh(W_x·x_t + W_h·h_{t-1} + b) that reads a short string one character at a time and predicts the next. Now every dark pixel is a learning rate at which this nanomodel successfully learns the string — click one and it tells you what it generates (it really does say "abcabc…"); click the bright chaos and it babbles, then blows up.

Two things had to be true to make it fractal, and both are lessons. First, recurrence: a memoryless next-character model has a smooth, boring trainability boundary — it's the hidden state fed back through W_h that makes the dynamics chaotic enough to fracture. Second, the loss: trained with ordinary cross-entropy, the model never diverges — softmax's gradients are bounded, so it's unconditionally stable (we checked out to learning rate 800). The explosive, fractal-making instability needs an unbounded loss, so the nanomodels are trained with squared error toward the next character. The edge of trainability is real, but you only see its fractal teeth when the loss can run away.

Doing it on the GPU

Every pixel is an independent training run with no shared state — which makes this perfectly parallel and a natural fit for a WebGPU compute shader: one thread per pixel, each iterating the tiny network and writing its escape colour to a buffer that a fullscreen pass blits. This page uses that path where WebGPU is available and falls back to a progressive CPU renderer otherwise, so it works everywhere; the GPU lets the resolution, step count, and network size climb toward the ranges in the original paper. The compute shader runs the identical update to the CPU reference, so both draw the same fractal.

sources