CNNs & Computer Vision

A convolution is a sliding dot product. You take a small filter (often 3x3 or 5x5), slide it across an input image, and at each position compute the dot product of the filter weights with the underlying patch. The result is a new 2D array called a feature map. The filter is shared across every spatial location, which is what gives the convolution its first defining property: translation equivariance. Shift the input one pixel right; the output also shifts one pixel right. The network has not had to learn the same pattern twice.

The second defining property is locality. Each output pixel only depends on a small neighbourhood of input pixels. Far-apart pixels can only interact by stacking multiple layers, which is how the receptive field grows. We will come back to this.

The third property is parameter sharing, which is locality plus translation equivariance taken seriously. A 3x3 filter on a 224x224 input is 9 weights, regardless of image size. A fully-connected layer on the same input would have 50,176 weights per output unit. This is the prior that makes vision tractable at scale.

A single conv layer in PyTorch takes a multi-channel input $(C_{\text{in}}, H, W)$ and produces a multi-channel output $(C_{\text{out}}, H', W')$. The weight tensor has shape $(C_{\text{out}}, C_{\text{in}}, k_H, k_W)$. Each output channel is its own filter, looking at all input channels at once.

conv = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, stride=1, padding=1)
out = conv(x)   # x: (N, 3, H, W) -> (N, 32, H, W)

def conv2d_naive(x, w, b, stride=1, pad=0):
    # x: (N, C_in, H, W). w: (C_out, C_in, kH, kW). b: (C_out,).
    N, C_in, H, W = x.shape
    C_out, _, kH, kW = w.shape
    H_out = (H + 2*pad - kH) // stride + 1
    W_out = (W + 2*pad - kW) // stride + 1
    x_padded = np.pad(x, ((0,0),(0,0),(pad,pad),(pad,pad)))
    out = np.zeros((N, C_out, H_out, W_out))
    for n in range(N):
        for c_out in range(C_out):
            for i in range(H_out):
                for j in range(W_out):
                    patch = x_padded[n, :, i*stride:i*stride+kH, j*stride:j*stride+kW]
                    out[n, c_out, i, j] = np.sum(patch * w[c_out]) + b[c_out]
    return out

That nested loop is what nn.Conv2d does, with all the loops fused inside a cuDNN kernel and the patch-extraction step done via im2col so the inner operation becomes a single matmul. Read this version once, then never write it again.

Two parameters you set on every conv layer. Both come from one question: how should the output shape compare to the input shape?

Padding is zero-padding the input borders before convolving. Without padding, a 3x3 filter on a 5x5 input gives a 3x3 output. You lose two pixels on each side per layer. Stack ten layers and you lose twenty pixels. With padding $P = (k-1)/2$ for an odd kernel size $k$, the output keeps the same spatial size as the input. This is called "same" padding. The alternative is "valid" (no padding).

Stride is how many pixels the filter jumps each step. Stride 1 is dense. Stride 2 halves the spatial dimensions, which is how you downsample without a pooling layer (modern ResNets do this).

The output shape formula, which you will use approximately five hundred times in this chapter:

Memorise it. Print it on a sticky note.

# Same padding, stride 1: keeps spatial dims
nn.Conv2d(64, 128, kernel_size=3, padding=1)

# Stride-2 conv: halves spatial dims, no pool needed
nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)

# Valid: shrinks the output
nn.Conv2d(64, 128, kernel_size=3, padding=0)

Pooling is a fixed (non-learned) operation that downsamples a feature map by taking the max or the mean over a small window. The standard choice is 2x2 max pooling with stride 2, which halves both spatial dimensions and keeps only the highest activation per window.

Why max specifically? Because the highest activation is the strongest evidence that a feature is present somewhere in that window. The exact spatial location is discarded. This is the move from translation equivariance (convolutions) to a degree of translation invariance. A cat shifted three pixels still triggers the same post-pool feature map.

Pooling is also a parameter-free shrink: no weights, just a deterministic reduction. Modern architectures have started replacing max pooling with strided convolution (ResNets, the entire post-2015 lineage), on the grounds that learned downsampling beats fixed downsampling. Both are still in production code.

pool = nn.MaxPool2d(kernel_size=2, stride=2)
out = pool(x)   # (N, C, H, W) -> (N, C, H/2, W/2)

def max_pool_2d(x, k=2, stride=2):
    N, C, H, W = x.shape
    H_out, W_out = (H - k)//stride + 1, (W - k)//stride + 1
    out = np.zeros((N, C, H_out, W_out))
    for i in range(H_out):
        for j in range(W_out):
            patch = x[:, :, i*stride:i*stride+k, j*stride:j*stride+k]
            out[:, :, i, j] = patch.reshape(N, C, -1).max(axis=-1)
    return out

The receptive field of a unit is the region of the input it depends on. For a single 3x3 conv layer, a unit's receptive field is 3x3. For two stacked 3x3 layers, each unit depends on a 5x5 region (the second layer sees a 3x3 grid of units, each of which sees its own 3x3 input patch, which overlap). For three stacked 3x3 layers, the receptive field is 7x7. Add a 2x2 stride-2 pool and the next layer's receptive field is doubled in each dimension.

This is one of the few CNN concepts where the arithmetic is load-bearing. Two stacked 3x3 convs cover the same receptive field as one 5x5 conv, but with fewer parameters (18 versus 25) and one more non-linearity in between. This is the entire argument behind VGG's design choice to use only 3x3 filters.

A useful piece of intuition: deep layers in a trained CNN look at large image regions, but they don't literally attend uniformly across that region. The effective receptive field is much smaller than the theoretical one, concentrated near the center. This is empirical (Luo et al. 2016) and matters when you start chasing why your model misses small or peripheral objects.

# Compute receptive field after a stack of layers
def receptive_field(layers: list[tuple[int, int]]) -> int:
    """layers: list of (kernel_size, stride). Returns RF in pixels."""
    rf, jump = 1, 1
    for k, s in layers:
        rf += (k - 1) * jump
        jump *= s
    return rf

# VGG-style first block: 3x3 conv, 3x3 conv, 2x2 pool stride 2
print(receptive_field([(3,1), (3,1), (2,2)]))   # 6
# After two such blocks
print(receptive_field([(3,1), (3,1), (2,2)] * 2))   # 16

LeNet-5 (LeCun et al. 1998) is the architecture that introduced approximately every idea in this chapter, two decades before they became dominant. Conv layers, average pooling, fully-connected head, gradient descent. It worked on the MNIST digits at 99% accuracy in 1998 and runs on a CPU in a few minutes today. Karpathy's "Recurse Center talk on LeCun 1989" (24-founder-blogs/karpathy-lecun1989) is worth reading after this section. He literally re-ran the 1989 paper's experiment on modern hardware and reproduced the results.

The architecture, in modern PyTorch:

class LeNet(nn.Module):
    def __init__(self, n_classes: int = 10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 6, kernel_size=5, padding=2), nn.Tanh(),
            nn.AvgPool2d(2, stride=2),
            nn.Conv2d(6, 16, kernel_size=5), nn.Tanh(),
            nn.AvgPool2d(2, stride=2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(16 * 5 * 5, 120), nn.Tanh(),
            nn.Linear(120, 84), nn.Tanh(),
            nn.Linear(84, n_classes),
        )
    def forward(self, x):
        return self.classifier(self.features(x))

Read it slowly. Two conv-pool blocks, three fully-connected layers, tanh nonlinearities. About 60k parameters. The reason this didn't take over the world in 1998 was hardware: training took weeks, datasets were tiny, and the rest of computer vision was busy on SVMs. The chapter on LeCun's original 1989 paper is worth reading once for the historical clarity.

AlexNet (Krizhevsky et al. 2012) won ImageNet by a margin so large that nobody outside the deep-learning community had calibrated for it. Top-5 error dropped from 26% to 16% in a single year. Everyone retrained their priors that week.

The architectural changes from LeNet were not philosophical. They were practical:

• ReLU instead of tanh, which removed the vanishing-gradient ceiling on depth and trained ~6x faster
• Dropout in the fully-connected layers to combat overfitting
• Two GPUs in parallel because no single GPU at the time had enough memory for the model
• Aggressive data augmentation (random crops, horizontal flips, PCA-based colour jitter)
• Local response normalisation which we have since stopped using (batch norm replaced it)

The architecture is a deeper LeNet: five conv layers, three FC layers, max pool, ReLU everywhere.

class AlexNet(nn.Module):
    def __init__(self, n_classes: int = 1000):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 96, kernel_size=11, stride=4, padding=2), nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(96, 256, kernel_size=5, padding=2), nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.Conv2d(256, 384, kernel_size=3, padding=1), nn.ReLU(inplace=True),
            nn.Conv2d(384, 384, kernel_size=3, padding=1), nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, padding=1), nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        self.classifier = nn.Sequential(
            nn.Dropout(),
            nn.Linear(256 * 6 * 6, 4096), nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096), nn.ReLU(inplace=True),
            nn.Linear(4096, n_classes),
        )
    def forward(self, x):
        return self.classifier(self.features(x).flatten(1))

About 60M parameters, most of them in those two 4096-dim FC layers. The first conv kernel is 11x11 with stride 4, which is a hand-crafted way of subsampling the 224x224 input quickly. VGG would later show that you can replace it with a stack of 3x3 convs and do better.

VGG (Simonyan and Zisserman 2014) is the architecture that proved the deep-learning intuition that "deeper is better, with simpler primitives". The whole network is 3x3 conv layers and 2x2 max pools. No 5x5, no 7x7, no fancy first-layer trick. Two-to-three 3x3 convs, then pool, then double the channels, repeat. VGG-16 has 16 weight layers; VGG-19 has 19.

The principled argument: two stacked 3x3 convs have the same receptive field as one 5x5 conv, but with $2 \cdot 9 = 18$ params per input/output channel pair versus $25$. Three stacked 3x3 convs equal one 7x7, with 27 versus 49 params. You get fewer parameters and more non-linearities for the same receptive field. The trade is more compute time (more layers, more activations to store), but better accuracy.

def vgg_block(n_convs: int, in_c: int, out_c: int) -> nn.Sequential:
    layers = []
    for _ in range(n_convs):
        layers.append(nn.Conv2d(in_c, out_c, kernel_size=3, padding=1))
        layers.append(nn.ReLU(inplace=True))
        in_c = out_c
    layers.append(nn.MaxPool2d(kernel_size=2, stride=2))
    return nn.Sequential(*layers)

class VGG16(nn.Module):
    def __init__(self, n_classes: int = 1000):
        super().__init__()
        cfg = [(2, 64), (2, 128), (3, 256), (3, 512), (3, 512)]
        in_c = 3
        blocks = []
        for n, out_c in cfg:
            blocks.append(vgg_block(n, in_c, out_c))
            in_c = out_c
        self.features = nn.Sequential(*blocks)
        self.classifier = nn.Sequential(
            nn.Flatten(), nn.Linear(512 * 7 * 7, 4096), nn.ReLU(inplace=True), nn.Dropout(),
            nn.Linear(4096, 4096), nn.ReLU(inplace=True), nn.Dropout(),
            nn.Linear(4096, n_classes),
        )
    def forward(self, x):
        return self.classifier(self.features(x))

VGG is ≈138M parameters. Most of them, again, in the FC head. The conv stack is only ≈15M. Modern architectures replace the giant FC head with global average pooling to a single conv classifier, dropping the parameter count by 10x with no accuracy hit. We will see this in ResNet.

GoogLeNet (Szegedy et al. 2014) introduced two ideas that have not gone away. First: instead of choosing one filter size per layer, run several in parallel and concatenate the results. Second: use 1x1 convolutions to mix channels and cheaply reduce dimensionality.

A 1x1 convolution is exactly what it sounds like: a filter of size 1x1xC_in that produces one output channel by linearly combining the input channels at a single spatial location. It has zero receptive field beyond a single pixel. What it does is channel mixing. It is also the only way to cheaply reduce the channel count without losing spatial information.

The Inception module computes four parallel paths from the same input: a 1x1 conv, a 3x3 conv (preceded by a 1x1 bottleneck), a 5x5 conv (preceded by a 1x1 bottleneck), and a 3x3 max pool followed by a 1x1 conv. All four outputs are concatenated along the channel dimension. The network chooses, layer by layer, how much of each filter size it wants.

class InceptionBlock(nn.Module):
    def __init__(self, in_c, c1, c2_reduce, c2, c3_reduce, c3, c4):
        super().__init__()
        # 1x1 path
        self.p1 = nn.Conv2d(in_c, c1, 1)
        # 1x1 -> 3x3 path
        self.p2 = nn.Sequential(
            nn.Conv2d(in_c, c2_reduce, 1), nn.ReLU(inplace=True),
            nn.Conv2d(c2_reduce, c2, 3, padding=1),
        )
        # 1x1 -> 5x5 path
        self.p3 = nn.Sequential(
            nn.Conv2d(in_c, c3_reduce, 1), nn.ReLU(inplace=True),
            nn.Conv2d(c3_reduce, c3, 5, padding=2),
        )
        # 3x3 pool -> 1x1 path
        self.p4 = nn.Sequential(
            nn.MaxPool2d(3, stride=1, padding=1),
            nn.Conv2d(in_c, c4, 1),
        )
    def forward(self, x):
        return torch.cat([F.relu(self.p1(x)), F.relu(self.p2(x)),
                          F.relu(self.p3(x)), F.relu(self.p4(x))], dim=1)

GoogLeNet was 22 layers deep but had fewer parameters than AlexNet (~5M versus 60M), because it ditched the giant FC head and used global average pooling. It also introduced auxiliary classifiers attached to intermediate layers, a trick for fighting vanishing gradients in deep networks before batch norm existed. Once batch norm landed, auxiliary classifiers became unnecessary and quietly disappeared.

Batch normalisation (Ioffe and Szegedy 2015) is what made very deep networks trainable. In a CNN, batch norm operates on the (N, C, H, W) tensor by computing per-channel statistics across the batch and spatial dimensions. For each channel $c$:

Two learned parameters per channel ($\gamma_c$, $\beta_c$). At inference time, you use running estimates of $\mu_c$ and $\sigma_c$ accumulated during training rather than batch statistics. This is the single most common bug in transferring a CNN to deployment: people forget to put the model in eval mode and batch norm uses the mini-batch's statistics, which during inference may be a batch of size 1 with degenerate variance.

# In a CNN block, batch norm goes between conv and ReLU
nn.Sequential(
    nn.Conv2d(64, 128, 3, padding=1, bias=False),   # bias=False because BN adds one
    nn.BatchNorm2d(128),
    nn.ReLU(inplace=True),
)

In 2015, He et al. tried to train a 34-layer plain CNN and got a worse result than an 18-layer version. This was not overfitting. The 34-layer model had higher training error too. Something about depth itself was breaking the optimisation.

The fix was the residual block. Instead of learning $H(x)$ directly, learn $F(x)$ such that $H(x) = F(x) + x$. The shortcut connection adds the input straight to the output, and the conv layers only have to model the residual — the difference from identity. This has a deep consequence: if the optimal function for a block is close to identity, gradient descent can find it by pushing $F$ toward zero, which is easy. The 34-layer plain network was failing because gradient descent could not find identity mappings in the deeper layers. ResNet handed identity to it for free.

The standard ResNet block (for ResNet-18 and ResNet-34):

import torch.nn.functional as F

class BasicBlock(nn.Module):
    expansion = 1
    def __init__(self, in_c: int, out_c: int, stride: int = 1):
        super().__init__()
        self.conv1 = nn.Conv2d(in_c, out_c, 3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_c)
        self.conv2 = nn.Conv2d(out_c, out_c, 3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_c)
        # If shapes differ, project the shortcut
        self.shortcut = nn.Identity()
        if stride != 1 or in_c != out_c:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_c, out_c, 1, stride=stride, bias=False),
                nn.BatchNorm2d(out_c),
            )
    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        return F.relu(out + self.shortcut(x))

ResNet-50 and beyond use the bottleneck block: 1x1 conv to reduce channels, 3x3 conv at the lower channel count, 1x1 conv to expand back. Same idea, four times fewer FLOPs (floating-point operations, the count of multiply-and-add steps the forward pass actually executes; the notebook's Part 4 budgets them precisely) at high channel counts. The deepest practical ResNet is around 152 layers; beyond that, you hit diminishing returns and modern alternatives (DenseNet, EfficientNet, ConvNeXt) start to look better.

The residual stream view of ResNets is interesting in its own right. Each block reads from the stream and writes a small correction back. The stream at depth $L$ is the sum of contributions from every block from 0 to $L$. This is exactly the framing Anthropic later applied to transformer residual streams in 22-anthropic-recent/2021-framework-index §residual-stream. The two architectures rhyme more than people noticed at the time.

SE blocks (Hu et al. 2017) add a tiny attention mechanism to a CNN block. After the conv layers, you take the feature map (N, C, H, W), globally average-pool to (N, C), pass through a two-layer MLP to (N, C) again, sigmoid the result, and multiply it back into the feature map channel-wise. You are learning a per-channel gate that says "this channel is relevant for this input, that channel is not".

class SqueezeExcitation(nn.Module):
    def __init__(self, channels: int, reduction: int = 16):
        super().__init__()
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channels, channels // reduction), nn.ReLU(inplace=True),
            nn.Linear(channels // reduction, channels), nn.Sigmoid(),
        )
    def forward(self, x):
        B, C, _, _ = x.shape
        w = self.fc(self.pool(x).view(B, C)).view(B, C, 1, 1)
        return x * w

This is input-dependent channel gating, not self-attention. Like attention, it produces data-dependent weights and uses them to reweight features. Unlike transformer self-attention, it does not form queries and keys or compute pairwise interactions between positions/channels: global average pooling compresses each channel, and an MLP emits one gate per channel. The analogy is useful at the level of dynamic reweighting, but the mechanisms are not equivalent. SE blocks gave ResNet a meaningful accuracy bump (winning ImageNet 2017) with minimal compute overhead. They are still used in production ConvNets.

This is the part where you stop treating the CNN as a black box. Olah, Mordvintsev and Schubert's work on 01-explorables/distill-feature-visualization and the follow-up Circuits thread (distill-circuits-early-vision, distill-circuits-curve-detectors, distill-circuits-zoom-in) showed that you can label individual units in a trained CNN by what they detect: Gabor-like edges at layer 1, textures at layer 2, curves and parallel lines at layer 3, dog snouts and car wheels at layer 5, semantic objects at layer 8.

Two techniques drive this. Feature visualisation generates an input image that maximally activates a specific unit, by gradient ascent on the input pixels. Dataset examples find the images in the training set that most activate a unit. Combine them and you get a hypothesis ("this unit fires for circles") plus the evidence ("here are the 9 training images that triggered it most"). When the two agree, you have an interpretation you can act on.

The Distill Circuits work goes one step further: it identifies families of units that work together (curve detectors at different orientations, parallel line detectors at different angles) and shows that they are connected by interpretable weights. Layer 3 curve detectors are constructed from layer 2 line detectors by literally adding curved templates. This is mechanistic interpretability (mech-interp: reverse-engineering the algorithm a network has actually learned by reading its weights and activations, the subject of another chapter) before mech-interp had a name, on a vision model. The lessons transferred to transformers later.

# Sketch of feature visualisation via gradient ascent
def visualise_unit(model, layer_idx: int, channel: int, steps: int = 500):
    x = torch.randn(1, 3, 224, 224, requires_grad=True)
    optim = torch.optim.Adam([x], lr=0.05)
    activations = {}
    def hook(mod, inp, out):
        activations['out'] = out
    h = list(model.modules())[layer_idx].register_forward_hook(hook)
    for _ in range(steps):
        optim.zero_grad()
        model(x)
        # Negate: we want to MAXIMISE the activation
        loss = -activations['out'][0, channel].mean()
        loss.backward()
        optim.step()
    h.remove()
    return x.detach()

In practice you regularise this heavily (decorrelate the input, jitter, total-variation penalty) so the generated image is interpretable rather than a high-frequency adversarial pattern. See 01-explorables/distill-feature-visualization §regularization for the production recipe.

Classification asks "what is in this image". Detection asks "what is in this image and where is it". The bounding-box output makes the problem harder, but the architectural primitives are the same convs and pools.

Two families. Region-proposal methods (R-CNN, Fast R-CNN, Faster R-CNN) generate candidate regions first and then classify each region. They are accurate and slow. Single-shot methods (YOLO, SSD, RetinaNet) treat detection as a dense regression problem: the CNN predicts, at every spatial location on a downsampled feature map, "is there an object here, and what is its class and box". They are fast and competitive on accuracy once you add focal loss and FPN.

# YOLO-style head: at each cell, predict (x, y, w, h, objectness, class_logits)
class YOLOHead(nn.Module):
    def __init__(self, in_c: int, n_anchors: int, n_classes: int):
        super().__init__()
        # 5 = x, y, w, h, objectness
        self.out = nn.Conv2d(in_c, n_anchors * (5 + n_classes), kernel_size=1)
    def forward(self, x):
        return self.out(x)

The actual training loss is intricate (anchor matching, IOU-based box regression, hard-negative mining, focal loss for class imbalance), and we are not going to derive it here. The point for this chapter is structural: detection re-uses every CNN primitive you have built, and the architectural innovation is in what to predict at every spatial cell, not in new operations.

Segmentation classifies every pixel. The output has the same spatial resolution as the input. The architectural challenge is that you cannot downsample everything away if you need pixel-level outputs back. U-Net (Ronneberger et al. 2015) solved this with a symmetric encoder-decoder: pool down to a low-resolution bottleneck, then upsample back, and at each scale add a skip connection from the encoder side to the decoder side. The skip connections carry high-frequency spatial information that the bottleneck would have destroyed.

class UNet(nn.Module):
    def __init__(self, n_classes: int = 2):
        super().__init__()
        def conv_block(in_c, out_c):
            return nn.Sequential(
                nn.Conv2d(in_c, out_c, 3, padding=1), nn.ReLU(inplace=True),
                nn.Conv2d(out_c, out_c, 3, padding=1), nn.ReLU(inplace=True),
            )
        self.down1 = conv_block(3, 64)
        self.down2 = conv_block(64, 128)
        self.down3 = conv_block(128, 256)
        self.bottleneck = conv_block(256, 512)
        self.up3 = nn.ConvTranspose2d(512, 256, 2, stride=2)
        self.dec3 = conv_block(512, 256)
        self.up2 = nn.ConvTranspose2d(256, 128, 2, stride=2)
        self.dec2 = conv_block(256, 128)
        self.up1 = nn.ConvTranspose2d(128, 64, 2, stride=2)
        self.dec1 = conv_block(128, 64)
        self.out = nn.Conv2d(64, n_classes, 1)
    def forward(self, x):
        d1 = self.down1(x)
        d2 = self.down2(F.max_pool2d(d1, 2))
        d3 = self.down3(F.max_pool2d(d2, 2))
        b = self.bottleneck(F.max_pool2d(d3, 2))
        u3 = self.dec3(torch.cat([self.up3(b), d3], dim=1))
        u2 = self.dec2(torch.cat([self.up2(u3), d2], dim=1))
        u1 = self.dec1(torch.cat([self.up1(u2), d1], dim=1))
        return self.out(u1)

That ConvTranspose2d is the upsampling operation: a "fractionally strided" or "transposed" convolution that produces a larger output from a smaller input. It is the only new primitive in U-Net. Modern segmentation often replaces it with bilinear upsampling followed by a regular conv, which avoids checkerboard artifacts at the cost of one extra layer.

Vision transformers (ViT, Dosovitskiy et al. 2020) replaced the convolutional inductive bias with the transformer's self-attention, where every patch can directly read from every other patch (the SE block in §11 did a 1D version of this over channels; ViT does it over the full set of patch tokens, which is what removes the hard-coded locality). The image is cut into fixed-size patches (typically 16x16), each patch is flattened and linearly projected into a token embedding, and the resulting sequence of patch tokens is fed to a standard transformer encoder. A learned [CLS] token's output embedding becomes the image's representation for classification.

The trade-off is well documented now. With enough data (JFT-300M scale), ViT outperforms ResNets on ImageNet. With ImageNet-1k alone (1.3M images), ViT underperforms ResNet because the convolutional inductive bias (locality, translation equivariance) does real work when data is scarce. The middle ground is hybrid models (ConvNeXt, Swin Transformer, MaxViT) that bring some convolutional locality back to the transformer.

class PatchEmbed(nn.Module):
    """Image -> sequence of patch tokens via a single strided conv."""
    def __init__(self, img_size: int = 224, patch: int = 16, d_model: int = 768):
        super().__init__()
        self.proj = nn.Conv2d(3, d_model, kernel_size=patch, stride=patch)
        self.n_patches = (img_size // patch) ** 2
    def forward(self, x):
        # x: (B, 3, H, W) -> (B, d_model, H/p, W/p) -> (B, n_patches, d_model)
        return self.proj(x).flatten(2).transpose(1, 2)

That single strided conv nn.Conv2d(3, d_model, kernel_size=16, stride=16) is the entire convolutional content of a ViT. Everything after it is the another chapter transformer architecture applied to image tokens.

This is the chapter's exit door. another chapter picks up here with the transformer details. The reason ViT works is that the transformer architecture is general enough that "image patches" and "tokens" are interchangeable, if you have enough data to learn the spatial relationships the convolution would have hard-coded. We will come back to that "if".