Making ONNX observable: residual streams for in-browser interpretability

Running an LLM in the browser via WebGPU (transformers.js + ONNX Runtime) is delightful — and completely opaque. A compiled ONNX graph exposes only its declared outputs: logits and KV-cache. There is nowhere to hang a forward hook. The hidden states exist for microseconds inside fused kernels, then vanish.

The trick: they're already computed

Optimized exports fuse each residual add into a SkipSimplifiedLayerNormalization node — which has an optional fourth output, input_skip_bias_sum: the residual stream itself. Layer i's residual is exactly that output on layer i+1's input-layernorm. The final layer's sum is one undeclared optional output away on the final norm node.

So the "surgery" is almost embarrassingly small: map the 30 tensor names, bridge each through a Cast to fp32 with a stable name (hidden.0…hidden.29), and append them to graph.output. Zero extra compute — the tensors were always flowing; they just weren't visible.

hidden.i = layers[i+1].input_layernorm.input_skip_bias_sum
validation vs PyTorch hooks:
  ✓ next-token argmax parity
  ✓ residual norms track reference (worst drift 12.1% — q4f16)
  ✓ 30 outputs, [batch, seq, 576]

What it unlocks

In model-graph, a patched model.forward reads hidden.* each step and computes residual norms, a seeded 3D projection (the trajectory view), and top-k — all client-side, from real tensors, on WebGPU. Serverless interpretability: the browser fetches weights from Hugging Face and inspects itself.

The honest boundary: per-block writes (attention/MLP separately) and per-head signals need more declared outputs — same technique, more tensor-name mapping. The export script and validation live in the repo (export_observable.py); try the UI at model-graph.com.