Visualising AI

Interactive explorers for deep-learning architectures. Each one takes the boxes of an architecture diagram and shows what they compute: every cell, kernel, weight and gradient can be hovered, and the arithmetic is written out. The networks are small on purpose, so that every number fits on screen.

The convolutional network is built. The Vision Transformer is next, then U-Net, CoAtNet and ConvNeXt.

How to read every page

Yellow is the kernel window as it moves forward over its input.
Green traces backwards: what a hovered cell sees, or the two sources of a weight's gradient.
Blue to orange is negative to positive. In 3D, cube size is magnitude.
termA dotted word has its definition on hover; every page ends with its Terms.
16×16 digit → conv → pool → conv → pool → average → 10 classes. Depth is channels.
01

Convolutional network

3 pages

Small learned filters slide over the image. Stacked, with pooling in between, they turn pixels into edges, edges into shapes, and shapes into a class.

It asks, at every position: does the pattern I learned appear here? Each layer sees a slightly wider patch than the one before, so the view grows a few pixels at a time.

Patches → tokens → attention, every token to every other → class token → 10 classes. One token's attention in green.
02

Vision Transformer (ViT)

next

Cut the image into patches, turn each patch into a vector called a token, and let every token look at every other one through attention.

Where a convolution's view grows a few pixels per layer, attention sees the whole image from the first layer and learns where to look. The price: its cost grows with the square of the number of patches, which is why patches are large.

What the page will show
  • Patch embedding as a convolution whose stride equals its size: the CNN pages' kernel, jumping instead of sliding.
  • Position added to each token, since attention by itself does not know where a patch came from.
  • One attention head written out: queries against keys, softmax, a weighted sum of values.
  • The attention map as a receptive field that is global in one step.

Builds on the CNN pages: stride, kernels, softmax.

Encoder down, decoder up, skip connections across. The output is a label per pixel.
03

U-Net, encoder and decoder

planned

An encoder shrinks the image into features, a decoder grows them back to full size, and skip connections hand the fine detail straight across.

The CNN here ends in 10 numbers for the whole image. A U-Net ends in an image: a label for every pixel, which is segmentation. The question changes from "what is this?" to "which pixels are it?".

What the page will show
  • The contracting path: the convolutions and pooling of the CNN pages, adding channels as the map shrinks.
  • Up-sampling by transposed convolution, the mirror image of a strided one.
  • Skip connections joined as extra channels, and what the output loses without them.
  • One output pixel's receptive field, reaching down through the bottleneck and across the skips.

Builds on the CNN pages: convolution, pooling, receptive field.

Convolution stages (C) while the map is large, attention stages (T) once it is small. Yellow: what one output reads.
04

CoAtNet, convolution then attention

planned

Convolution stages first, attention stages after: cheap local detail while the map is large, then global context once it is small.

It asks where in the network to stop looking locally and start looking everywhere, and what that switch costs. Its diagram is the one that started this project.

What the page will show
  • MBConv stages: expand the channels, a depthwise convolution, squeeze and excite, project back.
  • The switch from local to global, with the receptive field before and after it.
  • Attention with a relative position bias, which keeps a little of convolution's sense of neighbourhood.
  • The cost of each stage, which explains the order.

Builds on the CNN and ViT pages.

Patchify stem, then one block: 7×7 depthwise, expand ×4, project back, add the input.
05

ConvNeXt, a modernised CNN

planned

A plain convolutional network, rebuilt step by step with the design choices that made transformers work, until it matches them.

It asks whether attention was the point, or the recipe around it. Each change is small and can be switched on and off.

What the page will show
  • The patchify stem: a 4×4 convolution with stride 4, the ViT patch embedding again.
  • 7×7 depthwise convolution: one slice per channel and no sum across channels, against the 3×3×C kernel of the CNN pages.
  • 1×1 convolutions as the channel-mixing step, in an inverted bottleneck that widens four times and narrows back.
  • LayerNorm, GELU and the residual connection, and what each changes.

Builds on the CNN and ViT pages.