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
Convolutional network
3 pagesSmall 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.
Vision Transformer (ViT)
nextCut 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.
- 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.
U-Net, encoder and decoder
plannedAn 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?".
- 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.
CoAtNet, convolution then attention
plannedConvolution 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.
- 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.
ConvNeXt, a modernised CNN
plannedA 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.
- 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.