SOCR ≫ DSPA ≫ DSPA3 Topics ≫

library(ggplot2)
library(dplyr)
library(tidyr)
library(patchwork)
library(plotly)
if (TORCH_OK) library(torch)
if (TV_OK)    library(torchvision)
if (LUZ_OK)   library(luz)

Chapter 14 is published in five parts

Part Sections Content
Part 1: Foundations §14.0–§14.24 Tensors, automatic differentiation, the MLP, activations, initialization, normalization, optimizers, regularization, the training loop
Part 2 (this document) §14.25–§14.44 Convolution arithmetic, receptive fields, residual connections, transfer learning, CIFAR-10, segmentation, interpretability
Part 3: Sequence Models §14.45–§14.64 RNNs, LSTM/GRU, attention derived, Transformers, tokenization, text generation, neural forecasting
Part 4: Generative and Representation Learning §14.65–§14.84 Autoencoders, VAEs and the ELBO, GANs, diffusion, U-Net synthesis, self-supervised learning
Part 5: Generalization, Uncertainty, and Practice §14.85–§14.104 Double descent, calibration, ensembles, hyperparameter optimization, pruning, robustness

Section numbers are continuous across all five parts. Part 2 assumes Part 1 throughout: tensors and autograd (§14.1–14.2), the initialization derivation (§14.9), normalization (§14.10), and the training loop (§14.14).


How this part uses graphics

Two-dimensional figures use ggplot2; the equivalent plot_ly() code follows in a chunk marked eval=FALSE, echo=TRUE. Three-dimensional figures are evaluated plot_ly(), budgeted at six per part.

Heavy training is gated. HEAVY_EVAL defaults to FALSE, so CIFAR-10 and the segmentation case study display their code and load cached results rather than training during the build. Set it TRUE for a full run.


Learning objectives

After completing Part 2 you will be able to:

  1. Explain convolution as structured sparsity plus weight sharing, and compute the parameter saving over a dense layer.
  2. Compute output sizes for any combination of kernel, stride, padding, and dilation.
  3. Derive the receptive field of a unit and explain how it grows with depth.
  4. Compare pooling with strided convolution, and state what each discards.
  5. Distinguish equivariance from invariance, and demonstrate which one a convolutional layer has.
  6. Analyze the FLOP and memory cost of a convolutional layer and explain why CNNs are compute-bound where MLPs are memory-bound.
  7. Explain the degradation problem and why residual connections solve it, with the gradient-flow argument.
  8. Apply data augmentation and state precisely what it does and does not add.
  9. Design a transfer-learning protocol, what to freeze, what to fine-tune, and what BatchNorm does during it.
  10. Evaluate a segmentation model with the Dice coefficient, including its behavior on empty masks, and connect it to the metric family of Chapter 9.
  11. Produce and criticize saliency and Grad-CAM attributions.

Estimated time: 10–14 hours including exercises.


PART I: THE CONVOLUTIONAL PRIMITIVE

14.25 Convolution as structured sparsity

A dense layer connects every input to every output. For a \(224\times224\) RGB image flattened to \(150{,}528\) inputs and a hidden layer of \(1{,}000\) units, that is \(1.5\times10^{8}\) weights in the first layer alone, more parameters than ImageNet has training examples.

A convolutional layer replaces the dense connection with two constraints:

\[ \begin{aligned} \textbf{Local connectivity: }&\quad \text{each output depends only on a }k\times k\text{ neighbourhood}\\ \textbf{Weight sharing: }&\quad \text{the \emph{same} }k\times k\text{ kernel is applied at every spatial position} \end{aligned} \]

For input \(X\) with \(C_{\text{in}}\) channels and a kernel \(W\) of size \(k\times k\), the output at channel \(j\) and position \((u,v)\) is

\[\boxed{\;Y_{j,u,v}=b_j+\sum_{c=1}^{C_{\text{in}}}\sum_{p=1}^{k}\sum_{q=1}^{k}W_{j,c,p,q}\,X_{c,\,u+p-1,\,v+q-1}\;}\]

with parameter count \(C_{\text{out}}\big(C_{\text{in}}k^2+1\big)\), independent of the image size.

Common misconception: “a CNN filters the image and then classifies it.” That describes classical computer vision, where a human designs the filters, Sobel for edges, Gabor for texture, and a separate classifier consumes their outputs. The two stages are distinct and the filters are fixed.

In a CNN the kernels are parameters, learned by the same gradient descent that fits the classifier, jointly and end to end. Nothing designates a layer as “the filtering stage.” Early layers do converge on edge- and color-opponent detectors that resemble Gabor filters, that is an empirical finding about what gradient descent discovers, not an architectural stipulation, and later layers compose them into representations no one designed.

The consequence is practical: a convolutional layer is a structural prior (locality and translation-sharing), not a preprocessing step. It restricts the hypothesis space in a way suited to images, which is why it needs fewer examples than a dense layer to reach the same accuracy.

param_dense <- function(h, w, c_in, units) (h * w * c_in) * units + units
param_conv  <- function(k, c_in, c_out) c_out * (c_in * k^2 + 1)

data.frame(
  input = c("32 x 32 x 3 (CIFAR)", "224 x 224 x 3 (ImageNet)",
            "512 x 512 x 1 (MRI slice)"),
  dense_1000_units = format(c(param_dense(32,32,3,1000),
                              param_dense(224,224,3,1000),
                              param_dense(512,512,1,1000)), big.mark = ","),
  conv_3x3_64_filters = format(c(param_conv(3,3,64), param_conv(3,3,64),
                                 param_conv(3,1,64)), big.mark = ","),
  saving_factor = round(c(param_dense(32,32,3,1000)/param_conv(3,3,64),
                          param_dense(224,224,3,1000)/param_conv(3,3,64),
                          param_dense(512,512,1,1000)/param_conv(3,1,64))))

The saving is three to four orders of magnitude, and it grows with image size, the dense count scales as \(hw\), the convolutional count does not scale with the image at all.

sizes <- c(16, 32, 64, 128, 224, 320, 512)
chans <- c(1, 3, 8, 16, 32, 64)
Zdense <- outer(chans, sizes, function(c_, s) log10(param_dense(s, s, c_, 1000)))
Zconv  <- outer(chans, sizes, function(c_, s) log10(param_conv(3, c_, 64)))

plot_ly() |>
  add_surface(x = sizes, y = chans, z = Zdense, opacity = 0.9, showscale = FALSE,
              colorscale = "Reds", name = "Dense") |>
  add_surface(x = sizes, y = chans, z = Zconv, opacity = 0.9, showscale = FALSE,
              colorscale = "Blues", name = "Conv 3x3") |>
  layout(title = "Parameter count: dense (upper, red) vs. 3x3 convolution (lower, blue)",
         scene = list(xaxis = list(title = "Image side length"),
                      yaxis = list(title = "Input channels"),
                      zaxis = list(title = "log10 parameters")))

Rotate along the image-size axis. The dense surface climbs steeply, it scales as \(h w C_{\text{in}}\), while the convolutional surface is flat in image size and rises only with channel count. That flatness is the whole reason convolution made large images tractable.

# A learned kernel and a hand-designed one, applied to the same input
torch_manual_seed(11)
img <- torch_zeros(1, 1, 32, 32)
img[1, 1, 10:22, 10:22] <- 1                      # a bright square
img <- img + torch_randn(1, 1, 32, 32) * 0.08

sobel_x <- torch_tensor(array(c(-1,0,1, -2,0,2, -1,0,1), dim = c(1,1,3,3)))
edge <- with_no_grad(nnf_conv2d(img, sobel_x, padding = 1))

conv_learn <- nn_conv2d(1, 1, kernel_size = 3, padding = 1)
rand_out <- with_no_grad(conv_learn(img))

to_df <- function(t, lab) {
  m <- as.array(t$squeeze())
  expand.grid(row = 1:nrow(m), col = 1:ncol(m)) |>
    mutate(v = as.vector(m), panel = lab)
}
bind_rows(to_df(img, "Input"), to_df(edge, "Hand-designed Sobel-x"),
          to_df(rand_out, "Randomly initialized kernel")) |>
  mutate(panel = factor(panel, levels = c("Input", "Hand-designed Sobel-x",
                                          "Randomly initialized kernel"))) |>
  ggplot(aes(col, -row, fill = v)) +
  geom_raster() + facet_wrap(~ panel) + coord_fixed() +
  scale_fill_viridis_c(guide = "none") +
  labs(title = "The kernel is a parameter, not a design choice",
       subtitle = "Sobel is fixed by hand; the third panel's kernel is random and will be learned by gradient descent",
       x = NULL, y = NULL) +
  theme_dspa(9) + theme(axis.text = element_blank(), panel.grid = element_blank())

14.26 Convolution arithmetic

Four hyperparameters determine the output size.

\[\boxed{\;H_{\text{out}}=\left\lfloor\frac{H_{\text{in}}+2p-d(k-1)-1}{s}\right\rfloor+1\;}\]

with kernel \(k\), padding \(p\), stride \(s\), and dilation \(d\).

\[ \begin{aligned} \textbf{Kernel }k:&\quad \text{the neighbourhood size}\\ \textbf{Stride }s:&\quad \text{step between applications; }s>1\text{ downsamples by }\approx s\\ \textbf{Padding }p:&\quad \text{border added; }p=\lfloor k/2\rfloor\text{ with }s=d=1\text{ preserves size}\\ \textbf{Dilation }d:&\quad \text{gaps inside the kernel; enlarges the receptive field at fixed cost} \end{aligned} \]

Common misconception: “‘same’ padding preserves the image.” It preserves the output shape, and nothing about the border content. Zero-padding introduces artificial zeros that the kernel treats as data, so border outputs are computed from a mixture of real pixels and invented ones. Deep stacks accumulate this into a visible frame artifact, and in segmentation it systematically degrades boundary predictions, which is exactly where segmentation accuracy is usually judged.

Alternatives worth knowing: padding_mode = "reflect" or "replicate" extend the image plausibly rather than with zeros, and valid convolution (\(p=0\)) refuses to invent anything at the cost of shrinking the output, which is why the original U-Net used valid convolutions and cropped its skip connections.

out_size <- function(H, k, p = 0, s = 1, d = 1) floor((H + 2*p - d*(k-1) - 1)/s) + 1

cfg <- expand.grid(k = c(3, 5, 7), s = c(1, 2), p = c(0, NA), d = c(1, 2)) |>
  mutate(p = ifelse(is.na(p), (d * (k - 1)) %/% 2, p),
         H_in = 32, H_out = out_size(32, k, p, s, d),
         preserves_size = H_out == 32)
head(cfg[order(cfg$k, cfg$s, cfg$d), ], 12)
# Verify against torch rather than trusting the formula
verify <- function(k, p, s, d) {
  x <- torch_randn(1, 1, 32, 32)
  dim(nnf_conv2d(x, torch_randn(1, 1, k, k), padding = p, stride = s, dilation = d))[3]
}
data.frame(k = c(3,3,5,3), p = c(1,0,2,2), s = c(1,1,1,1), d = c(1,1,1,2),
           formula = mapply(out_size, 32, c(3,3,5,3), c(1,0,2,2), 1, c(1,1,1,2)),
           torch = mapply(verify, c(3,3,5,3), c(1,0,2,2), 1, c(1,1,1,2)))
# Where a dilated kernel actually samples
grid_pts <- function(k, d, cx = 8, cy = 8) {
  off <- (seq_len(k) - (k + 1)/2) * d
  expand.grid(dx = off, dy = off) |> mutate(x = cx + dx, y = cy + dy)
}
bind_rows(mutate(grid_pts(3, 1), cfg = "3x3, dilation 1  (RF 3)"),
          mutate(grid_pts(3, 2), cfg = "3x3, dilation 2  (RF 5)"),
          mutate(grid_pts(3, 4), cfg = "3x3, dilation 4  (RF 9)")) |>
  ggplot(aes(x, y)) +
  geom_tile(data = expand.grid(x = 1:15, y = 1:15), fill = "grey93",
            color = "white", linewidth = 0.3) +
  geom_tile(fill = "#3B7DD8", color = "white", linewidth = 0.4) +
  facet_wrap(~ cfg) + coord_fixed() +
  labs(title = "Dilation enlarges the receptive field without adding parameters",
       subtitle = "All three kernels have exactly 9 weights; only their sampling spacing differs",
       x = NULL, y = NULL) +
  theme_dspa(9) + theme(axis.text = element_blank(), panel.grid = element_blank())

14.27 Receptive fields

The receptive field of a unit is the region of the input that can affect it. It is what determines the largest structure a layer can possibly detect.

Stacking layers grows it. With kernel \(k_\ell\) and stride \(s_\ell\) at layer \(\ell\), the receptive field after \(L\) layers is

\[\boxed{\;r_L=1+\sum_{\ell=1}^{L}\big(k_\ell-1\big)\prod_{m=1}^{\ell-1}s_m\;}\]

Two regimes follow directly. With all strides 1, \(r_L=1+\sum(k_\ell-1)\) grows linearly in depth. With strides \(>1\), the product term makes it grow geometrically, which is why downsampling is how networks reach global context in a manageable number of layers.

rf <- function(kernels, strides) {
  r <- 1; jump <- 1
  for (i in seq_along(kernels)) { r <- r + (kernels[i] - 1) * jump
                                  jump <- jump * strides[i] }
  r
}
data.frame(
  architecture = c("10 x (3x3, s=1)", "10 x (3x3, s=1) with pool every 2",
                   "5 x (5x5, s=1)", "3 x (3x3, s=2)"),
  receptive_field = c(rf(rep(3,10), rep(1,10)),
                      rf(rep(3,10), rep(c(1,2), 5)),
                      rf(rep(5,5), rep(1,5)),
                      rf(rep(3,3), rep(2,3))),
  parameters_per_channel_pair = c(10*9, 10*9, 5*25, 3*9))

Two \(3\times3\) layers have the same receptive field as one \(5\times5\) layer (\(r=5\)) with fewer parameters (\(18\) versus \(25\) per channel pair) and an extra nonlinearity between them. That observation is the entire design principle of VGG (§14.31).

depths_rf <- 1:20
kernels_rf <- c(3, 5, 7, 9, 11)
Zrf <- outer(kernels_rf, depths_rf, Vectorize(function(k, L)
  rf(rep(k, L), rep(1, L))))
Zrf_stride <- outer(kernels_rf, depths_rf, Vectorize(function(k, L)
  rf(rep(k, L), rep(2, L))))

plot_ly() |>
  add_surface(x = depths_rf, y = kernels_rf, z = log10(Zrf_stride), opacity = 0.9,
              showscale = FALSE, colorscale = "Reds", name = "stride 2") |>
  add_surface(x = depths_rf, y = kernels_rf, z = log10(Zrf), opacity = 0.9,
              showscale = FALSE, colorscale = "Blues", name = "stride 1") |>
  layout(title = "Receptive field growth: stride 2 (upper, red) vs. stride 1 (lower, blue)",
         scene = list(xaxis = list(title = "Depth (layers)"),
                      yaxis = list(title = "Kernel size"),
                      zaxis = list(title = "log10 receptive field")))

Rotate along the depth axis. The stride-1 surface rises linearly on this log scale, that is, it grows like a polynomial. The stride-2 surface rises as a straight line in \(\log\), meaning geometric growth: 20 stride-2 layers reach a receptive field of millions of pixels, far past any real image, while 20 stride-1 layers reach only 41.

Common misconception: “compute the receptive field and you know what the network sees.” The effective receptive field is far smaller than the theoretical one. Contributions from the edge of the theoretical field pass through exponentially many paths that mostly cancel, so the actual influence decays roughly as a Gaussian from the centre and the effective radius grows only as \(O(\sqrt L)\) rather than \(O(L)\) (Luo et al., 2016).

The practical implication: computing \(r_L\) and concluding that a network “sees the whole image” overstates what it uses. Dilated convolutions, larger strides, and attention (Part 3) exist because stacking \(3\times3\) layers reaches global context far more slowly than the formula suggests.

# Measure the effective receptive field by gradient magnitude at the input
torch_manual_seed(21)
stack_n <- function(n) do.call(nn_sequential, unlist(lapply(seq_len(n), \(i)
  list(nn_conv2d(1, 1, 3, padding = 1), nn_relu())), recursive = FALSE))

erf_map <- function(n_layers, size = 65) {
  net <- stack_n(n_layers)
  x <- torch_zeros(1, 1, size, size, requires_grad = TRUE)
  out <- net(x)
  centre <- out[1, 1, (size + 1) %/% 2, (size + 1) %/% 2]
  centre$backward()
  abs(as.array(x$grad$squeeze()))
}
erf_df <- bind_rows(lapply(c(3, 8, 15), function(n) {
  m <- erf_map(n); m <- m / max(m)
  expand.grid(row = 1:nrow(m), col = 1:ncol(m)) |>
    mutate(v = as.vector(m),
           panel = sprintf("%d layers  (theoretical RF = %d)", n, rf(rep(3,n), rep(1,n))))
}))

ggplot(erf_df, aes(col, -row, fill = v)) +
  geom_raster() + facet_wrap(~ panel) + coord_fixed() +
  scale_fill_viridis_c(option = "inferno", guide = "none") +
  labs(title = "Effective receptive field, measured as input-gradient magnitude",
       subtitle = "The bright core is far smaller than the theoretical square, and roughly Gaussian",
       x = NULL, y = NULL) +
  theme_dspa(9) + theme(axis.text = element_blank(), panel.grid = element_blank())

# --- Interactive equivalent ------------------------------------------------
m <- erf_map(15); m <- m / max(m)
plot_ly(z = m, type = "heatmap", colorscale = "Inferno") |>
  layout(title = "Effective receptive field, 15 layers of 3x3 convolution",
         xaxis = list(scaleanchor = "y"))

14.28 Downsampling: pooling and strided convolution

Reducing spatial resolution serves three purposes: it grows the receptive field geometrically (§14.27), it cuts activation memory, and it discards spatial precision the task may not need.

\[ \begin{aligned} \textbf{Max pooling: }&\quad y_{u,v}=\max_{(p,q)\in\mathcal N(u,v)} x_{p,q} &&\text{no parameters}\\ \textbf{Average pooling: }&\quad y_{u,v}=\tfrac{1}{|\mathcal N|}\textstyle\sum_{(p,q)\in\mathcal N(u,v)} x_{p,q} &&\text{no parameters}\\ \textbf{Strided convolution: }&\quad \text{convolve with }s>1 &&\text{\emph{learned} downsampling}\\ \textbf{Global average pooling: }&\quad y_{c}=\tfrac{1}{HW}\textstyle\sum_{u,v} x_{c,u,v} &&\text{collapses space entirely} \end{aligned} \]

Common misconception: “pooling is a required part of a CNN.” It is one of several ways to downsample, and it is the least flexible: it has no parameters, so it cannot adapt what it discards to the task.

Strided convolution downsamples and learns the projection, and networks built entirely from strided convolutions match or beat pooled ones (Springenberg et al., 2015). Most architectures since ResNet use pooling sparingly, often a single early max-pool, then strided convolutions throughout.

Global average pooling deserves separate mention. Replacing the flatten plus dense head with a channel-wise spatial mean removes the largest parameter block in the network, VGG-16’s dense head is 90% of its parameters, and makes the network accept any input size, because the output no longer depends on \(H\) and \(W\). It also acts as a structural regularizer, forcing each channel to correspond to a class-relevant concept rather than to an arbitrary spatial position.

x_pool <- torch_randn(1, 16, 32, 32)

variants <- list(
  `max pool 2x2`        = nn_max_pool2d(2),
  `avg pool 2x2`        = nn_avg_pool2d(2),
  `strided conv 3x3 s2` = nn_conv2d(16, 16, 3, stride = 2, padding = 1),
  `global avg pool`     = nn_adaptive_avg_pool2d(1))

data.frame(
  method = names(variants),
  output_shape = vapply(variants, \(m)
    paste(dim(with_no_grad(m(x_pool)))[-1], collapse = " x "), character(1)),
  parameters = vapply(variants, \(m) {
    p <- m$parameters
    if (!length(p)) 0 else sum(vapply(p, \(q) prod(dim(q)), numeric(1)))
  }, numeric(1)),
  row.names = NULL)
# The dense head dominates a classical architecture's parameter count
conv_body <- 14714688       # VGG-16 convolutional layers
dense_head <- 25088 * 4096 + 4096 * 4096 + 4096 * 1000
gap_head <- 512 * 1000
data.frame(
  component = c("convolutional body", "dense head (VGG-style)",
                "global-average-pool head"),
  parameters = format(c(conv_body, dense_head, gap_head), big.mark = ","),
  share_of_VGG16 = c(round(conv_body/(conv_body+dense_head), 3),
                     round(dense_head/(conv_body+dense_head), 3), NA))

14.29 Complexity of a convolutional layer

For input \(C_{\text{in}}\times H\times W\), output \(C_{\text{out}}\times H'\times W'\), kernel \(k\times k\), and batch \(B\):

\[ \begin{aligned} \textbf{Parameters: }&\quad C_{\text{out}}\big(C_{\text{in}}k^2+1\big)\\ \textbf{FLOPs (forward): }&\quad \approx 2\,B\,C_{\text{out}}H'W'\,C_{\text{in}}k^2\\ \textbf{Activation memory: }&\quad 4\,B\,C_{\text{out}}H'W'\ \text{bytes at float32}\\ \textbf{Arithmetic intensity: }&\quad I=\dfrac{\text{FLOPs}}{\text{bytes moved}}\approx\dfrac{2C_{\text{in}}k^2}{4\big(C_{\text{in}}k^2/C_{\text{out}}+1\big)} \end{aligned} \]

CNNs are compute-bound where MLPs are memory-bound, and the reason is weight reuse. A dense layer uses each weight once per example, giving arithmetic intensity around \(0.5\) flops per byte, memory-bound, on the bandwidth roof of Chapter 10, §10.14.1. A convolutional layer reuses each weight \(H'W'\) times, so its intensity is tens to hundreds of flops per byte, placing it above the ridge point.

That difference is why GPUs help far more with CNNs than with MLPs: on the compute roof a GPU’s advantage is its flop rate, which is one to two orders of magnitude over a CPU; on the bandwidth roof the advantage is only the bandwidth ratio, roughly \(5\)\(20\times\).

conv_cost <- function(C_in, C_out, H, W, k, B = 32, s = 1) {
  Ho <- H %/% s; Wo <- W %/% s
  params <- C_out * (C_in * k^2 + 1)
  flops <- 2 * B * C_out * Ho * Wo * C_in * k^2
  act_bytes <- 4 * B * C_out * Ho * Wo
  weight_bytes <- 4 * params
  c(params = params, MFLOPs = flops/1e6, activation_MB = act_bytes/1e6,
    intensity = flops / (act_bytes + weight_bytes + 4*B*C_in*H*W))
}
dense_cost <- function(d_in, d_out, B = 32) {
  params <- d_in*d_out + d_out
  flops <- 2*B*d_in*d_out
  c(params = params, MFLOPs = flops/1e6, activation_MB = 4*B*d_out/1e6,
    intensity = flops / (4*(B*d_in + B*d_out + params)))
}
as.data.frame(rbind(
  `conv 3x3, 64->64, 56x56`   = conv_cost(64, 64, 56, 56, 3),
  `conv 3x3, 256->256, 14x14` = conv_cost(256, 256, 14, 14, 3),
  `conv 1x1, 256->64, 56x56`  = conv_cost(256, 64, 56, 56, 1),
  `dense 4096 -> 4096`        = dense_cost(4096, 4096))) |>
  mutate(across(everything(), \(z) signif(z, 4)))

Read the intensity column against the roofline: the convolutional rows sit well above a typical ridge point of 5–20, the dense row well below.

14.29.1 The \(1\times1\) convolution

A \(1\times1\) kernel has no spatial extent, so it might appear to do nothing. It is in fact a learned linear map across channels, applied identically at every spatial position, a per-pixel dense layer.

Its uses are structural. It changes channel count cheaply, so a \(3\times3\) convolution can be sandwiched between a \(1\times1\) that reduces channels and one that restores them, the bottleneck block that makes ResNet-50 and deeper feasible.

plain <- conv_cost(256, 256, 56, 56, 3)
bottleneck <- conv_cost(256, 64, 56, 56, 1) + conv_cost(64, 64, 56, 56, 3) +
              conv_cost(64, 256, 56, 56, 1)
data.frame(
  block = c("plain 3x3, 256 -> 256", "bottleneck 1x1-3x3-1x1"),
  parameters = format(c(plain[["params"]], bottleneck[["params"]]), big.mark = ","),
  MFLOPs = round(c(plain[["MFLOPs"]], bottleneck[["MFLOPs"]])),
  reduction = c(1, round(plain[["MFLOPs"]] / bottleneck[["MFLOPs"]], 2)))

The bottleneck computes an equivalent transformation for roughly a quarter of the cost. Depthwise separable convolution (MobileNet) pushes the same idea further, factoring the spatial and channel mixing entirely.

14.30 Equivariance and invariance

These two words are used interchangeably in casual writing and mean different things.

\[ \begin{aligned} \textbf{Equivariant: }&\quad f\big(T_\delta(x)\big)=T_\delta\big(f(x)\big) &&\text{shift the input, the output shifts}\\ \textbf{Invariant: }&\quad f\big(T_\delta(x)\big)=f(x) &&\text{shift the input, the output is unchanged} \end{aligned} \]

Common misconception: “CNNs are translation-invariant.” A convolutional layer is translation-equivariant, which is a different and weaker property. Shift the input by \(\delta\) pixels and every feature map shifts by \(\delta\) pixels, the representation moves with the object rather than ignoring where it is.

Invariance, when a network has it, comes from elsewhere: global pooling collapses the spatial dimensions and so discards position; local pooling gives approximate invariance to small shifts only; and training with random translations teaches approximate invariance without guaranteeing it.

Worse, the equivariance itself is imperfect in practice. Strided layers alias, they sample the feature map below its Nyquist rate, so a one-pixel shift of the input can change the output substantially, and CNN predictions are measurably unstable under small translations (Azulay & Weiss, 2019; Zhang, 2019). Anti-aliased downsampling, blur before subsampling, restores much of it.

torch_manual_seed(31)
base_img <- torch_zeros(1, 1, 40, 40)
base_img[1, 1, 14:26, 14:20] <- 1

shift_img <- function(t, dx) {
  out <- torch_zeros_like(t)
  H <- dim(t)[4]
  src <- max(1, 1 - dx):min(H, H - dx)
  out[, , , src + dx] <- t[, , , src]
  out
}

conv_eq <- nn_conv2d(1, 1, 5, padding = 2)
with_no_grad({
  f_shift_x  <- conv_eq(shift_img(base_img, 6))          # f(T(x))
  shift_f_x  <- shift_img(conv_eq(base_img), 6)          # T(f(x))
})
c(equivariance_max_abs_error =
    signif(as.numeric(torch_max(torch_abs(
      f_shift_x[, , , 8:34] - shift_f_x[, , , 8:34]))), 3),
  interpretation = "f(T(x)) = T(f(x)) up to border effects")
#>               equivariance_max_abs_error 
#>                                      "0" 
#>                           interpretation 
#> "f(T(x)) = T(f(x)) up to border effects"
# Invariance requires a pooling or aggregation step
with_no_grad({
  gap_orig  <- as.numeric(torch_mean(conv_eq(base_img)))
  gap_shift <- as.numeric(torch_mean(conv_eq(shift_img(base_img, 6))))
})
c(global_mean_original = signif(gap_orig, 5),
  global_mean_shifted = signif(gap_shift, 5),
  note = "global pooling is what produces (approximate) invariance")
#>                                       global_mean_original 
#>                                                  "0.02052" 
#>                                        global_mean_shifted 
#>                                                  "0.02052" 
#>                                                       note 
#> "global pooling is what produces (approximate) invariance"
# Strided pooling aliases: a one-pixel shift changes the output measurably
torch_manual_seed(33)
net_strided <- nn_sequential(nn_conv2d(1, 8, 3, padding = 1), nn_relu(),
                             nn_max_pool2d(2),
                             nn_conv2d(8, 8, 3, padding = 1), nn_relu(),
                             nn_max_pool2d(2),
                             nn_adaptive_avg_pool2d(1), nn_flatten(),
                             nn_linear(8, 1))
shifts <- 0:12
resp <- vapply(shifts, \(d)
  as.numeric(with_no_grad(net_strided(shift_img(base_img, d)))), numeric(1))

ggplot(data.frame(shift = shifts, output = resp), aes(shift, output)) +
  geom_hline(yintercept = resp[1], linetype = "dashed", color = "grey40") +
  geom_line(linewidth = 1, color = "steelblue") + geom_point(size = 2.2) +
  scale_x_continuous(breaks = shifts) +
  labs(title = "A truly invariant network would draw a flat line",
       subtitle = "Same object, translated one pixel at a time. Strided pooling aliases, so the output oscillates",
       x = "Horizontal shift (pixels)", y = "Network output") +
  theme_dspa()

c(output_range_across_shifts = signif(diff(range(resp)), 3),
  relative_variation = signif(diff(range(resp)) / abs(mean(resp)), 3))
#> output_range_across_shifts         relative_variation 
#>                    0.00397                    0.01110
# --- Interactive equivalent ------------------------------------------------
plot_ly(x = shifts, y = resp, type = "scatter", mode = "lines+markers",
        name = "Network output") |>
  add_lines(x = range(shifts), y = rep(resp[1], 2), name = "Perfect invariance",
            line = list(dash = "dash", color = "grey")) |>
  layout(title = "Sensitivity to translation",
         xaxis = list(title = "Shift (pixels)"),
         yaxis = list(title = "Output"))

The output should be constant if the network were invariant. It is not, and the oscillation has period 4, matching the two \(2\times2\) pooling stages. This is aliasing, not noise.


PART II: ARCHITECTURE

14.31 From LeNet to VGG

Three decades of architecture search converged on a small set of principles.

Network Year Depth Parameters The idea it contributed
LeNet-5 1998 7 60 K Convolution + pooling + dense, trained end to end
AlexNet 2012 8 60 M ReLU, dropout, GPU training, aggressive augmentation
VGG-16 2014 16 138 M Only \(3\times3\) kernels, stacked uniformly
GoogLeNet 2014 22 6.8 M Parallel multi-scale branches; \(1\times1\) bottlenecks
ResNet-50 2015 50 25 M Residual connections; depth beyond 100 layers
DenseNet 2017 121 8 M Concatenative rather than additive skips
ConvNeXt 2022 29–198 M CNN modernized with Transformer-era design choices

Two trends run through the table. Parameter counts fell after VGG, 138 M down to 25 M, because the dense head was replaced by global average pooling (§14.28) and bottlenecks reduced channel width (§14.29.1). And kernels shrank to \(3\times3\), because two stacked \(3\times3\) layers match one \(5\times5\) receptive field with fewer parameters and an extra nonlinearity (§14.27).

lenet5 <- nn_module(
  "LeNet5",
  initialize = function(n_classes = 10) {
    self$features <- nn_sequential(
      nn_conv2d(1, 6, 5, padding = 2), nn_relu(), nn_avg_pool2d(2),
      nn_conv2d(6, 16, 5), nn_relu(), nn_avg_pool2d(2))
    self$classifier <- nn_sequential(
      nn_flatten(), nn_linear(16*5*5, 120), nn_relu(),
      nn_linear(120, 84), nn_relu(), nn_linear(84, n_classes))
  },
  forward = function(x) self$classifier(self$features(x)))

m_lenet <- lenet5()
n_par <- function(m) sum(vapply(m$parameters, \(p) prod(dim(p)), numeric(1)))
c(total_parameters = n_par(m_lenet),
  feature_extractor = n_par(m_lenet$features),
  classifier_head = n_par(m_lenet$classifier),
  head_share = round(n_par(m_lenet$classifier) / n_par(m_lenet), 3),
  output_shape = paste(dim(with_no_grad(m_lenet(torch_randn(1,1,28,28)))),
                       collapse = " x "))
#>  total_parameters feature_extractor   classifier_head        head_share 
#>           "61706"            "2572"           "59134"           "0.958" 
#>      output_shape 
#>          "1 x 10"

Even in LeNet the dense head carries most of the parameters, the pattern that global average pooling later removed.

14.32 Residual connections

By 2015 the obstacle was not expressiveness but optimization. Stacking more plain layers made training accuracy worse.

Common misconception: “a deeper network that performs worse is overfitting.” The degradation problem is the opposite. A 56-layer plain network has higher training error than a 20-layer one (He et al., 2016), it fits the training data less well, which overfitting cannot explain.

It is not a representational limit either. The deeper network contains the shallower one: set the extra layers to the identity and the two are the same function. So a solution at least as good provably exists, and gradient descent fails to find it. The problem is that learning the identity map is hard for a stack of nonlinear layers, every layer must conspire to produce \(\mathcal H(x)=x\), and nothing in the parameterization makes that easy.

Residual connections change the parameterization so that the identity is free.

Instead of learning the target map \(\mathcal H(x)\) directly, a residual block learns the residual \(\mathcal F(x)=\mathcal H(x)-x\) and adds the input back:

\[\boxed{\;y=\mathcal F(x,\{W_i\})+x\;}\]

If the identity is optimal, the block only needs \(\mathcal F\to0\), driving weights toward zero, which weight decay already does. The hard case became the easy one.

The gradient-flow argument. With \(x_L=x_\ell+\sum_{i=\ell}^{L-1}\mathcal F(x_i)\), the chain rule gives

\[\frac{\partial\mathcal L}{\partial x_\ell}=\frac{\partial\mathcal L}{\partial x_L}\left(1+\frac{\partial}{\partial x_\ell}\sum_{i=\ell}^{L-1}\mathcal F(x_i)\right).\]

The 1 is the point. The gradient reaches layer \(\ell\) through an unobstructed additive path regardless of what the intervening \(\mathcal F\) do, so it cannot vanish as the product of Jacobians in a plain stack does (Part 1, §14.8).

This is the same mechanism as the LSTM cell state. In Chapter 12, §12.27 the update \(\mathbf c_t=\mathbf f_t\odot\mathbf c_{t-1}+\mathbf i_t\odot\tilde{\mathbf c}_t\) creates an additive path through time with derivative \(\operatorname{diag}(\mathbf f_t)\) rather than a weight matrix. A residual block creates an additive path through depth with derivative exactly 1. Different axis, identical idea: replace a product of Jacobians with a sum, so gradients survive.

residual_block <- nn_module(
  "ResidualBlock",
  initialize = function(ch) {
    self$body <- nn_sequential(
      nn_conv2d(ch, ch, 3, padding = 1, bias = FALSE), nn_batch_norm2d(ch),
      nn_relu(),
      nn_conv2d(ch, ch, 3, padding = 1, bias = FALSE), nn_batch_norm2d(ch))
    self$act <- nn_relu()
  },
  forward = function(x) self$act(self$body(x) + x))   # the "+ x" is the whole idea

plain_block <- nn_module(
  "PlainBlock",
  initialize = function(ch) {
    self$body <- nn_sequential(
      nn_conv2d(ch, ch, 3, padding = 1, bias = FALSE), nn_batch_norm2d(ch),
      nn_relu(),
      nn_conv2d(ch, ch, 3, padding = 1, bias = FALSE), nn_batch_norm2d(ch),
      nn_relu())
  },
  forward = function(x) self$body(x))

stack_blocks <- function(block_fn, n, ch = 16) {
  do.call(nn_sequential, lapply(seq_len(n), \(i) block_fn(ch)))
}
c(residual_block_params = n_par(residual_block(16)),
  plain_block_params = n_par(plain_block(16)),
  note = "identical parameter counts; only the forward pass differs")
#>                                       residual_block_params 
#>                                                      "4672" 
#>                                          plain_block_params 
#>                                                      "4672" 
#>                                                        note 
#> "identical parameter counts; only the forward pass differs"
grad_by_depth <- function(block_fn, n_blocks, ch = 16, seed = 41) {
  torch_manual_seed(seed)
  net <- stack_blocks(block_fn, n_blocks, ch)
  x <- torch_randn(8, ch, 16, 16)
  net(x)$sum()$backward()
  convs <- Filter(\(m) inherits(m, "nn_conv2d"), net$modules)
  g <- vapply(convs, \(m) if (is.null(m$weight$grad)) NA_real_ else
    as.numeric(torch_norm(m$weight$grad)), numeric(1))
  g[!is.na(g)]
}

gd <- bind_rows(
  data.frame(block = seq_along(grad_by_depth(residual_block, 24)),
             g = grad_by_depth(residual_block, 24), arch = "Residual"),
  data.frame(block = seq_along(grad_by_depth(plain_block, 24)),
             g = grad_by_depth(plain_block, 24), arch = "Plain"))

ggplot(gd, aes(block, pmax(g, 1e-20), color = arch)) +
  geom_line(linewidth = 0.9) + geom_point(size = 1.1) +
  scale_y_log10() +
  scale_color_manual(values = c(Residual = "#3B7DD8", Plain = "#D8433B")) +
  labs(title = "Gradient norm by layer through 24 blocks (48 convolutions)",
       subtitle = "Layer 1 is nearest the input. The plain stack's gradient decays; the residual stack's does not",
       x = "Convolution index (1 = closest to input)",
       y = "Gradient norm (log scale)", color = NULL) +
  theme_dspa()

gd |> summarise(first_layer = first(g), last_layer = last(g),
                decay_ratio = signif(last(g) / first(g), 3), .by = arch)
depths_res <- c(4, 8, 12, 16, 24, 32)
Zplain <- t(vapply(depths_res, function(n) {
  g <- grad_by_depth(plain_block, n)
  approx(seq_along(g)/length(g), log10(pmax(g, 1e-25)), xout = seq(0, 1, length.out = 20),
         rule = 2)$y
}, numeric(20)))
Zres <- t(vapply(depths_res, function(n) {
  g <- grad_by_depth(residual_block, n)
  approx(seq_along(g)/length(g), log10(pmax(g, 1e-25)), xout = seq(0, 1, length.out = 20),
         rule = 2)$y
}, numeric(20)))

plot_ly() |>
  add_surface(x = seq(0, 1, length.out = 20), y = depths_res, z = Zres,
              opacity = 0.9, showscale = FALSE, colorscale = "Blues",
              name = "Residual") |>
  add_surface(x = seq(0, 1, length.out = 20), y = depths_res, z = Zplain,
              opacity = 0.9, showscale = FALSE, colorscale = "Reds",
              name = "Plain") |>
  layout(title = "Gradient magnitude over relative depth and total depth: residual (blue) vs. plain (red)",
         scene = list(xaxis = list(title = "Relative position (0 = input)"),
                      yaxis = list(title = "Total blocks"),
                      zaxis = list(title = "log10 gradient norm")))

Rotate to the large-depth edge. The plain surface falls away toward the input side and the fall deepens as total depth grows, the gradient reaching early layers shrinks with every block added. The residual surface stays nearly flat in both directions. That flatness is what made 100-plus-layer networks trainable.

# The degradation problem: deeper plain networks reach WORSE training loss
train_stack <- function(block_fn, n_blocks, steps = 250, seed = 43) {
  torch_manual_seed(seed); set.seed(seed)
  X <- torch_randn(64, 16, 16, 16); Y <- torch_randn(64, 1)
  net <- nn_sequential(stack_blocks(block_fn, n_blocks),
                       nn_adaptive_avg_pool2d(1), nn_flatten(), nn_linear(16, 1))
  opt <- optim_adamw(net$parameters, lr = 3e-4)
  hist <- numeric(steps)
  for (s in seq_len(steps)) {
    opt$zero_grad(); l <- nnf_mse_loss(net(X), Y); l$backward(); opt$step()
    hist[s] <- as.numeric(l)
  }
  hist
}
deg <- bind_rows(lapply(c(4, 16, 32), function(n)
  bind_rows(data.frame(step = 1:250, loss = train_stack(plain_block, n),
                       arch = "Plain", depth = sprintf("%d blocks", n)),
            data.frame(step = 1:250, loss = train_stack(residual_block, n),
                       arch = "Residual", depth = sprintf("%d blocks", n)))))

ggplot(deg, aes(step, loss, color = arch)) +
  geom_line(linewidth = 0.8) +
  facet_wrap(~ factor(depth, levels = sprintf("%d blocks", c(4,16,32)))) +
  scale_y_log10() +
  scale_color_manual(values = c(Residual = "#3B7DD8", Plain = "#D8433B")) +
  labs(title = "Training loss, not validation loss",
       subtitle = "The plain network gets WORSE with depth on data it is trying to memorize -- an optimization failure, not overfitting",
       x = "Step", y = "Training MSE (log scale)", color = NULL) +
  theme_dspa(10)

deg |> filter(step == 250) |> select(arch, depth, final_loss = loss) |>
  pivot_wider(names_from = arch, values_from = final_loss) |>
  mutate(across(where(is.numeric), \(z) signif(z, 4)))

14.33 Normalization in convolutional networks

Part 1, §14.10 introduced batch and layer normalization. In a convolutional network the normalization axis is a design choice with visible consequences.

For an activation tensor of shape \((B,C,H,W)\):

\[ \begin{aligned} \textbf{BatchNorm: }&\quad \text{statistics over }(B,H,W)\text{, per channel} &&\text{2}C\text{ parameters}\\ \textbf{LayerNorm: }&\quad \text{over }(C,H,W)\text{, per example} &&\text{batch-independent}\\ \textbf{InstanceNorm: }&\quad \text{over }(H,W)\text{, per example per channel} &&\text{style transfer}\\ \textbf{GroupNorm: }&\quad \text{over }(C/G,H,W)\text{, per example per group} &&\text{batch-independent} \end{aligned} \]

BatchNorm’s dependence on the batch is a liability in vision. Segmentation and detection use large images and therefore small batches, often 2 to 8 per GPU, and at that size the batch statistics are too noisy to normalize with (Part 1, §14.10). GroupNorm was designed for exactly this case: it matches BatchNorm’s accuracy at batch size 32 and substantially beats it at batch size 2 (Wu & He, 2018).

Two further BatchNorm traps in vision. Its running statistics are updated only in train() mode and used only in eval() mode, so a mode error changes the numbers without raising an error. And during transfer learning (§14.35) a frozen BatchNorm layer still updates its running statistics unless it is explicitly put in eval() mode, a silent distribution shift toward the new data that “frozen” implies should not happen.

# xn <- torch_randn(4, 32, 16, 16) * 3 + 5
# norms <- list(BatchNorm = nn_batch_norm2d(32), LayerNorm = nn_layer_norm(c(32,16,16)),
#               InstanceNorm = nn_instance_norm2d(32), GroupNorm = nn_group_norm(8, 32))
# data.frame(
#   normalization = names(norms),
#   statistics_over = c("(B, H, W) per channel", "(C, H, W) per example",
#                       "(H, W) per example-channel", "(C/G, H, W) per example-group"),
#   depends_on_batch = c(TRUE, FALSE, FALSE, FALSE),
#   parameters = vapply(norms, \(m) {
#     p <- m$parameters; if (!length(p)) 0 else sum(vapply(p, \(q) prod(dim(q)), numeric(1)))
#   }, numeric(1)),
#   row.names = NULL)
# Define nn_instance_norm2d if it's missing in your torch version
if (!exists("nn_instance_norm2d", mode = "function")) {
  nn_instance_norm2d <- function(num_features, eps = 1e-5, momentum = 0.1, 
                                 affine = FALSE, track_running_stats = FALSE) {
    nn_module(
      initialize = function() {
        self$num_features <- num_features
        self$eps <- eps
        self$momentum <- momentum
        self$affine <- affine
        self$track_running_stats <- track_running_stats
        
        # Initialize learnable parameters if affine = TRUE
        if (self$affine) {
          self$weight <- nn_parameter(torch_ones(num_features))
          self$bias <- nn_parameter(torch_zeros(num_features))
        }
        # Initialize running stats if track_running_stats = TRUE
        if (self$track_running_stats) {
          self$running_mean <- torch_zeros(num_features)
          self$running_var <- torch_ones(num_features)
        }
      },
      forward = function(x) {
        nnf_instance_norm(
          x,
          running_mean = if (self$track_running_stats) self$running_mean else NULL,
          running_var = if (self$track_running_stats) self$running_var else NULL,
          weight = if (self$affine) self$weight else NULL,
          bias = if (self$affine) self$bias else NULL,
          use_input_stats = self$training || !self$track_running_stats,
          momentum = self$momentum,
          eps = self$eps
        )
      }
    )
  }
}

xn <- torch_randn(4, 32, 16, 16) * 3 + 5
norms <- list(BatchNorm = nn_batch_norm2d(32), LayerNorm = nn_layer_norm(c(32,16,16)),
              InstanceNorm = nn_instance_norm2d(32), GroupNorm = nn_group_norm(8, 32))

data.frame(
  normalization = names(norms),
  statistics_over = c("(B, H, W) per channel", "(C, H, W) per example",
                      "(H, W) per example-channel", "(C/G, H, W) per example-group"),
  depends_on_batch = c(TRUE, FALSE, FALSE, FALSE),
  parameters = vapply(norms, \(m) {
    p <- m$parameters; if (!length(p)) 0 else sum(vapply(p, \(q) prod(dim(q)), numeric(1)))
  }, numeric(1)),
  row.names = NULL)
# Stability of the normalized output as batch size shrinks
stability <- function(bs, reps = 60, seed = 51) {
  torch_manual_seed(seed)
  pop <- torch_randn(1024, 32, 8, 8)
  bn <- nn_batch_norm2d(32); gn <- nn_group_norm(8, 32)
  bn$train()
  out <- vapply(seq_len(reps), function(r) {
    idx <- sample(1024, bs)
    xb <- pop[idx, , , ]
    c(as.numeric(torch_mean(bn(xb)[1, , , ])), as.numeric(torch_mean(gn(xb)[1, , , ])))
  }, numeric(2))
  c(batch = bs, BatchNorm = sd(out[1, ]), GroupNorm = sd(out[2, ]))
}
st <- as.data.frame(do.call(rbind, lapply(c(2, 4, 8, 16, 32, 64), stability)))
st |> mutate(across(-batch, \(z) signif(z, 4)))
st |> pivot_longer(-batch, names_to = "method", values_to = "sd") |>
  ggplot(aes(batch, sd, color = method)) +
  geom_line(linewidth = 1) + geom_point(size = 2.4) +
  scale_x_log10(breaks = st$batch) + scale_y_log10() +
  scale_color_manual(values = c(BatchNorm = "#D8433B", GroupNorm = "#3B7DD8")) +
  labs(title = "Variability of a fixed example's normalized output",
       subtitle = "BatchNorm's output for one image depends on which others share its batch; GroupNorm's does not",
       x = "Batch size (log scale)", y = "SD across batches (log scale)",
       color = NULL) +
  theme_dspa()

The GroupNorm line is flat, an example’s normalized output does not depend on its batch-mates at all. The BatchNorm line rises sharply below batch 8, which is precisely the regime segmentation operates in (§14.37).


PART III: TRAINING VISION MODELS

14.34 Data augmentation

Augmentation applies label-preserving transformations to training inputs: flips, crops, rotations, color jitter, and their compositions.

Common misconception: “augmentation creates more training data.” It creates more training examples and no additional information. Every augmented image is a deterministic or stochastic function of one original, so the sample still contains exactly the information those originals carried.

What augmentation does is encode an invariance as a constraint on the hypothesis space. Training on horizontally flipped images tells the model that the label does not depend on left–right orientation, which shrinks the set of functions consistent with the data. That is regularization (Part 1, §14.13), and it is why augmentation reduces overfitting without reducing capacity.

The invariance must be true, and this is where augmentation goes wrong in practice. Horizontal flips are label-preserving for natural photographs and not for text, for digits (2 and 5), or for chest radiographs, where situs inversus is a finding rather than a nuisance. Rotation is safe for pathology slides, which have no canonical orientation, and unsafe for street scenes. Aggressive color jitter destroys the diagnostic signal in stained histology. Each transformation is a claim about the problem, and an augmentation pipeline copied from ImageNet asserts ImageNet’s claims about your data.

torch_manual_seed(61)
src <- torch_zeros(1, 24, 24)
src[1, 6:18, 8:12] <- 1; src[1, 6:8, 8:16] <- 1        # an asymmetric "F"-like shape

show_t <- function(t, lab) {
  m <- as.array(t$squeeze())
  expand.grid(row = 1:nrow(m), col = 1:ncol(m)) |>
    mutate(v = as.vector(m), panel = lab)
}
augs <- bind_rows(
  show_t(src, "original"),
  show_t(transform_hflip(src), "horizontal flip"),
  show_t(transform_rotate(src, 25), "rotate 25 deg"),
  show_t(transform_resized_crop(src, 3, 3, 18, 18, size = c(24, 24)), "random crop"))

ggplot(augs, aes(col, -row, fill = v)) +
  geom_raster() + facet_wrap(~ panel, nrow = 1) + coord_fixed() +
  scale_fill_viridis_c(guide = "none") +
  labs(title = "Each augmentation asserts an invariance",
       subtitle = "Whether the label survives the transformation is a claim about the domain, not a setting",
       x = NULL, y = NULL) +
  theme_dspa(9) + theme(axis.text = element_blank(), panel.grid = element_blank())

# Augmentation helps most where data are scarce, and saturates
# set.seed(63); torch_manual_seed(63)
# synth_task <- function(n, aug_strength, seed = 65) {
#   torch_manual_seed(seed); set.seed(seed)
#   gen <- function(m) {
#     lab <- sample(0:1, m, TRUE)
#     X <- torch_zeros(m, 1, 16, 16)
#     for (i in seq_len(m)) {
#       cx <- sample(4:11, 1); cy <- sample(4:11, 1)
#       if (lab[i] == 1) X[i, 1, cx:(cx+3), cy:(cy+1)] <- 1
#       else             X[i, 1, cx:(cx+1), cy:(cy+3)] <- 1
#     }
#     list(X = X + torch_randn(m, 1, 16, 16) * 0.25,
#          y = torch_tensor(matrix(lab, ncol = 1)))
#   }
#   tr <- gen(n); te <- gen(400)
#   net <- nn_sequential(nn_conv2d(1, 8, 3, padding = 1), nn_relu(), nn_max_pool2d(2),
#                        nn_conv2d(8, 16, 3, padding = 1), nn_relu(),
#                        nn_adaptive_avg_pool2d(1), nn_flatten(), nn_linear(16, 1))
#   opt <- optim_adamw(net$parameters, lr = 3e-3)
#   for (e in 1:60) {
#     net$train(); opt$zero_grad()
#     xb <- tr$X
#     if (aug_strength > 0) {                      # noise + random shift as augmentation
#       xb <- xb + torch_randn_like(xb) * aug_strength
#       sh <- sample(-2:2, 1)
#       if (sh != 0) xb <- torch_roll(xb, shifts = sh, dims = 4)
#     }
#     l <- nnf_binary_cross_entropy_with_logits(net(xb), tr$y)
#     l$backward(); opt$step()
#   }
#   net$eval()
#   mean(as.numeric(with_no_grad(torch_sigmoid(net(te$X)))) > 0.5 ==
#          as.numeric(te$y))
# }
# n_grid_aug <- c(30, 60, 120, 250, 500)
# a_grid <- c(0, 0.1, 0.25, 0.5, 0.8)
# Zaug <- outer(a_grid, n_grid_aug, Vectorize(synth_task))

set.seed(63); torch_manual_seed(63)
synth_task <- function(n, aug_strength, seed = 65) {
  torch_manual_seed(seed); set.seed(seed)
  gen <- function(m) {
    lab <- sample(0:1, m, TRUE)
    X <- torch_zeros(m, 1, 16, 16)
    for (i in seq_len(m)) {
      cx <- sample(4:11, 1); cy <- sample(4:11, 1)
      if (lab[i] == 1) X[i, 1, cx:(cx+3), cy:(cy+1)] <- 1
      else             X[i, 1, cx:(cx+1), cy:(cy+3)] <- 1
    }
    list(X = X + torch_randn(m, 1, 16, 16) * 0.25,
         # FIX: Force target to Float32 for BCE loss
         y = torch_tensor(matrix(as.numeric(lab), ncol = 1), dtype = torch_float32()))
  }
  tr <- gen(n); te <- gen(400)
  net <- nn_sequential(nn_conv2d(1, 8, 3, padding = 1), nn_relu(), nn_max_pool2d(2),
                       nn_conv2d(8, 16, 3, padding = 1), nn_relu(),
                       nn_adaptive_avg_pool2d(1), nn_flatten(), nn_linear(16, 1))
  opt <- optim_adamw(net$parameters, lr = 3e-3)
  for (e in 1:60) {
    net$train(); opt$zero_grad()
    xb <- tr$X
    if (aug_strength > 0) {                      # noise + random shift as augmentation
      xb <- xb + torch_randn_like(xb) * aug_strength
      sh <- sample(-2:2, 1)
      if (sh != 0) xb <- torch_roll(xb, shifts = sh, dims = 4)
    }
    l <- nnf_binary_cross_entropy_with_logits(net(xb), tr$y)
    l$backward(); opt$step()
  }
  net$eval()
  # FIX: Split evaluation to avoid console parser errors
  preds <- as.numeric(with_no_grad(torch_sigmoid(net(te$X))))
  mean((preds > 0.5) == as.numeric(te$y))
}

n_grid_aug <- c(30, 60, 120, 250, 500)
a_grid <- c(0, 0.1, 0.25, 0.5, 0.8)

# FIX: Swap outer() arguments so n_grid_aug maps to `n` and a_grid maps to `aug_strength`
Zaug <- outer(n_grid_aug, a_grid, Vectorize(synth_task))

plot_ly(x = n_grid_aug, y = a_grid, z = Zaug, type = "surface",
        colorscale = "Viridis",
        colorbar = list(title = "Test\naccuracy")) |>
  layout(title = "Test accuracy over training-set size and augmentation strength",
         scene = list(xaxis = list(title = "Training examples", type = "log"),
                      yaxis = list(title = "Augmentation strength"),
                       zaxis = list(title = "Test accuracy")))

Rotate along the sample-size axis. At small \(n\) the surface rises with augmentation strength, the invariance constraint substitutes for data it does not have. At large \(n\) it flattens and then falls: once the sample already covers the variation, further distortion only adds noise. There is no augmentation setting that is right independent of \(n\).

14.35 Transfer learning and fine-tuning

Training a modern vision model from scratch needs \(10^6\) labelled images and days of GPU time. Transfer learning reuses a network pretrained on a large corpus and adapts it, which is how nearly all applied vision is done.

\[ \begin{aligned} \textbf{Feature extraction: }&\quad \text{freeze the body, train a new head} &&\text{fast; small target sets}\\ \textbf{Fine-tuning: }&\quad \text{train some or all layers at a small LR} &&\text{better; needs more data}\\ \textbf{Discriminative LRs: }&\quad \text{smaller LR for earlier layers} &&\text{early features are more general} \end{aligned} \]

The layered structure is what makes this work. Early layers learn edges and color opponency that are generic across visual domains; later layers learn increasingly task-specific compositions. So the deeper the layer, the more it should be allowed to change.

Common misconception: “the model is pretrained, so it will transfer.” Transfer degrades with domain distance, and medical imaging is the standard cautionary case. ImageNet-pretrained features transfer well to other natural photographs and only weakly to grayscale radiographs, histology, or satellite imagery, where large pretrained models often match small models trained from scratch, at far greater cost (Raghu et al., 2019).

When it does help on medical images, the benefit is frequently attributable to better-conditioned initialization scale rather than to transferred features, the same effect Part 1, §14.9 obtains from He initialization. Test it: compare against a from-scratch baseline of matched capacity before attributing gains to the pretraining.

Common misconception: “requires_grad_(FALSE) freezes the layer.” It freezes the weights and not the buffers, and BatchNorm has both. Setting requires_grad_(FALSE) stops gradient updates to the weights. It does not stop a BatchNorm layer in train() mode from updating its running mean and variance, because those are buffers updated in the forward pass, not parameters updated by the optimizer.

The consequence is a silent distribution shift: the “frozen” body’s normalization statistics drift toward the new domain while its weights stay put, and the evaluation-mode behavior changes across epochs for reasons nothing in the training loop reveals. Freezing a body means both requires_grad_(FALSE) and putting those modules in eval() mode, and re-applying eval() after every model$train() call, since train() propagates to all children.

freeze_body <- function(module) {
  for (p in module$parameters) p$requires_grad_(FALSE)
  module$eval()                                   # BOTH steps are required
  invisible(module)
}

demo_body <- nn_sequential(nn_conv2d(3, 16, 3, padding = 1), nn_batch_norm2d(16),
                           nn_relu(), nn_adaptive_avg_pool2d(4))
x_tr <- torch_randn(16, 3, 32, 32) * 2 + 3        # a shifted "new domain"

# Weights frozen but module left in train() mode
for (p in demo_body$parameters) p$requires_grad_(FALSE)
demo_body$train()
rm_before <- as.numeric(torch_mean(demo_body[[2]]$running_mean))
invisible(demo_body(x_tr)); invisible(demo_body(x_tr))
rm_after_train <- as.numeric(torch_mean(demo_body[[2]]$running_mean))

# Properly frozen: weights AND eval mode
demo_body$eval()
invisible(demo_body(x_tr)); invisible(demo_body(x_tr))
rm_after_eval <- as.numeric(torch_mean(demo_body[[2]]$running_mean))

data.frame(
  state = c("initial", "after 2 forward passes in train() mode",
            "after 2 more in eval() mode"),
  mean_running_mean = signif(c(rm_before, rm_after_train, rm_after_eval), 5),
  changed = c(NA, rm_after_train != rm_before, rm_after_eval != rm_after_train))

The running statistics moved while every weight was frozen. That is the trap.

# How much to unfreeze depends on how much target data you have
transfer_sim <- function(n_target, n_unfrozen, seed = 71) {
  torch_manual_seed(seed); set.seed(seed)
  make <- function(m, shift = 0) {
    lab <- sample(0:1, m, TRUE)
    X <- torch_randn(m, 1, 16, 16) * 0.3 + shift
    for (i in seq_len(m))
      if (lab[i] == 1) X[i, 1, 5:11, 6:9] <- X[i, 1, 5:11, 6:9] + 1.5
      else             X[i, 1, 6:9, 5:11] <- X[i, 1, 6:9, 5:11] + 1.5
    
    # FIX 1: Force the target tensor to Float32 for BCE loss
    list(X = X, y = torch_tensor(matrix(as.numeric(lab), ncol = 1), dtype = torch_float32()))
  }
  # "Pretrain" on a large source set, then adapt to a small shifted target set
  src <- make(600, shift = 0); tgt <- make(n_target, shift = 0.8)
  te  <- make(400, shift = 0.8)
  blocks <- lapply(1:4, \(i) nn_sequential(
    nn_conv2d(if (i == 1) 1 else 8, 8, 3, padding = 1), nn_relu()))
  body <- do.call(nn_sequential, blocks)
  head <- nn_sequential(nn_adaptive_avg_pool2d(1), nn_flatten(), nn_linear(8, 1))
  net <- nn_sequential(body, head)
  opt <- optim_adamw(net$parameters, lr = 3e-3)
  for (e in 1:80) { opt$zero_grad()
    nnf_binary_cross_entropy_with_logits(net(src$X), src$y)$backward(); opt$step() }
  
  # Freeze the first (4 - n_unfrozen) blocks
  n_frozen <- 4 - n_unfrozen
  if (n_frozen > 0) for (i in seq_len(n_frozen))
    for (p in body[[i]]$parameters) p$requires_grad_(FALSE)
    
  opt2 <- optim_adamw(Filter(\(p) p$requires_grad, net$parameters), lr = 1e-3)
  for (e in 1:60) { opt2$zero_grad()
    nnf_binary_cross_entropy_with_logits(net(tgt$X), tgt$y)$backward(); opt2$step() }
  net$eval()
  mean((as.numeric(with_no_grad(torch_sigmoid(net(te$X)))) > 0.5) == as.numeric(te$y))
}

n_t <- c(20, 40, 80, 160, 320)
n_u <- 0:4

# FIX 2: Swap outer() arguments so n_t maps to n_target and n_u maps to n_unfrozen
Ztr <- outer(n_t, n_u, Vectorize(transfer_sim))

plot_ly(x = n_t, y = n_u, z = Ztr, type = "surface",
        colorscale = "Viridis", colorbar = list(title = "Target\naccuracy")) |>
  layout(title = "Transfer accuracy over target-set size and number of unfrozen blocks",
         scene = list(xaxis = list(title = "Target training examples", type = "log"),
                      yaxis = list(title = "Blocks unfrozen (0 = feature extraction)"),
                      zaxis = list(title = "Target test accuracy")))

The ridge runs diagonally. With few target examples, freezing most of the body wins, there is not enough data to fine-tune safely. As the target set grows, unfreezing more blocks wins. How much to unfreeze is a function of target-set size, not a fixed recipe.

# --- Fine-tuning a pretrained ResNet with torchvision (not evaluated) -------
library(torchvision)
model <- model_resnet18(pretrained = TRUE)

# 1. Freeze the body: BOTH requires_grad AND eval mode (see the trap above)
for (p in model$parameters) p$requires_grad_(FALSE)

# 2. Replace the head. A fresh head has random weights, so it must train first;
#    fine-tuning the body immediately would backpropagate large random-head
#    gradients into carefully pretrained features and destroy them.
n_features <- model$fc$in_features
model$fc <- nn_linear(n_features, n_classes)

# 3. Train the head only, at a normal learning rate
opt_head <- optim_adamw(model$fc$parameters, lr = 1e-3)

# 4. THEN unfreeze the last block and continue at a much smaller LR,
#    with discriminative rates: earlier layers change least.
for (p in model$layer4$parameters) p$requires_grad_(TRUE)
opt_all <- optim_adamw(list(
  list(params = model$layer4$parameters, lr = 1e-5),
  list(params = model$fc$parameters,     lr = 1e-4)))

# 5. Normalization must match what the model was pretrained with
#    (ImageNet channel statistics), or the frozen features see the wrong scale.
transform_fn <- function(img) {
  img |> transform_resize(256) |> transform_center_crop(224) |>
    transform_to_tensor() |>
    transform_normalize(mean = c(0.485, 0.456, 0.406),
                        std  = c(0.229, 0.224, 0.225))
}

Step 5 is the one most often omitted. A pretrained body expects inputs normalized with the statistics used during pretraining. Feeding raw \([0,1]\) pixels to an ImageNet-pretrained model shifts every activation and can cost more accuracy than the fine-tuning recovers.

14.36 Case study: CIFAR-10

CIFAR-10 is 60,000 \(32\times32\) color images in 10 balanced classes, small enough to train in minutes and large enough for the design choices above to matter.

# Downloads ~170 MB on first call and caches it
# # https://data.brainchip.com/dataset-mirror/cifar10/cifar-10-binary.tar.gz
# train_ds <- cifar10_dataset(root = dspa_cache_dir(), train = TRUE, download = TRUE,
#   transform = function(x) {
#     x |> transform_to_tensor() |>
#       transform_normalize(mean = c(0.4914, 0.4822, 0.4465),
#                           std  = c(0.2470, 0.2435, 0.2616))
#   })
# test_ds <- cifar10_dataset(root = dspa_cache_dir(), train = FALSE, download = TRUE,
#   transform = function(x) {
#     x |> transform_to_tensor() |>
#       transform_normalize(mean = c(0.4914, 0.4822, 0.4465),
#                           std  = c(0.2470, 0.2435, 0.2616))
#   })
# 
# # THREE-way split: the 10,000-image test set is touched once, at the end
# set.seed(1234)
# idx <- sample(length(train_ds))
# train_idx <- idx[1:45000]; valid_idx <- idx[45001:50000]
# 
# train_dl <- dataloader(dataset_subset(train_ds, train_idx), batch_size = 128,
#                        shuffle = TRUE)
# valid_dl <- dataloader(dataset_subset(train_ds, valid_idx), batch_size = 256)
# test_dl  <- dataloader(test_ds, batch_size = 256)
# c(train = length(train_idx), validation = length(valid_idx), test = length(test_ds))
cache_dir <- dspa_cache_dir()
tar_path <- file.path(cache_dir, "cifar-10-binary.tar.gz")
extracted_dir <- file.path(cache_dir, "cifar-10-batches-bin")

# Download from BrainChip mirror and extract if not already cached
if (!dir.exists(extracted_dir)) {
  dir.create(cache_dir, showWarnings = FALSE, recursive = TRUE)
  message("Downloading CIFAR-10 from BrainChip mirror...")
  download.file(
    url = "https://data.brainchip.com/dataset-mirror/cifar10/cifar-10-binary.tar.gz",
    destfile = tar_path,
    mode = "wb"
  )
  message("Extracting archive...")
  untar(tar_path, exdir = cache_dir)
}

# Load datasets using the local cache with download = FALSE
train_ds <- cifar10_dataset(root = cache_dir, train = TRUE, download = FALSE,
  transform = function(x) {
    x |> transform_to_tensor() |>
      transform_normalize(mean = c(0.4914, 0.4822, 0.4465),
                          std  = c(0.2470, 0.2435, 0.2616))
  })

test_ds <- cifar10_dataset(root = cache_dir, train = FALSE, download = FALSE,
  transform = function(x) {
    x |> transform_to_tensor() |>
      transform_normalize(mean = c(0.4914, 0.4822, 0.4465),
                          std  = c(0.2470, 0.2435, 0.2616))
  })

# THREE-way split: the 10,000-image test set is touched once, at the end
set.seed(1234)
idx <- sample(length(train_ds))
train_idx <- idx[1:45000]; valid_idx <- idx[45001:50000]

train_dl <- dataloader(dataset_subset(train_ds, train_idx), batch_size = 128,
                       shuffle = TRUE)
valid_dl <- dataloader(dataset_subset(train_ds, valid_idx), batch_size = 256)
test_dl  <- dataloader(test_ds, batch_size = 256)
c(train = length(train_idx), validation = length(valid_idx), test = length(test_ds))
# A small residual network: every element from Sections 14.25-14.33
res_block_s <- nn_module(
  "ResBlockS",
  initialize = function(c_in, c_out, stride = 1) {
    self$body <- nn_sequential(
      nn_conv2d(c_in, c_out, 3, stride = stride, padding = 1, bias = FALSE),
      nn_batch_norm2d(c_out), nn_relu(),
      nn_conv2d(c_out, c_out, 3, padding = 1, bias = FALSE),
      nn_batch_norm2d(c_out))
    # A projection is needed when the shortcut's shape does not match
    self$shortcut <- if (stride != 1 || c_in != c_out)
      nn_sequential(nn_conv2d(c_in, c_out, 1, stride = stride, bias = FALSE),
                    nn_batch_norm2d(c_out)) else nn_identity()
    self$act <- nn_relu()
  },
  forward = function(x) self$act(self$body(x) + self$shortcut(x)))

small_resnet <- nn_module(
  "SmallResNet",
  initialize = function(n_classes = 10, width = 32) {
    self$stem <- nn_sequential(nn_conv2d(3, width, 3, padding = 1, bias = FALSE),
                               nn_batch_norm2d(width), nn_relu())
    self$stage1 <- nn_sequential(res_block_s(width, width),
                                 res_block_s(width, width))
    self$stage2 <- nn_sequential(res_block_s(width, 2*width, stride = 2),
                                 res_block_s(2*width, 2*width))
    self$stage3 <- nn_sequential(res_block_s(2*width, 4*width, stride = 2),
                                 res_block_s(4*width, 4*width))
    # Global average pooling, not a dense head (Section 14.28)
    self$head <- nn_sequential(nn_adaptive_avg_pool2d(1), nn_flatten(),
                               nn_linear(4*width, n_classes))
  },
  forward = function(x)
    self$head(self$stage3(self$stage2(self$stage1(self$stem(x))))))

m_cifar <- small_resnet()
c(parameters = format(n_par(m_cifar), big.mark = ","),
  output_shape = paste(dim(with_no_grad(m_cifar(torch_randn(2,3,32,32)))),
                       collapse = " x "))
#>   parameters output_shape 
#>    "696,618"     "2 x 10"
# Augmentation is applied to TRAINING only -- never to validation or test
augment <- function(batch) {
  if (runif(1) > 0.5) batch <- transform_hflip(batch)
  pad <- nnf_pad(batch, c(4,4,4,4), mode = "reflect")     # reflect, not zeros
  ox <- sample(0:8, 1); oy <- sample(0:8, 1)
  pad[, , (oy+1):(oy+32), (ox+1):(ox+32)]
}

fit_cifar <- function(epochs = 30, lr = 0.1, use_aug = TRUE) {
  torch_manual_seed(1234)
  net <- small_resnet()
  opt <- optim_sgd(net$parameters, lr = lr, momentum = 0.9, weight_decay = 5e-4)
  sched <- lr_one_cycle(opt, max_lr = lr, epochs = epochs,
                        steps_per_epoch = length(train_dl))
  hist <- data.frame()
  for (e in seq_len(epochs)) {
    net$train(); tot <- 0; n <- 0
    coro::loop(for (b in train_dl) {
      xb <- if (use_aug) augment(b[[1]]) else b[[1]]
      opt$zero_grad()
      l <- nnf_cross_entropy(net(xb), b[[2]])
      l$backward(); opt$step(); sched$step()
      tot <- tot + as.numeric(l) * b[[2]]$size(1); n <- n + b[[2]]$size(1)
    })
    net$eval(); correct <- 0; m <- 0
    with_no_grad(coro::loop(for (b in valid_dl) {
      correct <- correct + as.numeric(sum(net(b[[1]])$argmax(dim = 2) == b[[2]]))
      m <- m + b[[2]]$size(1)
    }))
    hist <- rbind(hist, data.frame(epoch = e, train_loss = tot/n,
                                   valid_acc = correct/m))
  }
  list(net = net, history = hist)
}
res_aug   <- dspa_cache("cifar_aug",   fit_cifar(use_aug = TRUE))
res_noaug <- dspa_cache("cifar_noaug", fit_cifar(use_aug = FALSE))
bind_rows(mutate(res_aug$history,   run = "with augmentation"),
          mutate(res_noaug$history, run = "no augmentation")) |>
  ggplot(aes(epoch, valid_acc, color = run)) +
  geom_line(linewidth = 1) +
  scale_color_manual(values = c(`with augmentation` = "#3B7DD8",
                                 `no augmentation` = "#D8433B")) +
  labs(title = "CIFAR-10 validation accuracy",
       subtitle = "Same architecture, seed, and schedule. Only the augmentation differs",
       x = "Epoch", y = "Validation accuracy", color = NULL) +
  theme_dspa()

# The test set is scored ONCE, after all decisions are final
score_test <- function(net) {
  net$eval(); correct <- 0; m <- 0
  with_no_grad(coro::loop(for (b in test_dl) {
    correct <- correct + as.numeric(sum(net(b[[1]])$argmax(dim = 2) == b[[2]]))
    m <- m + b[[2]]$size(1) }))
  correct / m
}
data.frame(run = c("with augmentation", "no augmentation"),
           best_validation = c(max(res_aug$history$valid_acc),
                               max(res_noaug$history$valid_acc)),
           test_accuracy = c(score_test(res_aug$net), score_test(res_noaug$net)))

Report the class-wise confusion, not only the aggregate. CIFAR-10 is balanced, so aggregate accuracy is not misleading about prevalence, but it still hides which classes fail. Cat–dog and automobile–truck confusions dominate the error, and an aggregate figure conceals that a deployment distinguishing exactly those pairs would perform far worse than the headline number suggests. The metric family of Chapter 9, §9.4 applies unchanged.


PART IV: SEGMENTATION AND INTERPRETATION

14.37 Segmentation and the Dice coefficient

Classification assigns one label per image. Segmentation assigns one per pixel, the task behind tumour delineation, organ volumetry, and lesion quantification.

The reference architecture is the U-Net (Ronneberger et al., 2015): a contracting encoder that builds context, an expanding decoder that restores resolution, and skip connections carrying high-resolution features across at matching scales.

The skip connections are not an optimization device here. In a ResNet the additive shortcut exists to let gradients flow (§14.32). In a U-Net the skips are concatenative and exist to recover spatial precision: the encoder discards location as it downsamples, and only the corresponding encoder feature map still knows where the boundary was. Without them the decoder can say what is present and not exactly where, which is the entire quantity segmentation is scored on.

conv_block <- nn_module("ConvBlock",
  initialize = function(c_in, c_out) {
    self$net <- nn_sequential(
      nn_conv2d(c_in, c_out, 3, padding = 1, bias = FALSE),
      nn_group_norm(min(8, c_out), c_out), nn_relu(),   # GroupNorm: small batches
      nn_conv2d(c_out, c_out, 3, padding = 1, bias = FALSE),
      nn_group_norm(min(8, c_out), c_out), nn_relu())
  },
  forward = function(x) self$net(x))

unet_small <- nn_module("UNetSmall",
  initialize = function(c_in = 1, base = 16) {
    self$enc1 <- conv_block(c_in, base)
    self$enc2 <- conv_block(base, base*2)
    self$enc3 <- conv_block(base*2, base*4)
    self$pool <- nn_max_pool2d(2)
    self$up2  <- nn_conv_transpose2d(base*4, base*2, 2, stride = 2)
    self$dec2 <- conv_block(base*4, base*2)      # base*4 = upsampled + skip
    self$up1  <- nn_conv_transpose2d(base*2, base, 2, stride = 2)
    self$dec1 <- conv_block(base*2, base)
    self$out  <- nn_conv2d(base, 1, 1)           # 1x1 conv -> one logit per pixel
  },
  forward = function(x) {
    e1 <- self$enc1(x)
    e2 <- self$enc2(self$pool(e1))
    e3 <- self$enc3(self$pool(e2))
    d2 <- self$dec2(torch_cat(list(self$up2(e3), e2), dim = 2))   # concatenate
    d1 <- self$dec1(torch_cat(list(self$up1(d2), e1), dim = 2))
    self$out(d1)                                  # LOGITS, not probabilities
  })

m_unet <- unet_small()
c(parameters = format(n_par(m_unet), big.mark = ","),
  input = "1 x 64 x 64",
  output = paste(dim(with_no_grad(m_unet(torch_randn(2,1,64,64))))[-1],
                 collapse = " x "))
#>    parameters         input        output 
#>     "117,073" "1 x 64 x 64" "1 x 64 x 64"

14.37.1 The Dice coefficient

\[\boxed{\;D=\frac{2|X\cap Y|}{|X|+|Y|}=\frac{2\,TP}{2\,TP+FP+FN}\;}\]

Dice is the \(F_1\) score. Writing \(F_1=\frac{2\cdot\text{precision}\cdot\text{recall}}{\text{precision}+\text{recall}}\) and substituting \(\text{precision}=\frac{TP}{TP+FP}\), \(\text{recall}=\frac{TP}{TP+FN}\) gives exactly \(\frac{2TP}{2TP+FP+FN}\). The segmentation literature and the classification literature named the same quantity twice.

That identity carries its properties over from Chapter 9, §9.4. Dice ignores true negatives entirely, which is why it is used instead of accuracy for segmentation, where a tumour may occupy 1% of the image and predicting “background everywhere” scores 99% accurate and 0 Dice.

Two facts about Dice that determine how it must be used:

It is undefined on empty masks. If the ground truth has no tumour and the prediction correctly has none, \(D=\frac{0}{0}\). The convention is to define \(D=1\) for that case, but a naive implementation returns NaN and averaging over slices silently drops them, inflating the reported mean by removing the easiest cases. Empty-mask slices are common in volumetric MRI, where most slices contain no lesion.

The hard form is not differentiable, so training uses the soft Dice with predicted probabilities \(p\) in place of a thresholded mask, plus a smoothing constant:

\[\mathcal L_{\text{Dice}}=1-\frac{2\sum_i p_iy_i+\epsilon}{\sum_i p_i+\sum_i y_i+\epsilon}\]

dice_hard <- function(pred, target, empty_value = 1) {
  tp <- sum(pred == 1 & target == 1)
  denom <- sum(pred == 1) + sum(target == 1)
  if (denom == 0) return(empty_value)        # BOTH empty: perfect agreement
  2 * tp / denom
}
dice_soft_loss <- function(logits, target, eps = 1) {
  p <- torch_sigmoid(logits)
  num <- 2 * torch_sum(p * target) + eps
  den <- torch_sum(p) + torch_sum(target) + eps
  1 - num / den
}

# The empty-mask case, handled three ways
empty_pred <- rep(0, 100); empty_true <- rep(0, 100)
naive_dice <- 2*sum(empty_pred == 1 & empty_true == 1) /
              (sum(empty_pred) + sum(empty_true))
data.frame(
  handling = c("naive 2TP/(|X|+|Y|)", "convention D = 1", "smoothed, eps = 1"),
  value = c(naive_dice, dice_hard(empty_pred, empty_true),
            2*0/(0 + 0 + 1) + 1/(0+0+1)),
  note = c("NaN: silently dropped by mean(), inflating the average",
           "correct: two empty sets agree perfectly",
           "smoothing makes it finite and differentiable"))
# Why accuracy is useless for segmentation, and Dice is not
set.seed(81)
lesion_fracs <- c(0.5, 0.2, 0.05, 0.01, 0.002)
cmp <- do.call(rbind, lapply(lesion_fracs, function(f) {
  n <- 10000
  truth <- rbinom(n, 1, f)
  all_bg <- rep(0, n)                               # predict background everywhere
  data.frame(lesion_fraction = f,
             accuracy_of_trivial = mean(all_bg == truth),
             dice_of_trivial = dice_hard(all_bg, truth, empty_value = 0))
}))
cmp |> mutate(across(-lesion_fraction, \(z) round(z, 4)))
cmp |> pivot_longer(-lesion_fraction, names_to = "metric", values_to = "v") |>
  ggplot(aes(lesion_fraction, v, color = metric)) +
  geom_line(linewidth = 1) + geom_point(size = 2.4) +
  scale_x_log10() +
  scale_color_manual(values = c(accuracy_of_trivial = "#D8433B",
                                 dice_of_trivial = "#3B7DD8"),
                      labels = c("Accuracy", "Dice")) +
  labs(title = "A model that predicts 'no lesion' everywhere",
       subtitle = "Accuracy rises toward 1 as lesions become rarer; Dice correctly stays at 0",
       x = "Lesion fraction of the image (log scale)", y = "Score", color = NULL) +
  theme_dspa()

14.37.2 Protocol for medical segmentation

Split by patient, never by slice. A volumetric scan yields dozens of adjacent 2-D slices that are nearly identical. Splitting at the slice level puts slice \(k\) in training and slice \(k+1\) in test, so the model is scored on images it has effectively seen, the leakage of Chapter 9, §9.17.1 and Chapter 12, §12.17, in its most consequential form.

Reported Dice under slice-level splitting can exceed patient-level Dice by a wide margin, and the difference is entirely artifact. The grouping variable is the patient, and where scans come from several sites, splitting by site as well tests the generalization that actually matters clinically.

Weighted sampling toward large lesions decalibrates the model. Oversampling slices with big tumours is a reasonable response to class imbalance, and it changes the effective prior. The resulting predicted probabilities no longer correspond to frequencies in the real population, the model is trained on a distribution where lesions are common and deployed where they are rare.

Two consequences follow. Any probability threshold tuned on the resampled validation set is wrong for deployment. And the calibration assessment of Chapter 9, §9.8 must be performed on a validation set with the natural prevalence, not the resampled one. Recalibrating afterwards, Platt scaling on a natural-prevalence held-out set, restores the correspondence.

# A small synthetic segmentation task, trained end to end
set.seed(83); torch_manual_seed(83)
make_seg <- function(n, size = 64) {
  X <- torch_randn(n, 1, size, size) * 0.35
  Y <- torch_zeros(n, 1, size, size)
  for (i in seq_len(n)) {
    r <- sample(6:12, 1); cx <- sample((r+2):(size-r-2), 1); cy <- sample((r+2):(size-r-2), 1)
    gx <- matrix(rep(1:size, size), size); gy <- t(gx)
    mask <- ((gx - cx)^2 + (gy - cy)^2) <= r^2
    X[i,1,,] <- X[i,1,,] + torch_tensor(mask * 1.6)
    Y[i,1,,] <- torch_tensor(mask * 1)
  }
  list(X = X, Y = Y)
}
tr_seg <- make_seg(120); va_seg <- make_seg(40)

net_seg <- unet_small(base = 8)
opt_seg <- optim_adamw(net_seg$parameters, lr = 3e-3)
for (e in 1:60) {
  net_seg$train(); opt_seg$zero_grad()
  logit <- net_seg(tr_seg$X)
  # Combined loss: BCE stabilizes early training, Dice targets the metric
  l <- nnf_binary_cross_entropy_with_logits(logit, tr_seg$Y) +
       dice_soft_loss(logit, tr_seg$Y)
  l$backward(); opt_seg$step()
}
net_seg$eval()
pred_va <- with_no_grad(torch_sigmoid(net_seg(va_seg$X)))
d_scores <- vapply(1:40, \(i) dice_hard(
  as.array(pred_va[i,1,,]) > 0.5, as.array(va_seg$Y[i,1,,]) == 1), numeric(1))
c(mean_dice = round(mean(d_scores), 4), sd_dice = round(sd(d_scores), 4),
  min_dice = round(min(d_scores), 4))
#> mean_dice   sd_dice  min_dice 
#>    0.9750    0.0234    0.9187
k <- which.min(abs(d_scores - median(d_scores)))
to_grid <- function(m, lab) {
  a <- as.array(m)
  expand.grid(row = 1:nrow(a), col = 1:ncol(a)) |>
    mutate(v = as.vector(a), panel = lab)
}
bind_rows(to_grid(va_seg$X[k,1,,], "Input"),
          to_grid(va_seg$Y[k,1,,], "Ground-truth mask"),
          to_grid(pred_va[k,1,,], "Predicted probability")) |>
  mutate(panel = factor(panel, levels = c("Input", "Ground-truth mask",
                                          "Predicted probability"))) |>
  ggplot(aes(col, -row, fill = v)) +
  geom_raster() + facet_wrap(~ panel) + coord_fixed() +
  scale_fill_viridis_c(guide = "none") +
  labs(title = sprintf("Median-Dice validation case (D = %.3f)", d_scores[k]),
       subtitle = "The output is a probability per pixel; the threshold is a separate decision",
       x = NULL, y = NULL) +
  theme_dspa(9) + theme(axis.text = element_blank(), panel.grid = element_blank())

14.38 Interpretability

Three families of post-hoc attribution, in increasing sophistication.

\[ \begin{aligned} \textbf{Vanilla saliency: }&\quad S(x)=\Big|\tfrac{\partial f_c(x)}{\partial x}\Big| &&\text{one backward pass}\\ \textbf{Grad-CAM: }&\quad L^c=\mathrm{ReLU}\Big(\textstyle\sum_k\alpha_k^cA^k\Big),\ \ \alpha_k^c=\tfrac{1}{Z}\textstyle\sum_{u,v}\tfrac{\partial f_c}{\partial A^k_{u,v}} &&\text{class-discriminative}\\ \textbf{Integrated gradients: }&\quad IG_i=(x_i-x_i')\!\int_0^1\!\tfrac{\partial f(x'+\alpha(x-x'))}{\partial x_i}d\alpha &&\text{satisfies completeness} \end{aligned} \]

Grad-CAM weights the final convolutional feature maps by their average gradient, giving a coarse but class-specific heat map. Integrated gradients accumulates gradients along a path from a baseline \(x'\) to \(x\) and satisfies completeness, the attributions sum to \(f(x)-f(x')\), which vanilla saliency does not.

Common misconception: “the saliency map shows what the model looked at.” Attribution maps are seductive because they produce a picture that looks like an explanation, and several published failure modes should temper that.

Some are independent of the model. Certain saliency methods produce essentially unchanged maps when the network’s weights are randomized, or when the labels are randomized, so they are acting as edge detectors on the input rather than reporting anything the model learned (Adebayo et al., 2018). Run those sanity checks before trusting a map.

They are not stable. Imperceptible input perturbations that leave the prediction unchanged can produce a completely different attribution (Ghorbani et al., 2019).

They show association, not mechanism. A bright region indicates that the output is locally sensitive to those pixels. It does not establish that the model uses the clinically meaningful feature there, the well-documented cases where a model keyed on scanner text, laterality markers, or a chest-drain artifact rather than pathology all produced plausible-looking heat maps over the lesion’s vicinity.

Use attributions to generate hypotheses about failure, then test those hypotheses by intervention: occlude the region, retrain without the artifact, or evaluate on a site where the artifact is absent.

torch_manual_seed(91)
# A small classifier on the synthetic segmentation data: is the circle large?
cls_net <- nn_module("Cls",
  initialize = function() {
    self$features <- nn_sequential(
      nn_conv2d(1, 8, 3, padding = 1), nn_relu(), nn_max_pool2d(2),
      nn_conv2d(8, 16, 3, padding = 1), nn_relu(), nn_max_pool2d(2),
      nn_conv2d(16, 16, 3, padding = 1), nn_relu())
    self$head <- nn_sequential(nn_adaptive_avg_pool2d(1), nn_flatten(),
                               nn_linear(16, 1))
  },
  forward = function(x) self$head(self$features(x)))()

y_big <- torch_tensor(matrix(as.numeric(
  vapply(1:120, \(i) sum(as.array(tr_seg$Y[i,1,,])), numeric(1)) > 250), ncol = 1))
opt_c <- optim_adamw(cls_net$parameters, lr = 3e-3)
for (e in 1:120) {
  opt_c$zero_grad()
  nnf_binary_cross_entropy_with_logits(cls_net(tr_seg$X), y_big)$backward()
  opt_c$step()
}

grad_cam <- function(net, x) {
  A <- net$features(x); A$retain_grad()
  score <- net$head(A)$squeeze()
  score$backward()
  alpha <- torch_mean(A$grad, dim = c(3, 4), keepdim = TRUE)   # spatial mean
  cam <- torch_relu(torch_sum(alpha * A, dim = 2, keepdim = TRUE))
  cam <- nnf_interpolate(cam, size = c(dim(x)[3], dim(x)[4]), mode = "bilinear",
                         align_corners = FALSE)
  as.array(cam$squeeze())
}
saliency <- function(net, x) {
  x <- x$clone()$detach()$requires_grad_(TRUE)
  net(x)$squeeze()$backward()
  abs(as.array(x$grad$squeeze()))
}

xi <- tr_seg$X[3, , , , drop = FALSE]
cam_map <- grad_cam(cls_net, xi)
sal_map <- saliency(cls_net, xi)

# SANITY CHECK: randomize the weights and recompute
torch_manual_seed(93)
rand_net <- nn_module("Cls2", initialize = cls_net$initialize,
                      forward = cls_net$forward)()
cam_rand <- grad_cam(rand_net, xi)

nrm <- function(m) (m - min(m)) / (max(m) - min(m) + 1e-12)
bind_rows(to_grid(torch_tensor(as.array(xi[1,1,,])), "Input"),
          to_grid(torch_tensor(nrm(sal_map)), "Vanilla saliency"),
          to_grid(torch_tensor(nrm(cam_map)), "Grad-CAM (trained)"),
          to_grid(torch_tensor(nrm(cam_rand)), "Grad-CAM (RANDOM weights)")) |>
  mutate(panel = factor(panel, levels = c("Input", "Vanilla saliency",
                                          "Grad-CAM (trained)",
                                          "Grad-CAM (RANDOM weights)"))) |>
  ggplot(aes(col, -row, fill = v)) +
  geom_raster() + facet_wrap(~ panel, nrow = 1) + coord_fixed() +
  scale_fill_viridis_c(option = "inferno", guide = "none") +
  labs(title = "Attribution maps, with the model-randomization sanity check",
       subtitle = "If the random-weight panel resembles the trained one, the method is reporting the input, not the model",
       x = NULL, y = NULL) +
  theme_dspa(9) + theme(axis.text = element_blank(), panel.grid = element_blank())

c(correlation_trained_vs_random_cam =
    signif(cor(as.vector(cam_map), as.vector(cam_rand)), 3),
  interpretation = "high correlation would mean the map is not model-dependent")
#>                            correlation_trained_vs_random_cam 
#>                                                      "0.277" 
#>                                               interpretation 
#> "high correlation would mean the map is not model-dependent"
# --- Interactive equivalent ------------------------------------------------
plot_ly(z = nrm(cam_map), type = "heatmap", colorscale = "Inferno") |>
  layout(title = "Grad-CAM attribution", xaxis = list(scaleanchor = "y"))

PART V: SYNTHESIS

14.39 Complexity summary

\(B\) batch, \(C\) channels, \(H\times W\) spatial size, \(k\) kernel, \(L\) depth, \(G\) groups.

Operation Parameters FLOPs Activation memory Note
Dense \(d_{\text{in}}\!\to\!d_{\text{out}}\) \(d_{\text{in}}d_{\text{out}}\) \(2Bd_{\text{in}}d_{\text{out}}\) \(4Bd_{\text{out}}\) \(I\approx0.5\): memory-bound
Conv \(k\times k\) \(C_{\text{out}}(C_{\text{in}}k^2{+}1)\) \(2BC_{\text{out}}H'W'C_{\text{in}}k^2\) \(4BC_{\text{out}}H'W'\) \(I\gg1\): compute-bound
Conv \(1\times1\) \(C_{\text{out}}(C_{\text{in}}{+}1)\) \(2BC_{\text{out}}HWC_{\text{in}}\) \(4BC_{\text{out}}HW\) Channel mixing only
Bottleneck \(1\!\times\!1\)\(3\!\times\!3\)\(1\!\times\!1\) \(\approx\!\tfrac14\) plain \(\approx\!\tfrac14\) plain similar ResNet-50+
Depthwise separable \(C_{\text{in}}k^2+C_{\text{in}}C_{\text{out}}\) \(\approx\!\tfrac{1}{k^2}\) of dense conv similar MobileNet
Max / avg pool \(2\times2\) \(0\) \(O(BCHW)\) \(4BCHW/4\) No learning
Global average pool \(0\) \(O(BCHW)\) \(4BC\) Removes the dense head
BatchNorm \(2C\) \(O(BCHW)\) \(4BCHW\) Batch-coupled
GroupNorm \(2C\) \(O(BCHW)\) \(4BCHW\) Batch-independent
Residual add \(0\) \(O(BCHW)\) Gradient path with derivative 1
U-Net skip (concat) \(0\) doubles the decoder input Restores spatial precision
Receptive field (stride 1) \(r_L=1+\sum(k_\ell-1)\): linear
Receptive field (stride \(s\)) geometric in depth
Effective receptive field \(O(\sqrt L)\), not \(O(L)\)

Four consequences.

Weight reuse is what makes CNNs compute-bound. A convolution applies each weight \(H'W'\) times, pushing arithmetic intensity above the roofline ridge (Chapter 10, §10.14.1), which is why GPUs deliver one to two orders of magnitude on CNNs and only the bandwidth ratio on MLPs.

Activation memory dominates in vision. A \(3\times3\) convolution on a \(224\times224\) input with 64 channels has 37 K parameters and 3.2 M activations per image. Early layers of a CNN are where the memory goes, which is why downsampling early is standard.

Downsampling buys receptive field geometrically. Reaching a receptive field of 200 pixels takes 100 stride-1 layers or about 7 stride-2 layers.

The effective receptive field grows as \(O(\sqrt L)\), so the theoretical figure overstates what the network uses, the gap that dilation and attention exist to close.


14.40 Common pitfalls

# Pitfall Consequence Fix
1 Treating convolution as designed filtering Misses that kernels are learned jointly with the head It is a structural prior, not preprocessing
2 Assuming “same” padding preserves information Zero-padding invents border data; frame artifacts reflect/replicate, or valid convolution
3 Reading the theoretical receptive field as what the net uses Effective field is \(O(\sqrt L)\), roughly Gaussian Measure it by input gradient
4 Calling CNNs translation-invariant They are equivariant; invariance comes from pooling And strided layers alias, breaking even that
5 Treating pooling as mandatory It has no parameters and cannot adapt Strided convolution; global average pooling
6 Keeping a VGG-style dense head ~90% of parameters for little gain Global average pooling
7 Reading a deeper net’s worse training loss as overfitting It is the degradation problem — an optimization failure Residual connections
8 Using BatchNorm at batch size 2–8 Statistics too noisy to normalize with GroupNorm or LayerNorm
9 Forgetting BatchNorm’s train()/eval() modes Running statistics silently wrong Switch modes explicitly
10 “Freezing” a body without eval() Running statistics still drift requires_grad_(FALSE) and eval()
11 Fine-tuning immediately with a random head Large random gradients destroy pretrained features Train the head first, then unfreeze
12 Omitting the pretraining normalization statistics Frozen features see the wrong input scale Match the pretrained mean/std
13 Assuming pretrained features transfer to any domain Weak on medical, satellite, microscopy Compare against from-scratch of matched capacity
14 Copying an ImageNet augmentation pipeline Each transform asserts an invariance Verify each is label-preserving in your domain
15 Augmenting the validation or test set The estimate no longer describes deployment Augment training only
16 Believing augmentation adds information It adds examples, not information It is a constraint on the hypothesis space
17 Splitting a volumetric dataset by slice Adjacent slices are near-duplicates; severe leakage Split by patient, and by site where relevant
18 Reporting accuracy for segmentation 99% accurate at 1% lesion fraction, 0 Dice Dice / IoU, which ignore true negatives
19 Naive Dice on empty masks NaN, dropped by mean(), inflating the average Define \(D=1\) when both are empty; smooth
20 Thresholding soft Dice during training Not differentiable Soft Dice on probabilities with \(\epsilon\)
21 Weighted sampling toward large lesions, then reporting probabilities Decalibrated against natural prevalence Recalibrate on a natural-prevalence set
22 Treating saliency maps as explanations Some are unchanged by weight randomization Run the sanity checks; test by intervention
23 Reporting only aggregate accuracy Hides which class pairs fail Per-class confusion
24 Tuning on the test set across architecture trials The test estimate becomes selection-optimized Three-way split; test touched once

14.41 Practice problems

Problem 1: Verify the output-size formula against torch

Solution
grid_cfg <- expand.grid(k = c(1,3,5,7), p = 0:3, s = 1:3, d = 1:2)
grid_cfg$formula <- with(grid_cfg, floor((64 + 2*p - d*(k-1) - 1)/s) + 1)
grid_cfg$torch <- mapply(function(k,p,s,d) {
  dim(nnf_conv2d(torch_randn(1,1,64,64), torch_randn(1,1,k,k),
                 padding = p, stride = s, dilation = d))[3]
}, grid_cfg$k, grid_cfg$p, grid_cfg$s, grid_cfg$d)
c(configurations = nrow(grid_cfg),
  all_match = all(grid_cfg$formula == grid_cfg$torch),
  max_discrepancy = max(abs(grid_cfg$formula - grid_cfg$torch)))
#>  configurations       all_match max_discrepancy 
#>              96               1               0
subset(grid_cfg, formula == 64)[1:6, ]
All configurations agree. The rows with output 64 are the size-preserving ones: \(p=d(k-1)/2\) with \(s=1\).

Problem 2: Measure the effective receptive field’s growth rate

Solution
# ----------------------------------------------------------------------
# Reproducibility
# ----------------------------------------------------------------------
torch_manual_seed(1234)

# ----------------------------------------------------------------------
# Network without a final ReLU
# ----------------------------------------------------------------------
# ----------------------------------------------------------------------
# CNN for ERF experiment
#
# Key changes:
#   1. several hidden channels instead of only one
#   2. LeakyReLU instead of ReLU
#   3. no activation after final convolution
# ----------------------------------------------------------------------
stack_n <- function(n_layers, width = 8L, negative_slope = 0.01) {
  stopifnot(n_layers >= 1)
  layers <- list()
  for (i in seq_len(n_layers)) {
    in_channels <-
      if (i == 1L) 1L else width
    out_channels <-
      if (i == n_layers) 1L else width
    layers[[length(layers) + 1L]] <-
      nn_conv2d(in_channels  = in_channels, out_channels = out_channels,
        kernel_size  = 3, padding = 1)

    # Activation only between convolutions.
    # LeakyReLU is intentional here:
    # it avoids exactly zero gradient paths.
    if (i < n_layers) {
      layers[[length(layers) + 1L]] <- nn_leaky_relu(negative_slope = negative_slope)
    }
  }
  nn_sequential(!!!layers)
}


# ----------------------------------------------------------------------
# Gradient map
# Returns an ordinary base-R numeric matrix.
# ----------------------------------------------------------------------
erf_map <- function(n_layers, size = 81L, width = 8L) {
  net <- stack_n(n_layers = n_layers, width = width)

  x <- torch_randn(1, 1, size, size,  requires_grad = TRUE  )

  out <- net(x)

  ctr <- (size + 1L) %/% 2L

  centre <- out[1, 1,  ctr, ctr]

  centre$backward()

  if (is.null(x$grad)) {
    stop("Autograd did not generate x$grad.")
  }

  # ------------------------------------------------------------
  # Torch -> ordinary R conversion
  # ------------------------------------------------------------
  grad <- x$grad$
    detach()$
    abs()$
    squeeze()$
    cpu()

  m <- torch::as_array(grad)

  # Explicitly force base-R numeric matrix storage
  m <- matrix(as.numeric(m), nrow = size, ncol = size)

  m
}

# ----------------------------------------------------------------------
# Monte-Carlo estimate at each depth
# ----------------------------------------------------------------------
erf_radius <- function(n_layers, size = 81L, width = 8L, n_rep = 10L, seed = 1234L) {

  radii <- numeric(n_rep)

  for (r in seq_len(n_rep)) {

    # Reproducible but different realization
    torch_manual_seed(
      seed +
        1000L * as.integer(n_layers) +
        as.integer(r)
    )

    radii[r] <- erf_radius_once(
      n_layers = n_layers,
      size = size,
      width = width
    )
  }

  mean(radii)
}

# ----------------------------------------------------------------------
# Effective receptive-field RMS radius
# ----------------------------------------------------------------------
erf_radius_once <- function(n_layers, size = 81L, width = 8L) {

  m <- erf_map(n_layers = n_layers, size = size, width = width)

  stopifnot(is.matrix(m), is.numeric(m), all(dim(m) == c(size, size)))

  if (any(!is.finite(m))) {
    stop("Non-finite values detected at depth ", n_layers)
  }

  total_mass <- sum(m)

  if (!is.finite(total_mass) || total_mass <= 0) {
    stop("Zero gradient mass at depth ", n_layers,
      ". This should be extremely unlikely with LeakyReLU."
    )
  }

  # Normalize gradient mass
  m <- m / total_mass

  ctr <- (size + 1L) %/% 2L

  coords <- seq_len(size)

  gx <- matrix(rep(coords, size), nrow = size, ncol = size)

  gy <- t(gx)

  r2 <- (gx - ctr)^2 + (gy - ctr)^2

  sqrt(sum(m * r2))
}


# ----------------------------------------------------------------------
# Compare effective and theoretical radii
# ----------------------------------------------------------------------
depths_e <- c(2, 4, 8, 12, 16, 20)

p2 <- data.frame(depth = depths_e, 
                 effective_radius = vapply(depths_e, erf_radius, numeric(1)),
                 theoretical_radius = vapply(depths_e, function(n) {
                   (rf(rep(3, n), rep(1, n)) - 1) / 2
                   },
                   numeric(1))
                 )

p2$ratio <-  p2$effective_radius /  p2$theoretical_radius

p2 |>
  mutate(
    across(
      where(is.numeric),
      function(z) round(z, 3)
    )
  )
# ----------------------------------------------------------------------
# Power-law fit
# ----------------------------------------------------------------------
fit_pow <- lm(
  log(effective_radius) ~ log(depth),
  data = p2
)

beta <- unname(coef(fit_pow)[["log(depth)"]])

c(
  fitted_exponent = round(beta, 3),
  predicted_by_theory = 0.5
)
#>     fitted_exponent predicted_by_theory 
#>               0.593               0.500
# ----------------------------------------------------------------------
# Plot
# ----------------------------------------------------------------------
ggplot(p2, aes(x = depth)) +

  geom_line(
    aes(
      y = theoretical_radius,
      color = "Theoretical"
    ),
    linewidth = 1
  ) +

  geom_line(
    aes(
      y = effective_radius,
      color = "Effective (measured)"
    ),
    linewidth = 1
  ) +

  geom_point(
    aes(
      y = effective_radius,
      color = "Effective (measured)"
    ),
    size = 2.4
  ) +

  scale_x_log10(
    breaks = depths_e
  ) +

  scale_y_log10() +

  scale_color_manual(
    values = c(
      Theoretical = "#D8433B",
      `Effective (measured)` = "#3B7DD8"
    )
  ) +

  labs(
    title =
      "Theoretical and effective receptive-field radius",

    subtitle =
      paste(
        "On log-log axes, the theoretical slope is 1;",
        "the effective receptive field is expected to approach 1/2."
      ),

    x = "Depth (log scale)",
    y = "Radius in pixels (log scale)",
    color = NULL
  ) +

  theme_dspa()

The fitted exponent is close to \(0.5\), confirming \(O(\sqrt L)\) growth against the theoretical \(O(L)\). The ratio falls with depth, so the gap widens exactly where one would most want to rely on the formula.

Problem 3: Reproduce the degradation problem

Solution
depths_d <- c(2, 6, 12, 20, 30)
p3 <- do.call(rbind, lapply(depths_d, function(n) {
  data.frame(blocks = n,
             plain_final = tail(train_stack(plain_block, n, steps = 200), 1),
             residual_final = tail(train_stack(residual_block, n, steps = 200), 1))
}))
p3 |> mutate(across(-blocks, \(z) signif(z, 4)),
             plain_got_worse = c(NA, diff(plain_final) > 0),
             residual_got_worse = c(NA, diff(residual_final) > 0))
The plain network’s training loss increases with depth past a point, neither overfitting nor a capacity limit, since the deeper network contains the shallower one as a special case. The residual network’s does not.

Problem 4: Quantify translation sensitivity

Solution
torch_manual_seed(101)
build_net <- function(anti_alias) {
  down <- if (anti_alias)
    nn_sequential(nn_avg_pool2d(2, stride = 1, padding = 0), nn_max_pool2d(2))
  else nn_max_pool2d(2)
  nn_sequential(nn_conv2d(1, 8, 3, padding = 1), nn_relu(), down,
                nn_conv2d(8, 8, 3, padding = 1), nn_relu(), nn_max_pool2d(2),
                nn_adaptive_avg_pool2d(1), nn_flatten(), nn_linear(8, 1))
}
sweep_shift <- function(net) vapply(0:12, \(d)
  as.numeric(with_no_grad(net(shift_img(base_img, d)))), numeric(1))

r_plain <- sweep_shift(build_net(FALSE))
r_blur  <- sweep_shift(build_net(TRUE))
data.frame(
  network = c("plain max-pool", "blur then max-pool"),
  output_range = signif(c(diff(range(r_plain)), diff(range(r_blur))), 4),
  relative_variation = signif(c(diff(range(r_plain))/abs(mean(r_plain)),
                                diff(range(r_blur))/abs(mean(r_blur))), 4))
bind_rows(data.frame(shift = 0:12, out = r_plain, net = "plain max-pool"),
          data.frame(shift = 0:12, out = r_blur, net = "blur then max-pool")) |>
  ggplot(aes(shift, out, color = net)) +
  geom_line(linewidth = 1) + geom_point(size = 2) +
  facet_wrap(~ net, scales = "free_y") +
  scale_color_manual(values = c("#D8433B", "#3B7DD8"), guide = "none") +
  labs(title = "Output under one-pixel translations",
       subtitle = "Blurring before subsampling reduces aliasing and flattens the response",
       x = "Shift (pixels)", y = "Output") +
  theme_dspa(10)

Anti-aliased downsampling reduces the oscillation substantially. A flat line would be true invariance, and neither network achieves it, invariance is approximate at best.

Problem 5: Dice on empty masks

Solution
set.seed(103)
n_slices <- 200
empty_frac <- 0.6                                  # typical for volumetric MRI
truths <- lapply(1:n_slices, function(i) {
  m <- matrix(0, 32, 32)
  if (runif(1) > empty_frac) m[10:20, 12:22] <- 1
  m })
preds <- lapply(truths, function(m) {
  p <- m
  if (sum(m) > 0) { p[10:20, 12:22] <- 0; p[11:21, 13:23] <- 1 }   # slight offset
  p })

naive <- vapply(seq_len(n_slices), function(i) {
  d <- 2*sum(preds[[i]] == 1 & truths[[i]] == 1) / (sum(preds[[i]]) + sum(truths[[i]]))
  d }, numeric(1))
correct <- vapply(seq_len(n_slices), function(i)
  dice_hard(preds[[i]], truths[[i]], empty_value = 1), numeric(1))

data.frame(
  handling = c("naive (NaN on empty)", "convention D = 1 on empty"),
  n_scored = c(sum(!is.nan(naive)), length(correct)),
  n_dropped = c(sum(is.nan(naive)), 0),
  mean_dice = signif(c(mean(naive, na.rm = TRUE), mean(correct)), 4))
The naive version silently drops the empty slices, the ones the model gets perfectly right, and reports a mean over only the hard cases. Here that understates performance; had the model produced false positives on empty slices, dropping them would overstate it. Either way the reported mean is not a mean over the data.

Problem 6: Run the saliency sanity check

Solution
sanity <- function(seed) {
  torch_manual_seed(seed)
  rn <- nn_module("C3", initialize = cls_net$initialize, forward = cls_net$forward)()
  c(cam = cor(as.vector(grad_cam(cls_net, xi)), as.vector(grad_cam(rn, xi))),
    sal = cor(as.vector(saliency(cls_net, xi)), as.vector(saliency(rn, xi))))
}
s <- vapply(1:5, sanity, numeric(2))
data.frame(
  method = c("Grad-CAM", "Vanilla saliency"),
  mean_corr_trained_vs_random = signif(rowMeans(s), 3),
  sd = signif(apply(s, 1, sd), 3),
  passes_sanity_check = rowMeans(s) < 0.5)
A high correlation between the trained and randomly initialized model’s maps means the method is reporting the input’s structure, not the model’s reasoning, the method fails the check and its maps carry no information about what was learned. Report this diagnostic alongside any attribution figure.

14.42 Checkpoint

  1. Your \(56\)-layer plain CNN has higher training loss than your 20-layer one. What is happening, and what is the fix?
  2. Segmentation Dice on the validation set is 0.94; on a second hospital’s data it is 0.61. Name three candidate causes.
  3. You froze a pretrained backbone and validation accuracy still drifts across epochs. Why?
  4. Your classifier is 99.2% accurate at detecting a lesion present in 0.8% of pixels. Is it good?
  5. Grad-CAM highlights the tumour. Does that mean the model uses the tumour?
  6. Why does a GPU help far more with a CNN than with an MLP of the same parameter count?
Answers
  1. The degradation problem, an optimization failure, not overfitting. Overfitting would raise validation loss while training loss fell; here the deeper network fits the training data worse. It is not a capacity limit either, since the 56-layer network contains the 20-layer one (set the extra layers to the identity), so a solution at least as good provably exists. Gradient descent cannot find it because learning an identity map through a stack of nonlinear layers is hard, and gradients decay as a product of Jacobians. Residual connections re-parameterize the block as \(y=\mathcal F(x)+x\), making the identity free (\(\mathcal F\to0\)) and creating an additive gradient path with derivative exactly 1.
  2. Domain shift, leakage in the original split, or a prevalence change. Domain shift is most likely: different scanner vendor, field strength, acquisition protocol, or reconstruction kernel changes the intensity distribution the frozen normalization statistics assume. Leakage, if the original split was by slice rather than by patient, adjacent near-identical slices spanned train and validation, so 0.94 was never a generalization estimate. Prevalence and case mix, if lesions are smaller or rarer at the second site, Dice falls even at constant per-pixel performance, since Dice depends on lesion size. Check the split protocol first, then compare intensity histograms across sites, then stratify Dice by lesion volume.
  3. BatchNorm’s running statistics are buffers, not parameters. requires_grad_(FALSE) stops the optimizer from updating weights; it does nothing to the running mean and variance, which are updated in the forward pass whenever the module is in train() mode. So the “frozen” backbone’s normalization is silently drifting toward the new domain, and since those statistics are what eval() uses, validation accuracy changes for reasons invisible in the training loop. Freeze both: requires_grad_(FALSE) and eval() on those modules, re-applied after every model$train(), which propagates to children.
  4. Almost certainly not. At 0.8% prevalence, predicting “background everywhere” scores 99.2%, the model may have learned nothing. Accuracy is dominated by true negatives, which is exactly why segmentation uses Dice or IoU, both of which exclude \(TN\) and would score the trivial model at 0. Report Dice with its distribution across cases (not just the mean), handle empty masks by the \(D=1\) convention rather than letting them become NaN, and add the precision–recall view of Chapter 9, §9.7, whose chance level is the prevalence rather than 0.5.
  5. No. A bright region means the output is locally sensitive to those pixels, which is association, not mechanism. Three checks before believing it. Run the model-randomization sanity check: if a randomly initialized network produces a similar map, the method is reporting the input’s edges rather than anything learned. Check stability: imperceptible perturbations that leave the prediction unchanged can completely change the attribution. And test by intervention, occlude the region, or retrain on data without the suspected artifact, since documented failures where models keyed on scanner text, laterality markers, or chest drains all produced plausible-looking maps near the pathology.
  6. Arithmetic intensity. A dense layer uses each weight once per example, giving roughly \(0.5\) flops per byte moved, below the roofline ridge point, so it runs at memory bandwidth and a GPU’s advantage is only the bandwidth ratio, about \(5\)\(20\times\). A convolutional layer reuses each weight \(H'W'\) times, giving tens to hundreds of flops per byte, above the ridge, so it runs at the compute roof, where a GPU’s advantage is its flop rate, one to two orders of magnitude. Weight sharing is not only a parameter-count argument; it is what makes convolution the operation GPUs are good at.

14.43 Summary

The convolutional primitive

  • Convolution is local connectivity plus weight sharing, a structural prior for images, with kernels learned rather than designed. Parameter count is independent of image size.
  • Output size is \(\lfloor(H+2p-d(k-1)-1)/s\rfloor+1\). “Same” padding preserves shape, not border information.
  • The receptive field grows linearly with stride-1 depth and geometrically with striding, but the effective field grows only as \(O(\sqrt L)\) and is roughly Gaussian.
  • Two \(3\times3\) layers beat one \(5\times5\): same receptive field, fewer parameters, an extra nonlinearity.
  • CNNs are equivariant, not invariant. Invariance comes from pooling or augmentation, and strided layers alias, so even equivariance is imperfect.
  • \(1\times1\) convolutions are learned channel mixing; they make bottlenecks and cut cost roughly fourfold.
  • Weight reuse makes convolution compute-bound where dense layers are memory-bound, the reason GPUs help far more with CNNs.

Architecture

  • Global average pooling removes the dense head, which was ~90% of VGG’s parameters, and admits any input size.
  • The degradation problem is an optimization failure, visible in training loss. Residual connections make the identity free and create a gradient path with derivative exactly 1, the same mechanism as the LSTM cell state, applied to depth rather than time.
  • BatchNorm fails at small batches; GroupNorm is batch-independent and is the right default for segmentation and detection.

Training

  • Augmentation adds examples, not information. Each transformation asserts an invariance that must be true in your domain.
  • Transfer learning: train the head first, then unfreeze progressively at smaller learning rates, and match the pretraining normalization statistics. How much to unfreeze depends on target-set size.
  • “Frozen” does not freeze BatchNorm, requires_grad_(FALSE) and eval().
  • Pretrained features transfer weakly to distant domains; compare against a from-scratch baseline of matched capacity.

Segmentation and interpretation

  • U-Net skips are concatenative and restore spatial precision, a different purpose from residual skips.
  • Dice is the \(F_1\) score and ignores true negatives, which is why it, not accuracy, is the segmentation metric. It is undefined on empty masks (\(D=1\) by convention) and needs the soft form to be differentiable.
  • Split by patient, never by slice. Weighted sampling toward large lesions decalibrates the model against natural prevalence.
  • Saliency maps are not explanations. Run the model-randomization sanity check, and test attributions by intervention.

14.44 Where Part 2 leads

Continue with Content
Part 3: Sequence Models RNNs and LSTM gating, attention derived from the alignment problem, Transformers, and the \(O(T^2d)\)-versus-\(O(Td^2)\) comparison. Attention is the answer to §14.27’s problem: reaching global context without stacking depth
Part 4: Generative and Representation Learning The U-Net of §14.37 reappears as the backbone of diffusion models; autoencoders, the ELBO, GANs, self-supervised pretraining
Part 5: Generalization, Uncertainty, and Practice Calibration of the CNNs built here, double descent, ensembles, Bayesian hyperparameter optimization, pruning

Earlier chapters this part depended on

  • Part 1: tensors, autograd, initialization, normalization, the training loop
  • Chapter 9: the metric family, grouped splitting, calibration
  • Chapter 10: the roofline and arithmetic intensity
  • Chapter 12: the LSTM additive path, which residual connections mirror
  • Chapter 13: gradient flow, optimizers, automatic differentiation

Session information

sessionInfo()
#> R version 4.3.3 (2024-02-29 ucrt)
#> Platform: x86_64-w64-mingw32/x64 (64-bit)
#> Running under: Windows 11 x64 (build 26200)
#> 
#> Matrix products: default
#> 
#> 
#> locale:
#> [1] LC_COLLATE=English_United States.utf8 
#> [2] LC_CTYPE=English_United States.utf8   
#> [3] LC_MONETARY=English_United States.utf8
#> [4] LC_NUMERIC=C                          
#> [5] LC_TIME=English_United States.utf8    
#> 
#> time zone: America/New_York
#> tzcode source: internal
#> 
#> attached base packages:
#> [1] stats     graphics  grDevices utils     datasets  methods   base     
#> 
#> other attached packages:
#> [1] luz_0.5.2         torchvision_0.6.0 torch_0.13.0      plotly_4.12.1    
#> [5] patchwork_1.3.0   tidyr_1.3.1       dplyr_1.1.4       ggplot2_4.0.1    
#> 
#> loaded via a namespace (and not attached):
#>  [1] sass_0.4.9         generics_0.1.3     hms_1.1.3          digest_0.6.37     
#>  [5] magrittr_2.0.3     evaluate_1.0.3     grid_4.3.3         RColorBrewer_1.1-3
#>  [9] fastmap_1.2.0      jsonlite_1.8.9     processx_3.8.6     progress_1.2.3    
#> [13] ps_1.9.0           httr_1.4.7         purrr_1.0.2        crosstalk_1.2.1   
#> [17] viridisLite_0.4.2  scales_1.4.0       coro_1.0.4         codetools_0.2-20  
#> [21] jquerylib_0.1.4    cli_3.6.3          rlang_1.1.5        crayon_1.5.3      
#> [25] bit64_4.0.5        withr_3.0.2        cachem_1.1.0       yaml_2.3.10       
#> [29] otel_0.2.0         tools_4.3.3        zeallot_0.1.0      vctrs_0.6.5       
#> [33] R6_2.6.1           lifecycle_1.0.5    htmlwidgets_1.6.4  fs_1.6.5          
#> [37] bit_4.0.5          pkgconfig_2.0.3    callr_3.7.6        bslib_0.9.0       
#> [41] pillar_1.10.1      gtable_0.3.6       data.table_1.16.4  glue_1.8.0        
#> [45] Rcpp_1.0.14        xfun_0.52          tibble_3.2.1       tidyselect_1.2.1  
#> [49] rstudioapi_0.18.0  knitr_1.51         farver_2.1.2       htmltools_0.5.8.1 
#> [53] labeling_0.4.3     rmarkdown_2.31     compiler_4.3.3     prettyunits_1.2.0 
#> [57] S7_0.2.1