Convolutional networks

What it is and when you reach for it

A convolutional layer is a linear layer with two constraints imposed before training starts: locality, each output looks only at a small neighborhood of the input, and weight sharing, the same kernel is reused at every position. Both constraints remove parameters rather than add machinery, and that is the point. A fully connected layer from a 28×28 image to a same-sized bank of 6 feature maps would need about 3.7 million weights; the convolutional version in LeNet does the same job with 156, because a 5×5 kernel per input-output channel pair is all it stores. The removed parameters encode an assumption, called an inductive bias: an edge detector useful in the top-left corner is useful everywhere, so learn it once. When that assumption matches the data, grids where nearby values are related and patterns are translation-invariant (images, spectrograms, board states, 1-D sensor streams), the network learns faster from less data than any unconstrained model could. When it does not match, say tabular columns with no spatial order, convolution buys nothing. Relative to its neighbors: the MLP is convolution with both constraints deleted, and the transformer deletes locality but keeps position-independent weights, then lets attention decide dynamically what is near what. A longer conceptual treatment lives in my notes at /ai/notes/convolutional-networks; this page is the implementation-first version.

The math

The sliding dot product

For a single channel, the output at position (i, j) is the dot product of a k×k kernel w with the input patch anchored there: out[i, j] = Σ_aΣ_b w[a, b] · x[i·s + a, j·s + b], with stride s and the sums running over the kernel. A tiny example makes it concrete. Take a 4×4 input and the 2×2 kernel [[1, 0], [0, −1]], which computes x[i, j] − x[i+1, j+1], a crude diagonal-edge detector:

input           kernel        output (3×3, stride 1, no padding)
1 2 0 1         1  0          1−1  2−3  0−1        0 −1 −1
0 1 3 1         0 −1          0−1  1−0  3−0   =   −1  1  3
2 1 0 0                       2−0  1−1  0−2        2  0 −2
1 0 1 2

With multiple channels the dot product also runs over the input channel axis: a kernel of shape (C_in, k, k) produces one output channel, and a layer stacks C_out such kernels into a weight tensor of shape (C_out, C_in, k, k). Each output channel is one learned pattern detector evaluated everywhere. Strictly speaking this operation is cross-correlation, not the signal-processing convolution, which would flip the kernel first; deep learning frameworks skip the flip because a learned kernel can just as easily learn the flipped weights. The name stuck anyway.

Output-size arithmetic, worked numerically

For input size n, kernel k, padding p, and stride s, the output size along each spatial axis is

out = ⌊(n + 2p − k) / s⌋ + 1

The formula counts anchor positions: the padded input has n + 2p pixels, the last anchor must leave room for the kernel (subtract k, then the +1 restores the first position), and stride s keeps every s-th anchor. Worked through the numbers this page actually uses: a 28×28 digit through a 5×5 kernel with p = 2, s = 1 gives (28 + 4 − 5)/1 + 1 = 28, size preserved, which is what "same" padding means and why p = (k−1)/2 is the magic value for odd k. A 2×2 max pool with stride 2 gives (28 − 2)/2 + 1 = 14. The second conv, 5×5 with no padding, gives (14 − 5)/1 + 1 = 10, and the second pool (10 − 2)/2 + 1 = 5. One stride-2 example from a real stem: ResNet's first layer is a 7×7 kernel, p = 3, s = 2 on a 224×224 image, so ⌊(224 + 6 − 7)/2⌋ + 1 = ⌊223/2⌋ + 1 = 112, and the floor is doing real work there. Getting this arithmetic wrong is the single most common way a hand-built network fails to run, so it is worth being able to do it in your head.

Receptive fields

A unit deep in the network depends on a patch of the original input called its receptive field, and the patch grows layer by layer according to r_l = r_l−1 + (k_l − 1) · j_l−1, where j is the jump, the product of all strides so far. Tracing LeNet: after conv 5×5 (r = 5, j = 1), pool 2×2 stride 2 (r = 6, j = 2), conv 5×5 (r = 6 + 4·2 = 14, j = 2), pool (r = 16, j = 4). So each unit in the final 5×5 feature map sees a 16×16 patch of the digit, and the fully connected head that reads all 25 positions sees everything. This is also the arithmetic behind a famous design shift: two stacked 3×3 convolutions have the same 5×5 receptive field as one 5×5 layer but cost 2·9·C² = 18C² weights instead of 25C² and insert an extra nonlinearity between them. That observation is essentially the VGG paper, and it is why kernels larger than 3×3 became rare for a decade. One honest caveat: the effective receptive field, where the gradient actually concentrates, is measurably smaller than this theoretical box and roughly Gaussian around the center, so treat the formula as an upper bound.

Pooling and batch norm

Pooling summarizes each small window by one number, usually the max, halving resolution while keeping the strongest response. It buys a little local translation invariance (a feature that shifts by a pixel usually still wins its window) and, more importantly, it lets the layers above cover more of the image with the same kernel size, which is exactly the jump term in the receptive-field recursion. Modern networks often replace pooling with stride-2 convolutions, which do the same downsampling with learned weights; the arithmetic is identical.

Batch normalization, at the concept level, standardizes each channel using the statistics of the current batch: for channel c, compute the mean and variance over all samples and all spatial positions, normalize to zero mean and unit variance, then apply a learned per-channel scale γ and shift β so the network keeps its expressive range. Normalizing per channel rather than per unit is forced by weight sharing: a channel is one feature evaluated everywhere, so it gets one set of statistics. The practical payoff is that activations stay in a healthy range regardless of depth, which permits higher learning rates and makes deep stacks trainable. The practical trap is that it behaves differently at training time (batch statistics) and at inference time (running averages accumulated during training), which is exactly what the train/eval mode switch in every framework controls.

Implementation, twice

conv2d from scratch: convolution is a matmul wearing a coat

The classic way to implement convolution, and still one of the ways production backends do it, is im2col: copy every k×k×C receptive field out into a column of a matrix, so the whole convolution collapses into a single matrix multiply between the flattened kernels and that column matrix. It spends memory (each pixel is duplicated up to k² times) to buy the one operation GPUs are best at. The PyTorch tab uses F.unfold, which is im2col by its framework name; the JAX tab shows the production path, lax.conv_general_dilated, the primitive underneath every JAX conv layer, plus an educational im2col built from pad, slice, and einsum. Both tabs end by asserting agreement with the framework's own convolution on a fixed seed, which is the check that makes the exercise trustworthy.

import torch
import torch.nn.functional as F

def conv2d_im2col(x, w, b=None, stride=1, padding=0):
    """x: (N, C_in, H, W), w: (C_out, C_in, kH, kW).

    F.unfold copies every receptive field into a column, ordered
    channel-major then kernel-position, which is exactly the order
    w.view(C_out, -1) flattens to, so one batched matmul finishes it.
    """
    N, C, H, W = x.shape
    C_out, _, kH, kW = w.shape
    H_out = (H + 2 * padding - kH) // stride + 1
    W_out = (W + 2 * padding - kW) // stride + 1
    cols = F.unfold(x, (kH, kW), stride=stride, padding=padding)
    #    (N, C*kH*kW, L) with L = H_out*W_out
    out = w.view(C_out, -1) @ cols          # (C_out,K) @ (N,K,L) -> (N,C_out,L)
    if b is not None:
        out = out + b.view(1, -1, 1)
    return out.view(N, C_out, H_out, W_out)

# verify against the library on a fixed seed
torch.manual_seed(0)
x = torch.randn(2, 3, 8, 8)
w = torch.randn(4, 3, 3, 3)
b = torch.randn(4)
mine = conv2d_im2col(x, w, b, stride=2, padding=1)
ref  = F.conv2d(x, w, b, stride=2, padding=1)
assert mine.shape == ref.shape == (2, 4, 4, 4)   # (8+2-3)//2+1 = 4
assert torch.allclose(mine, ref, atol=1e-5)

import jax
import jax.numpy as jnp
from jax import lax

def conv2d(x, w, stride=1, padding=0):
    """The production path: this primitive underlies every JAX conv
    layer, and XLA lowers it to the backend's native convolution."""
    return lax.conv_general_dilated(
        x, w,
        window_strides=(stride, stride),
        padding=[(padding, padding), (padding, padding)],
        dimension_numbers=("NCHW", "OIHW", "NCHW"))

def conv2d_im2col(x, w, stride=1, padding=0):
    """Educational im2col: same answer from pad + slice + einsum.

    The loops run over kernel offsets (k*k iterations, tiny and
    unrolled under jit), never over image positions, so every slice
    below is a full strided view of the image: still vectorized.
    """
    N, C, H, W = x.shape
    C_out, _, kH, kW = w.shape
    H_out = (H + 2 * padding - kH) // stride + 1
    W_out = (W + 2 * padding - kW) // stride + 1
    xp = jnp.pad(x, ((0, 0), (0, 0), (padding,) * 2, (padding,) * 2))
    patches = jnp.stack(
        [xp[:, :, i:i + stride * H_out:stride, j:j + stride * W_out:stride]
         for i in range(kH) for j in range(kW)],
        axis=2)                                   # (N, C, kH*kW, H_out, W_out)
    cols = patches.reshape(N, C * kH * kW, H_out * W_out)
    out = jnp.einsum("ok,nkl->nol", w.reshape(C_out, -1), cols)
    return out.reshape(N, C_out, H_out, W_out)

# verify one against the other on a fixed key
k1, k2 = jax.random.split(jax.random.PRNGKey(0))
x = jax.random.normal(k1, (2, 3, 8, 8))
w = jax.random.normal(k2, (4, 3, 3, 3))
a = conv2d(x, w, stride=2, padding=1)
b = conv2d_im2col(x, w, stride=2, padding=1)
assert a.shape == (2, 4, 4, 4)                   # (8+2-3)//2+1 = 4
assert jnp.allclose(a, b, atol=1e-4)

The im2col versions are forward-only on purpose: the backward pass through convolution is itself a convolution (with the kernel spatially flipped, against the output gradient), and autograd derives it for free through unfold and einsum, which you can confirm by calling grad through either function.

A LeNet-style classifier, idiomatically

The second block is the network that started the field: two conv-pool stages that shrink 28×28 down to a 16-channel 5×5 map, then an MLP head that reads it, about 60k parameters in total. The shape pipeline, with the arithmetic from above annotated:

input 1×28×28
  │ conv 5×5, pad 2, 6 filters      →  6×28×28    (28+4−5)/1+1 = 28
  │ ReLU, maxpool 2×2 s2            →  6×14×14    (28−2)/2+1  = 14
  │ conv 5×5, 16 filters            → 16×10×10    (14−5)/1+1  = 10
  │ ReLU, maxpool 2×2 s2            → 16×5×5      (10−2)/2+1  = 5
  │ flatten                         → 400
  │ dense 120 → dense 84 → dense 10 → logits

For the JAX tab I chose flax (Linen) over a raw params pytree, and the reason is where the teaching value lives. The neural network page already builds and threads a pytree by hand, and for an MLP that is instructive; for a CNN the pytree stops teaching and starts being bookkeeping, since every conv layer carries a differently shaped kernel and the natural next step, batch norm, adds mutable running statistics that plain jax.grad has no slot for. Flax's Module handles initialization from a single RNG key, nests the parameters with readable names, and has a first-class place for that mutable state, while everything still compiles down to the same lax conv primitive shown above.

import torch
from torch import nn

class LeNet(nn.Module):
    """LeNet-5 shaped, for 1×28×28 inputs. NCHW, PyTorch's native layout."""
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 6, 5, padding=2), nn.ReLU(),   # 28 -> 28
            nn.MaxPool2d(2),                            # 28 -> 14
            nn.Conv2d(6, 16, 5), nn.ReLU(),             # 14 -> 10
            nn.MaxPool2d(2),                            # 10 -> 5
        )
        self.head = nn.Sequential(
            nn.Flatten(),
            nn.Linear(16 * 5 * 5, 120), nn.ReLU(),
            nn.Linear(120, 84), nn.ReLU(),
            nn.Linear(84, num_classes),                 # raw logits
        )

    def forward(self, x):
        return self.head(self.features(x))

model = LeNet()
opt = torch.optim.Adam(model.parameters(), lr=1e-3)

def train_step(x, y):
    opt.zero_grad()
    loss = nn.functional.cross_entropy(model(x), y)
    loss.backward()
    opt.step()
    return loss.item()

# smoke test on a fake batch
x, y = torch.randn(64, 1, 28, 28), torch.randint(0, 10, (64,))
print(train_step(x, y))   # ~2.3 = ln(10): uniform over 10 classes

import jax
import jax.numpy as jnp
import flax.linen as nn
import optax

class LeNet(nn.Module):
    """LeNet-5 shaped, for 28×28×1 inputs. NHWC, flax's native layout."""
    num_classes: int = 10

    @nn.compact
    def __call__(self, x):
        x = nn.Conv(6, (5, 5), padding="SAME")(x)          # 28 -> 28
        x = nn.relu(x)
        x = nn.max_pool(x, (2, 2), strides=(2, 2))         # 28 -> 14
        x = nn.Conv(16, (5, 5), padding="VALID")(x)        # 14 -> 10
        x = nn.relu(x)
        x = nn.max_pool(x, (2, 2), strides=(2, 2))         # 10 -> 5
        x = x.reshape(x.shape[0], -1)                      # 400
        x = nn.relu(nn.Dense(120)(x))
        x = nn.relu(nn.Dense(84)(x))
        return nn.Dense(self.num_classes)(x)               # raw logits

model = LeNet()
params = model.init(jax.random.PRNGKey(0), jnp.zeros((1, 28, 28, 1)))
tx = optax.adam(1e-3)
opt_state = tx.init(params)

def loss_fn(params, x, y):
    logits = model.apply(params, x)
    return optax.softmax_cross_entropy_with_integer_labels(logits, y).mean()

@jax.jit
def train_step(params, opt_state, x, y):
    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
    updates, opt_state = tx.update(grads, opt_state)
    return optax.apply_updates(params, updates), opt_state, loss

# smoke test on a fake batch
k1, k2 = jax.random.split(jax.random.PRNGKey(1))
x = jax.random.normal(k1, (64, 28, 28, 1))
y = jax.random.randint(k2, (64,), 0, 10)
params, opt_state, loss = train_step(params, opt_state, x, y)
print(loss)               # near ln(10) ≈ 2.3: still guessing over 10 classes

Using it on a real shape of problem

Point either training step at MNIST (60,000 grayscale 28×28 digits, batches of 64, a couple of epochs) and the behavior is textbook. The first printed loss sits near ln 10 ≈ 2.30, the entropy of guessing uniformly over ten classes; within a few hundred steps it is under 0.5, and after one epoch a LeNet at this scale reaches roughly 98% test accuracy, with 99% in reach after a few more, exact numbers varying with seed and machine. Two diagnostics are worth building the habit of checking. First, that initial loss: if it is not close to ln(num_classes), the head or the loss wiring is wrong before training has anything to do with it. Second, overfit a single batch: any correct CNN of this size should drive 64 samples to essentially zero loss within a few hundred steps, and failure to do so means a bug, not a capacity problem. The same architecture with the first conv widened to 3 input channels trains on CIFAR-10 unchanged, though it plateaus far lower there, which is the historically accurate signal that depth, normalization, and residual connections had to be invented.

Applications

The convolutional lineage is the spine of computer vision. LeNet read a meaningful share of US bank checks in the late 1990s; AlexNet's 2012 ImageNet win is the event that made deep learning mainstream; VGG established the all-3×3 design; ResNet's residual connections made 100-plus-layer networks trainable and remain the default baseline backbone. On top of classification backbones sit the detection family, Faster R-CNN and the YOLO line, still dominant wherever real-time boxes matter, and segmentation, where U-Net's encoder-decoder with skip connections is so entrenched in medical imaging that entire frameworks exist just to serve variants of it. Convolutions also read audio: a spectrogram is an image whose axes are time and frequency, and conv frontends power keyword spotting on phones, speaker ID, and the input stem of Whisper, which downsamples the mel spectrogram with two stride-2 convolutions before any transformer layer sees it.

In the transformer era convs did not disappear; they moved into the hybrids. Vision transformers commonly use convolutional stems or patch-embedding convs; ConvNeXt showed that a modernized pure-conv network matches same-scale ViTs; MobileNet and EfficientNet variants still own on-device inference because depthwise convolutions are cheap in exactly the way mobile hardware likes; and the U-Net inside Stable Diffusion's original architecture interleaves convolutional residual blocks with attention. The bias survives because it is true: at the lowest layers of vision, features really are local and translation-invariant, and paying attention's quadratic price to rediscover that is wasteful.

Against the real libraries

The im2col conv above is honest about the algorithm and silent about everything a production stack adds. Backends like cuDNN choose among several algorithms per layer shape (im2col-style GEMM, FFT for large kernels, Winograd for 3×3, direct kernels), use tensor cores and fused epilogues, and never materialize the k²-duplicated column matrix the naive version pays for. torchvision supplies the canonical model implementations (ResNet, MobileNet, Faster R-CNN, and the rest) with pretrained ImageNet weights, the dataset and transform plumbing, and reference training recipes whose accuracy numbers are documented per model. timm is the model zoo of record beyond that: over a thousand architectures behind one create_model interface, with pretrained weights and the training tricks (augmentation recipes, EMA, schedulers) that reproduce published results, which is why nearly every vision paper's baseline table is built on it. On the JAX side, the flax repository's examples directory contains the reference MNIST and ImageNet training pipelines that most JAX vision code descends from. My notes on the PyTorch codebase itself are at /oss/pytorch.

When is from-scratch enough? For understanding, always; for production, essentially never, and the line is sharper here than for the MLP because convolution performance is dominated by algorithm selection and kernel fusion you cannot replicate in a few lines. The from-scratch version earns its keep as a verification oracle and as the mental model that explains behavior: why channels-last memory layout speeds up training, why a 1×1 conv is just a per-pixel linear layer, why grouped and depthwise convs are cheap. The verification recipe is in the code above: fixed seed, random x and w, compare your implementation against F.conv2d or lax.conv_general_dilated with allclose at about 1e-5 in float32, and include a strided, padded case, because stride-1 no-padding agreement catches almost nothing. If your conv matches the library on a padded stride-2 case with non-square everything, your indexing is right; that single test retires the whole class of off-by-one bugs.

Traps and misconceptions

"Convolution" is cross-correlation. Frameworks do not flip the kernel, so a hand-derived signal-processing result can disagree with conv2d until you notice. For learning it is irrelevant, the network learns flipped weights if it needs them, but it matters whenever you port classical filters or compare against scipy, whose convolve2d does flip.

Trusting "same" padding without doing the arithmetic. Same padding preserves size only at stride 1. At stride 2 a "same"-padded 28 input gives 14, and for even inputs the padding is asymmetric, so different frameworks can place the extra pixel on different sides and produce subtly different outputs from identical weights. Porting bugs between PyTorch (explicit p) and TensorFlow or flax ("SAME" strings) usually live exactly here.

Forgetting eval mode with batch norm. In train mode batch norm uses current-batch statistics; forget model.eval() and inference results depend on batch composition, including the degenerate single-sample case where batch variance is meaningless. The symptom is a model that "works in the notebook, fails in the service." The same switch controls dropout, so the habit pays twice.

Confusing theoretical with effective receptive field. The recursion in the math section gives the maximum patch a unit can see; the gradient-weighted effective field is much smaller and Gaussian-shaped, so a network whose theoretical field covers the image can still fail on genuinely global structure. When objects are large relative to the field that actually matters, the fixes are depth, dilation, or attention, not hope.

Layout mixups between NCHW and NHWC. PyTorch defaults to channels-first, flax and TensorFlow to channels-last, and a tensor reshaped across that boundary rather than transposed produces scrambled images that still have the right shape, so nothing crashes: the model just trains badly. When porting weights, permute axes explicitly and verify with a fixed-input forward-pass comparison before training anything.