Training GNNs on WebGPU, and the Ways the Numbers Lie

Hand-written WGSL for GCN, GraphSAGE and GIN in the browser. The kernels were the easy part. The hard parts were a leakage trap that makes link prediction look nearly perfect, an oversmoothing metric that was wrong before the model was, and a test suite that had been skipping itself on every machine.

Author

Benedict Thekkel

Published

September 16, 2026

model_playground runs ML models in the browser, on the visitor’s GPU. It is live at playground.benedictthekkel.com. Most of it deploys pretrained models through Transformers.js and ONNX Runtime Web, and none of that is my kernel code. This post is about the part that is: the graph neural network pages, built between 2026-09-12 and 2026-09-16, where training runs on WGSL compute shaders I wrote by hand.

One shader for three architectures

GCN, GraphSAGE and GIN look different on paper. In the forward pass they are the same operation: for each node, a scaled sum of its own features and its neighbours’ features. They differ only in the scales.

So there is one aggregation shader, over the adjacency in CSR form, one invocation per node:

for (var f = 0u; f < dims.nFeat; f = f + 1u) {
  var acc = selfScale * x[base + f];
  for (var e = start; e < end; e = e + 1u) {
    let j = colIdx[e];
    acc = acc + beta[j] * x[j * dims.nFeat + f];
  }
  out[base + f] = outScale * acc;
}

GCN’s symmetric normalisation, SAGE’s mean and GIN’s sum are all choices of selfScale, beta and outScale. And because the self-looped adjacency is symmetric, the backward pass is the same dispatch with two scale vectors swapped. One shader, both directions, three models.

Around it, five more small shaders: a naive matmul in 16x16 workgroups with no tiling, elementwise ops, transpose, scale, and a point-cloud renderer used elsewhere. 275 lines of WGSL in all. GAT is the exception: its attention coefficients are learned per edge, so it runs on a CPU propagator.

The dataset ships as CSR too. Cora’s 2,708 nodes and 1,433 features are 15.5 MB as dense float32; as a sparse structure it is 161 KB.

How you know the kernels are right

Two checks, at different layers:

Kernel parity. The GPU aggregation runs against a CPU reference implementation, in a real browser, and the worst element-wise error must be under 1e-5. Dropping the self-loop as a deliberate fault sends it to 0.87, so the test can actually fail.
Finite-difference gradient checks on every hand-written backward pass: GCN, SAGE, GIN and GAT at one, two and three layers, with and without dropout. They are evaluated at a point with non-zero bias, so no pre-activation sits exactly on ReLU’s kink, where the numerical gradient is meaningless.

While wiring those up I found that every WebGPU end-to-end test in the repo had been silently skipping itself, on every machine. The suite probed for navigator.gpu before navigating, from about:blank. WebGPU is only exposed in a secure context, and about:blank in that harness was not one, so the probe always said “no GPU” and every spec politely skipped. Green CI, zero coverage. A skip is a result too, and it is worth asserting on how many tests actually ran.

Oversmoothing: the metric was wrong first

Stack enough GCN layers and every node’s representation converges toward the same vector. That is oversmoothing, and the page has a slider for depth so you can watch it happen.

The first measure I used was the obvious one: mean pairwise cosine similarity over all node pairs. If representations converge, that should climb with depth. It came out as

0.720, 0.643, 0.663, 0.579, 0.677

which is not a trend. The metric was answering the wrong question. Oversmoothing is a local phenomenon first: connected nodes become indistinguishable. Unrelated nodes on opposite sides of the graph can stay far apart for a long time while every neighbourhood collapses internally.

The replacement measures similarity between adjacent nodes, on row-normalised representations, which is the Dirichlet energy up to a constant. With it, the picture is monotonic and matches the accuracy:

GCN layers	Test accuracy	Neighbour similarity
2	0.776	0.947
4	0.687	0.958
6	0.496	0.973
8	0.518	0.978

Accuracy is itself not monotonic (8 layers edges 6), which is a reminder that one run per depth is noise-sensitive. The similarity column is the one that moves cleanly.

GIN at 8 layers is a different failure wearing the same symptom: 0.064 accuracy. That is not smoothing. GIN sums neighbours without normalising, the sums overflow, and ReLU zeroes everything downstream. Every representation is dead. So the page reports the dead fraction separately and names the failure when more than half the units are dead, instead of letting a single accuracy number blame the wrong cause.

Link prediction: the leak that improves the score

Link prediction hides some edges, trains on the rest, and asks the model to score the hidden edges above random non-edges. The trap is that a GNN encoder reads the graph. If a held-out edge is still in the adjacency the encoder aggregates over, the model is not predicting it. It is looking at it. The expected symptom is an AUC around 0.99, a number that looks like success.

So the split was written to make that impossible before any model was trained on it. Four things have to be true, and the first one is the only one most write-ups mention:

Remove both directed copies of each held-out edge. CSR stores an undirected edge twice.
Recompute the degrees from the training graph. GCN normalises by \(D^{-1/2}\). A degree vector computed on the full graph still encodes the held-out edge, as a count, in every message that touches its endpoints.
Never hold out an edge that would isolate a node. On Cora, up to 485 degree-1 nodes would be stranded with nothing to aggregate. Those edges stay in training.
Sample negatives against the full graph, so a “negative” is never secretly a held-out positive.

for (const e of order) {
  if (nTest >= wantTest && nVal >= wantVal) break;
  const u = edges[2 * e]; const v = edges[2 * e + 1];
  // Removing this would leave an endpoint with nothing to aggregate over.
  if (degree[u] <= 1 || degree[v] <= 1) { rescued++; continue; }
  degree[u]--; degree[v]--;
  // ...
}
// Negatives are drawn against the *full* graph
const testNeg = sampleNegatives(rowPtr, colIdx, nNodes, nTest, rand);

And the end-to-end test asserts a band, not a floor: test AUC above 0.85 and below 0.985. A score that high on Cora is not a better model. It is the leak coming back, and the test fails on it.

The hyperparameters mattered more than I expected. Reusing the node-classification recipe gave 0.737. The graph auto-encoder recipe gave 0.879, then 0.906 at 250 epochs, and 0.925 with 0.2 dropout, which is what ships. For reference, the published GAE result on Cora is 0.910, so this is in the right place rather than suspiciously above it. GPU and CPU paths agree to within 0.003.

Takeaway

None of the three stories is about WGSL. The kernels were checked by parity and gradient tests and were right. Every real problem was in what the numbers meant:

Number	What was wrong
Pairwise cosine vs depth	Measured a global property of a local failure
GIN accuracy at depth 8	Blamed smoothing for overflow and dead ReLUs
Link-prediction AUC	Would reward leakage, so it is bounded above in the test
Green WebGPU test suite	Nothing had run

The habit that catches all four is the same: before trusting a metric, ask what the wrong answer would look like, and make sure a test can tell the two apart.