06
Neural Network — watch it learn
The capstone: a 2-12-12-1 MLP — 205 parameters — trained entirely on the GPU. Each frame is one training step: base loss, a forward-difference gradient (one thread per parameter, no reduction), a momentum-SGD update over ping-ponged weights, then a kernel paints the decision field at every pixel. The CPU only reads back the loss HUD.
LIVE · GPU
This demo needs WebGPU
click canvas to re-init · dataset cycles
epoch — · loss — · accuracy —
fps —
use system.collections.array
use system.math
use system.io
// Network shape: 2 inputs -> 12 -> 12 -> 1 output.
const IN = 2
const H1 = 12
const H2 = 12
// Weight layout inside the flat parameter buffer (205 floats):
// W1 [24] @0 b1 [12] @24 W2 [144] @36 b2 [12] @180 W3 [12] @192 b3 @204
const W1 = 0
const B1 = 24
const W2 = 36
const B2 = 180
const W3 = 192
const B3 = 204
const PARAMS = 205
// One slot past the weights carries the epoch count, so it ping-pongs with them
// and resets exactly when they do. The host reads it out for the HUD.
const EPOCH = PARAMS
const SLOTS = PARAMS + 1
// Optimizer steps taken per displayed frame. The reference host takes three, and
// the convergence rate a viewer sees is the product of this and the step size.
const EPOCHS_PER_FRAME = 3
// How long each shape stays on screen, and how many shapes the demo cycles
// through. The net converges in well under a second, so a long hold would leave
// the field frozen for most of the loop; this holds the settled boundary just
// long enough to read before re-forming on the next shape.
const SECS_PER_MODE = 6
const MODES = 3
// Dataset: 100 points, each (x, y, label in {0, 1}); PER_ARM points per class.
// Training runs on a single GPU thread that walks the whole batch each epoch, so
// the per-frame cost is linear in the sample count; 100 points keeps the frame
// well inside budget while still defining two clean spiral arms (and the rings /
// XOR shapes). Fewer points, not fewer pixels, is what buys frame rate here — the
// GPU saturates its lanes at any sane resolution, so the field render is nearly
// free next to the serial training walk.
const SAMPLES = 100
const PER_ARM = 50
const DATA = SAMPLES * 3
// Decision-field surface: a 960×540 (16:9) RGBA image. The render pass evaluates
// the whole net and scans every data point per pixel, so the pixel count is the
// dominant per-frame cost; 960×540 keeps the field and dots crisp (the browser
// scales the canvas to fit) while staying well within frame budget.
const CW = 960
const CH = 540
const PIXELS = CW * CH
const PAINT = PIXELS * 4
// Weights (plus the trailing epoch slot) and velocities, ping-ponged a -> b.
gpu var w_a = Array<f32, SLOTS>()
gpu var w_b = Array<f32, SLOTS>()
gpu var v_a = Array<f32, PARAMS>()
gpu var v_b = Array<f32, PARAMS>()
// Dataset points and the [loss, accuracy, epoch] HUD stats.
gpu var data = Array<f32, DATA>()
gpu var stats = Array<f32, 3>()
// Decision-field paint output (960×540 RGBA).
gpu var paint = Array<f32, PAINT>()
fn rand_unit(key i32) f32
let h = hash_u32(key as u32)
return (h as f32) / (4294967295.0 as f32)
// Initial spread of parameter `i`, small and per-layer in the Xavier-ish manner
// the reference net uses. Biases start at exactly zero, so they scale by 0.
fn init_scale(i i32) f32
if i < B1: return 0.65
if i >= W2 and i < B2: return 0.5
if i >= W3 and i < B3: return 0.68
return 0.0
// Which shape is on screen at time `t`: spiral (0), then rings (1), then xor (2).
fn mode_at(t f32) i32
return (floor(t / (SECS_PER_MODE as f32)) as i32) % MODES
// Seed weights from that spread; zero the velocities.
forall i in 0..PARAMS
w_a[i] = ((rand_unit(i + 1) * 2.0 - 1.0) * init_scale(i)) as f32
v_a[i] = 0.0
// Seed a two-arm spiral: points 0..PER_ARM are class 0, the rest class 1.
forall s in 0..SAMPLES
let c = s / PER_ARM
let i = s - c * PER_ARM
let t = ((i as f32) / (PER_ARM as f32)) as f32
let r = (0.13 + 0.74 * t) as f32
let th = (t * 4.4 + (c as f32) * (3.14159265 as f32)) as f32
let nx = ((rand_unit(s * 2 + 1) - 0.5) * 0.07) as f32
let ny = ((rand_unit(s * 2 + 2) - 0.5) * 0.07) as f32
data[s * 3] = (r * (cos(th) as f32) + nx) as f32
data[s * 3 + 1] = (r * (sin(th) as f32) + ny) as f32
data[s * 3 + 2] = c as f32
// Clear the paint buffer. This top-level 2-D pass over the exact display extent
// tells the web-gpu backend the canvas is 960×540 (16:9); the per-frame render
// pass is 1-D, which alone would leave the canvas shape ambiguous.
forall px, py in 0..CW, 0..CH
let base = (py * CW + px) * 4
paint[base] = 0.012
paint[base + 1] = 0.02
paint[base + 2] = 0.046
paint[base + 3] = 1.0
gpu frame
// Pass 1: several full-batch backprop epochs. A single thread takes the
// previous frame's weights from `w_a`/`v_a`, and for each epoch forwards
// every sample, backpropagates the cross-entropy error into the exact
// gradient `g`, and takes one momentum SGD step. The result is published to
// `w_b`/`v_b`, which the host ping-pongs back each frame. Weights, gradient
// and activations are all per-invocation scratch, which is what makes the
// whole training loop expressible on-device: no cross-thread reduction and
// no host round-trip.
forall t in 0..1
// The epochs run against per-invocation copies of the weights. A pass may
// not both read and write a device buffer — the compiler rejects that as a
// race, and rightly, since it cannot know this pass is a single lane — so
// the ping-pong happens once at the edges of the frame rather than once
// per epoch.
var w = Array<f32, PARAMS>()
var vel = Array<f32, PARAMS>()
var c = 0
while c < PARAMS
w[c] = w_a[c]
vel[c] = v_a[c]
c = c + 1
var g = Array<f32, PARAMS>()
var h1 = Array<f32, H1>()
var h2 = Array<f32, H2>()
var d1 = Array<f32, H1>()
var d2 = Array<f32, H2>()
var epoch = 0
while epoch < EPOCHS_PER_FRAME
// The gradient accumulates across the batch, so it must start each
// epoch at zero rather than carrying the previous epoch's sum.
var z = 0
while z < PARAMS
g[z] = 0.0
z = z + 1
var s = 0
while s < SAMPLES
let x0 = data[s * 3]
let x1 = data[s * 3 + 1]
let lbl = data[s * 3 + 2]
// Forward pass.
var j = 0
while j < H1
let a1 = (w[W1 + j * IN] * x0 + w[W1 + j * IN + 1] * x1 + w[B1 + j]) as f32
h1[j] = tanh(a1) as f32
j = j + 1
j = 0
while j < H2
var a2 = w[B2 + j]
var k = 0
while k < H1
a2 = (a2 + w[W2 + j * H1 + k] * h1[k]) as f32
k = k + 1
h2[j] = tanh(a2) as f32
j = j + 1
var y = w[B3]
j = 0
while j < H2
y = (y + w[W3 + j] * h2[j]) as f32
j = j + 1
let p = (1.0 / (1.0 + (exp(0.0 - y) as f32))) as f32
// Backward pass: accumulate the exact gradient into `g`.
let ds = (p - lbl) as f32
g[B3] = (g[B3] + ds) as f32
j = 0
while j < H2
g[W3 + j] = (g[W3 + j] + ds * h2[j]) as f32
d2[j] = (ds * w[W3 + j] * (1.0 - h2[j] * h2[j])) as f32
j = j + 1
j = 0
while j < H2
g[B2 + j] = (g[B2 + j] + d2[j]) as f32
var k = 0
while k < H1
g[W2 + j * H1 + k] = (g[W2 + j * H1 + k] + d2[j] * h1[k]) as f32
k = k + 1
j = j + 1
var k = 0
while k < H1
var acc = 0.0
j = 0
while j < H2
acc = (acc + d2[j] * w[W2 + j * H1 + k]) as f32
j = j + 1
d1[k] = (acc * (1.0 - h1[k] * h1[k])) as f32
k = k + 1
k = 0
while k < H1
g[B1 + k] = (g[B1 + k] + d1[k]) as f32
g[W1 + k * IN] = (g[W1 + k * IN] + d1[k] * x0) as f32
g[W1 + k * IN + 1] = (g[W1 + k * IN + 1] + d1[k] * x1) as f32
k = k + 1
s = s + 1
// Momentum SGD (lr 0.25, momentum 0.9), gradient averaged over the
// batch. No weight decay: the spiral needs the sharp, high-curvature
// boundary that L2 would smooth.
let sc = (0.25 / (SAMPLES as f32)) as f32
var i = 0
while i < PARAMS
let nv = (0.9 * vel[i] - sc * g[i]) as f32
vel[i] = nv
w[i] = (w[i] + nv) as f32
i = i + 1
epoch = epoch + 1
// Publish the frame's result. Clicking the canvas restarts training from a
// fresh net instead of the one just refined; the keys are offset by the
// frame clock so each restart draws different weights rather than
// replaying the first run, and the offset wraps well inside 32-bit range.
// A dataset change restarts it for the same reason a trainer would not
// fine-tune across unrelated tasks: the converged weights are large enough
// to saturate every tanh, leaving no gradient to fit the new shape with.
let before = frame.time - frame.dt
let prev = before if before > 0.0 else 0.0
let restart = frame.clicked or mode_at(frame.time) != mode_at(prev)
let key_base = ((((frame.time * 97.0) as i32) % 9973) * PARAMS) as i32
var o = 0
while o < PARAMS
if restart
w_b[o] = ((rand_unit(key_base + o + 1) * 2.0 - 1.0) * init_scale(o)) as f32
v_b[o] = 0.0
else
w_b[o] = w[o]
v_b[o] = vel[o]
o = o + 1
// Epochs taken by the net now on screen: it counts from zero for each
// freshly seeded one, so it measures this shape's training, not the
// session's.
let epochs_before = 0.0 if restart else w_a[EPOCH]
w_b[EPOCH] = (epochs_before + (EPOCHS_PER_FRAME as f32)) as f32
// Pass 2: post-step loss + accuracy over the batch, plus the epoch count
// the train pass just advanced, for the HUD readback.
forall t in 0..1
var h1 = Array<f32, H1>()
var h2 = Array<f32, H2>()
var loss = 0.0
var correct = 0.0
var s = 0
while s < SAMPLES
let x0 = data[s * 3]
let x1 = data[s * 3 + 1]
let lbl = data[s * 3 + 2]
var j = 0
while j < H1
let a1 = (w_b[W1 + j * IN] * x0 + w_b[W1 + j * IN + 1] * x1 + w_b[B1 + j]) as f32
h1[j] = tanh(a1) as f32
j = j + 1
j = 0
while j < H2
var a2 = w_b[B2 + j]
var k = 0
while k < H1
a2 = (a2 + w_b[W2 + j * H1 + k] * h1[k]) as f32
k = k + 1
h2[j] = tanh(a2) as f32
j = j + 1
var y = w_b[B3]
j = 0
while j < H2
y = (y + w_b[W3 + j] * h2[j]) as f32
j = j + 1
let p = clamp(1.0 / (1.0 + (exp(0.0 - y) as f32)), 0.000001, 0.999999) as f32
loss = (loss - (lbl * (log(p) as f32) + (1.0 - lbl) * (log(1.0 - p) as f32))) as f32
let hit = 1.0 if (p - 0.5) * (lbl - 0.5) > 0.0 else 0.0
correct = (correct + hit) as f32
s = s + 1
stats[0] = (loss / (SAMPLES as f32)) as f32
stats[1] = (correct / (SAMPLES as f32)) as f32
stats[2] = w_b[EPOCH]
// Pass 3: regenerate the dataset for the current mode. A timer cycles the
// shape every SECS_PER_MODE seconds (spiral -> rings -> xor), so the field
// visibly re-forms for each one from the net the train pass just reseeded.
// Within a mode the layout is a deterministic function of the point index,
// so it is stable across frames and the net can learn it. `train` stays the
// first pass (it owns the weight ping-pong); this pass feeds the render below
// and next frame's training, so the shape only lags a frame at a boundary.
forall s in 0..SAMPLES
let mode = mode_at(frame.time)
var dx = 0.0
var dy = 0.0
var lbl = 0.0
if mode == 0
// Two-arm spiral: points 0..PER_ARM are class 0, the rest class 1.
let c = s / PER_ARM
let i = s - c * PER_ARM
let tt = ((i as f32) / (PER_ARM as f32)) as f32
let r = (0.13 + 0.74 * tt) as f32
let th = (tt * 4.4 + (c as f32) * (3.14159265 as f32)) as f32
dx = (r * (cos(th) as f32) + (rand_unit(s * 2 + 1) - 0.5) * 0.07) as f32
dy = (r * (sin(th) as f32) + (rand_unit(s * 2 + 2) - 0.5) * 0.07) as f32
lbl = c as f32
if mode == 1
// Concentric rings: an inner disk (class 1) in an outer annulus,
// split evenly across the batch.
let inner = s * 2 < SAMPLES
let rad = (sqrt(rand_unit(100003 + s * 3 + 1)) * 0.34) as f32 if inner else (0.58 + rand_unit(100003 + s * 3 + 1) * 0.27) as f32
let ang = (rand_unit(100003 + s * 3 + 2) * (6.2831853 as f32)) as f32
dx = (rad * (cos(ang) as f32)) as f32
dy = (rad * (sin(ang) as f32)) as f32
lbl = 1.0 if inner else 0.0
if mode == 2
// XOR: class 1 where x and y share a sign, class 0 otherwise.
let px = ((rand_unit(200003 + s * 3 + 1) * 2.0 - 1.0) * 0.85) as f32
let py = ((rand_unit(200003 + s * 3 + 2) * 2.0 - 1.0) * 0.85) as f32
dx = px
dy = py
lbl = 1.0 if px * py > 0.0 else 0.0
data[s * 3] = dx
data[s * 3 + 1] = dy
data[s * 3 + 2] = lbl
// Pass 4: decision field. Evaluate the trained net per pixel, shade by
// class and confidence, add outward-flowing confidence bands and a pulsing
// decision-boundary seam, then overlay the data points.
forall px in 0..PIXELS
let ix = px % CW
let iy = px / CW
// Aspect-corrected input square: y spans [-1, 1], x widens by 16:9.
let gx = ((((ix as f32) + 0.5) / (CW as f32)) - 0.5) * 2.0 * ((CW as f32) / (CH as f32)) as f32
let gy = (0.5 - (((iy as f32) + 0.5) / (CH as f32))) * 2.0 as f32
var h1 = Array<f32, H1>()
var h2 = Array<f32, H2>()
var j = 0
while j < H1
let a1 = (w_b[W1 + j * IN] * gx + w_b[W1 + j * IN + 1] * gy + w_b[B1 + j]) as f32
h1[j] = tanh(a1) as f32
j = j + 1
j = 0
while j < H2
var a2 = w_b[B2 + j]
var k = 0
while k < H1
a2 = (a2 + w_b[W2 + j * H1 + k] * h1[k]) as f32
k = k + 1
h2[j] = tanh(a2) as f32
j = j + 1
var y = w_b[B3]
j = 0
while j < H2
y = (y + w_b[W3 + j] * h2[j]) as f32
j = j + 1
let p = (1.0 / (1.0 + (exp(0.0 - y) as f32))) as f32
// Class side (blue below 0.5, gold above) lifted off a deep navy base by
// the squared confidence, so the uncertain band around the boundary stays
// dark and the settled regions glow.
let conf = (abs(p - 0.5) * 2.0) as f32
let sr = 0.16 if p < 0.5 else 1.0
let sg = 0.31 if p < 0.5 else 0.84
let sb = 0.82 if p < 0.5 else 0.24
let fill = (0.10 + 0.60 * conf * conf) as f32
var cr = (0.012 + sr * fill) as f32
var cg = (0.02 + sg * fill) as f32
var cb = (0.046 + sb * fill) as f32
// Confidence iso-contours flowing outward from the boundary: the field
// keeps moving even once the weights have nearly settled.
let bands = (abs(fract(p * 9.0 - frame.time * 0.25) - 0.5)) as f32
let band = (smoothstep(0.44, 0.5, bands) * 0.10) as f32
cr = (cr + sr * band) as f32
cg = (cg + sg * band) as f32
cb = (cb + sb * band) as f32
// Decision boundary: a bright seam at p = 0.5, breathing on a slow pulse.
let e = (abs(p - 0.5) / 0.035) as f32
let edge = (exp(0.0 - e * e)) as f32
let pulse = (0.72 + 0.28 * (sin(frame.time * 2.4) as f32)) as f32
cr = (cr + 0.90 * edge * 0.55 * pulse) as f32
cg = (cg + 0.95 * edge * 0.55 * pulse) as f32
cb = (cb + 1.0 * edge * 0.55 * pulse) as f32
// Data-point overlay: small filled dots colored by their true label, each
// ringed by a dark rim so they read against the field. Samples outside the
// pixel's x band are rejected on one load, before the distance math.
var s = 0
while s < SAMPLES
let dx = (gx - data[s * 3]) as f32
if dx * dx < 0.00032
let dy = (gy - data[s * 3 + 1]) as f32
let dd = (dx * dx + dy * dy) as f32
if dd < 0.00032
let rr = (sqrt(dd)) as f32
let lbl = data[s * 3 + 2]
let pr = 0.45 if lbl < 0.5 else 1.0
let pg = 0.62 if lbl < 0.5 else 0.88
let pb = 1.0 if lbl < 0.5 else 0.40
let disc = (1.0 - smoothstep(0.0155, 0.0178, rr)) as f32
cr = (cr + (0.01 - cr) * disc) as f32
cg = (cg + (0.02 - cg) * disc) as f32
cb = (cb + (0.05 - cb) * disc) as f32
let core = (1.0 - smoothstep(0.0100, 0.0148, rr)) as f32
cr = (cr + (pr - cr) * core) as f32
cg = (cg + (pg - cg) * core) as f32
cb = (cb + (pb - cb) * core) as f32
s = s + 1
let base = px * 4
paint[base] = cr
paint[base + 1] = cg
paint[base + 2] = cb
paint[base + 3] = 1.0
Related docs: why the CPU only reads back the loss HUD and the gpu frame that runs one training step per tick.
RUN IT YOURSELF
From this page to your own GPU
Four steps. You'll need a WebGPU-capable browser (Chrome or Edge 113+, or Safari 18+).
-
1
Install Miri
Build the compiler from source (full install guide):
git clone https://github.com/miri-lang/miri.git cd miri && cargo build --releaseThe binary lands at
target/release/miri. -
2
Grab the program
Hit copy program above and save it as
neural.mi. -
3
Compile it to WebGPU
miri build neural.mi --target web-gpu --out neural-webOut comes a self-contained
index.html— the runtime and every compiled WGSL kernel are inlined. -
4
Open it
Double-click
neural-web/index.html. It runs straight fromfile://— same interaction as the preview above.
No browser needed to try it: miri run neural.mi runs the
same kernels on your local GPU through Metal, Vulkan or DX12. Same language, same code, three
backends and a browser — that's the point.