SOCR ≫

DSPA ≫

DSPA3 Topics ≫

DSPA3: Data Science and Predictive Analytics (UMich HS650)

Chapter 14, Part 1: Deep Learning Foundations - Tensors, Networks, and Training

SOCR/MCAIM (Ivo Dinov)

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

---

Chapter 14 is published in five parts

Deep learning is the book’s largest single topic, and a single document would neither compile nor read well. Each part below stands alone and cross-links to the others.

Part	Sections	Content
Part 1 (this document)	§14.1–§14.20	Tensors, automatic differentiation, the MLP, activations, initialization, normalization, optimizers, regularization, and the training loop
Part 2: Convolutional Networks and Vision	§14.25–§14.44	Convolution arithmetic, receptive fields, residual connections, transfer learning, CIFAR-10, interpretability
Part 3: Sequence Models: Recurrence and Attention	§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, 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, so a reference to §14.44 is unambiguous wherever you encounter it.

---

How this chapter uses graphics

Every two-dimensional figure is drawn with ggplot2 and rendered statically. Immediately after each one, the equivalent plot_ly() code appears in a chunk marked eval=FALSE, echo=TRUE.

Every three-dimensional figure is drawn with plot_ly() and evaluated, with a budget of six per part to control output size. Deep learning is full of genuinely two-parameter effects, activation variance over (depth, initialization gain), gradient magnitude over (depth, activation function), validation loss over (learning rate, batch size), and each of those is a surface whose shape is the point.

---

Learning objectives

After completing Part 1 you will be able to:

1. Explain what a tensor is, why it is not merely a multidimensional array, and how device and dtype affect computation.
2. Use torch’s autograd, and connect it to the reverse-mode differentiation of Chapter 13, §13.25.
3. Count a network’s parameters, FLOPs, and activation memory, and identify which is the binding constraint.
4. Fix the layer notation \(a_k^{(\ell)}=f\big(\sum_i w_{k,i}^{(\ell)}a_i^{(\ell-1)}+b_k^{(\ell)}\big)\) and pair each output activation with its matching loss.
5. State the universal approximation theorem precisely, including what it does not claim, and relate it to the functional completeness of NAND.
6. Explain why depth can be exponentially more efficient than width.
7. Choose an activation function from its derivative behaviour, and explain saturation and dying units.
8. Derive Xavier and He initialization from variance preservation, and demonstrate what happens when the scale is wrong.
9. Explain what batch and layer normalization actually do, and why the original explanation is disputed.
10. Distinguish Adam from AdamW, and explain why decoupled weight decay matters.
11. Distinguish samples, batches, epochs, and steps, and explain what a larger batch does and does not buy.
12. Write a training loop from first principles, diagnose its failure modes, and evaluate the result against strong non-neural baselines.

Estimated time: 10–14 hours including exercises. Prerequisites: Chapter 6 (perceptrons, backpropagation, §6.5), Chapter 9 (nested validation and calibration), Chapter 10 (memory hierarchy and arithmetic intensity, §10.14), and Chapter 13 (gradient methods, adaptive optimizers, automatic differentiation).

---

PART I: THE COMPUTATIONAL SUBSTRATE

14.0 The R deep-learning landscape

Before any code, one decision determines whether everything that follows is reproducible: which stack.

Package	Status	Python needed?	Verdict
torch	Active; binds LibTorch directly	No	Primary stack for these five parts
luz	Active; high-level fit() over torch	No	Used for the concise training loop
keras3	Active; multi-backend Keras 3	Yes, via reticulate	Alternative, shown eval=FALSE
tensorflow	Active	Yes	Backend for keras3
reticulate	Active	Yes	Python interop, when genuinely needed
nnet	Base-adjacent; one hidden layer	No	Fine for a shallow net; not deep learning
neuralnet	Maintained; small MLPs	No	Chapter 6 uses it
h2o	Active	No (needs Java)	Cluster-oriented; different use case
mxnet / MXNetR	Retired (Apache Attic, 2023)	—	Do not use
darch, deepnet, rcppDL, deepr	Archived from CRAN	—	Do not use
kerasR	Superseded by keras/keras3	Yes	Do not use

The decision: torch + luz. R torch binds LibTorch, the same C++ library underlying PyTorch, and installs from CRAN, downloading a pre-compiled binary. It requires no Python interpreter, no reticulate, no conda or virtualenv, and supports CUDA and Apple Metal.

That matters for more than convenience. A chapter whose code depends on a Python environment named "r-tensorflow" existing on the reader’s machine is not reproducible; it is a set of instructions that happen to have worked somewhere. Everything in these five parts runs from a clean R installation.

keras3 remains a reasonable choice, and where its idiom differs materially the equivalent appears in an eval=FALSE chunk. Note that Keras 3 is a multi-backend front end, it runs on TensorFlow, JAX, or PyTorch, so “Keras versus TensorFlow” is no longer the choice it once was.

Several frameworks that dominated earlier accounts of this subject are no longer live options: Theano ended development in 2017, Caffe2 merged into PyTorch in 2018, CNTK was archived by Microsoft in 2019, and Lua Torch was superseded by PyTorch in 2017. Of the classic list, TensorFlow/Keras and PyTorch are what remain.

14.1 Tensors

A tensor is a multidimensional array bundled with three things an ordinary R array does not carry: a device, a dtype, and a gradient history.

\[ \begin{aligned} \textbf{rank 0: }&\ \text{scalar} & \textbf{rank 1: }&\ \text{vector }\in\mathbb R^{n}\\ \textbf{rank 2: }&\ \text{matrix }\in\mathbb R^{m\times n} & \textbf{rank 3: }&\ \text{e.g. }(\text{time},\text{batch},\text{feature})\\ \textbf{rank 4: }&\ (\text{batch},\text{channel},\text{height},\text{width}) & & \end{aligned} \]

Common misconception: “a tensor is just an array with more dimensions.” The extra dimensions are the least interesting part. What distinguishes a tensor in a deep-learning framework is that it records the operations performed on it, so that a gradient can be replayed backward through them — and that it carries a device, so the same code executes on CPU or GPU.

The practical consequence is that x + y is not one operation but two: an arithmetic result and a node appended to a computational graph. That second effect is invisible until you call backward(), and it is the entire mechanism of §14.2.

x <- torch_tensor(matrix(1:6, nrow = 2))
c(shape = paste(dim(x), collapse = " x "),
  dtype = as.character(x$dtype),
  device = as.character(x$device))

#>   shape   dtype  device 
#> "2 x 3"  "Long"   "cpu"

# Broadcasting: shapes are aligned from the trailing dimension
a <- torch_randn(3, 1)      # 3 x 1
b <- torch_randn(1, 4)      # 1 x 4
c(broadcast_result = paste(dim(a + b), collapse = " x "))

#> broadcast_result 
#>          "3 x 4"

# dtype is a memory and precision decision, not a formality
sizes <- data.frame(
  dtype = c("float64", "float32", "float16"),
  bytes_per_element = c(8, 4, 2),
  MB_for_10M_elements = c(8, 4, 2) * 1e7 / 1e6,
  note = c("R's default; rarely needed for training",
           "the deep-learning standard",
           "mixed precision; halves memory and doubles throughput"))
sizes

float32 is the default for a reason. R uses double (float64) everywhere, and moving a model to float64 doubles both memory traffic and parameter storage while buying accuracy that gradient-based training does not use, the stochastic gradient noise of Chapter 13, §13.9 swamps the difference. Modern training goes the other way, to mixed precision: float16 for the forward and backward passes with a float32 master copy of the weights, which halves memory and roughly doubles throughput on hardware with tensor cores.

14.1.1 Devices and the transfer cost

dev <- if (cuda_is_available()) {
  "cuda"
} else if (tryCatch(backends_mps_is_available(), error = function(e) FALSE)) {
  "mps"
} else {
  "cpu" # Add a safe fallback default (e.g., CPU)
}

c(selected_device = dev)

#> selected_device 
#>           "cpu"

n_bench <- 2000
A <- torch_randn(n_bench, n_bench)
t_cpu <- system.time(for (i in 1:3) A$matmul(A))[["elapsed"]] / 3

if (dev != "cpu") {
  Ag <- A$to(device = dev)
  invisible(Ag$matmul(Ag))                       # warm up: first call compiles kernels
  t_gpu <- system.time({ for (i in 1:3) Ag$matmul(Ag)
                         if (dev == "cuda") cuda_synchronize() })[["elapsed"]] / 3
  t_move <- system.time(A$to(device = dev))[["elapsed"]]
  data.frame(operation = c("matmul on CPU", paste("matmul on", dev), "host-to-device copy"),
             seconds = round(c(t_cpu, t_gpu, t_move), 4),
             gflops = round(c(2 * n_bench^3 / t_cpu, 2 * n_bench^3 / t_gpu, NA) / 1e9, 1))
} else {
  data.frame(operation = "matmul on CPU", seconds = round(t_cpu, 4),
             gflops = round(2 * n_bench^3 / t_cpu / 1e9, 1))
}

Matrix multiplication has arithmetic intensity growing as \(O(n)\) (Chapter 10, §10.14.1), which puts it firmly in the compute-bound regime where a GPU’s advantage is real. Elementwise operations have intensity near \(0.1\) and are memory-bound on both devices, moving them to a GPU buys the bandwidth ratio, not the core-count ratio.

14.2 Automatic differentiation

Chapter 13, §13.25 established that reverse-mode automatic differentiation computes a full gradient at a constant multiple of the cost of one function evaluation, independent of the number of parameters, and that backpropagation is reverse-mode AD specialized to a layered graph. torch supplies exactly that machinery.

\[\text{forward: build the graph}\quad\longrightarrow\quad\text{backward: accumulate }\frac{\partial\mathcal L}{\partial v}\text{ at every node}\]

# f(x, y) = sin(x*y) + exp(x); the same function differentiated by hand in
# Chapter 13's reverse-mode demonstration
x <- torch_tensor(1.1, requires_grad = TRUE)
y <- torch_tensor(2.3, requires_grad = TRUE)

f <- torch_sin(x * y) + torch_exp(x)
f$backward()

x1 <- 1.1; y1 <- 2.3
data.frame(
  partial = c("df/dx", "df/dy"),
  torch_autograd = signif(c(as.numeric(x$grad), as.numeric(y$grad)), 10),
  analytic = signif(c(y1 * cos(x1 * y1) + exp(x1), x1 * cos(x1 * y1)), 10))

One backward sweep produced both partials, and would produce all of them for a million-parameter model at the same cost.

# The graph is built by the forward pass and freed by backward() unless retained
w <- torch_randn(3, 2, requires_grad = TRUE)
b <- torch_zeros(2, requires_grad = TRUE)
inp <- torch_randn(5, 3)

out <- torch_relu(inp$matmul(w) + b)
loss <- out$sum()

c(loss_requires_grad = loss$requires_grad,
  w_grad_before_backward = is.null(w$grad))

#>     loss_requires_grad w_grad_before_backward 
#>                   TRUE                  FALSE

loss$backward()
c(w_grad_shape = paste(dim(w$grad), collapse = " x "),
  b_grad_shape = paste(dim(b$grad), collapse = " x "))

#> w_grad_shape b_grad_shape 
#>      "3 x 2"          "2"

Gradients accumulate; they are not overwritten. Calling backward() twice without clearing adds the second gradient to the first. Every training loop must therefore call optimizer$zero_grad() (or zero_grad() on the module) before each backward pass. Omitting it is the single most common bug in hand-written training loops, and its symptom, a loss that decreases for a few steps and then diverges, looks exactly like a learning rate that is too large.

The accumulation is deliberate: it is what makes gradient accumulation over micro-batches possible, letting a large effective batch size fit in limited memory.

z <- torch_tensor(2.0, requires_grad = TRUE)
(z^2)$backward();  g1 <- as.numeric(z$grad)      # d(z^2)/dz = 2z = 4
(z^2)$backward();  g2 <- as.numeric(z$grad)      # accumulates: 4 + 4 = 8
z$grad$zero_()

#> torch_tensor
#>  0
#> [ CPUFloatType{1} ]

(z^2)$backward();  g3 <- as.numeric(z$grad)      # cleared first: 4

c(first_backward = g1, second_backward_accumulated = g2,
  after_zero_grad = g3, analytic = 4)

#>              first_backward second_backward_accumulated 
#>                           4                           8 
#>             after_zero_grad                    analytic 
#>                           4                           4

# Evaluation must not build a graph: with_no_grad() saves memory and time
m <- nn_linear(100, 100)
xin <- torch_randn(256, 100)

t_grad <- system.time(for (i in 1:50) m(xin))[["elapsed"]]
t_nograd <- system.time(with_no_grad({ for (i in 1:50) m(xin) }))[["elapsed"]]
c(with_graph_seconds = round(t_grad, 4),
  no_grad_seconds = round(t_nograd, 4),
  speedup = round(t_grad / t_nograd, 2))

#> with_graph_seconds    no_grad_seconds            speedup 
#>                  0                  0                NaN

14.3 Counting parameters, FLOPs, and memory

Three quantities describe a network’s cost, and they scale differently.

For a fully connected layer mapping \(d_{\text{in}}\to d_{\text{out}}\) with batch size \(B\):

\[ \begin{aligned} \textbf{Parameters: }&\quad d_{\text{in}}d_{\text{out}}+d_{\text{out}}\\ \textbf{Forward FLOPs: }&\quad \approx 2Bd_{\text{in}}d_{\text{out}}\\ \textbf{Activation memory: }&\quad Bd_{\text{out}}\ \text{elements, \emph{retained} until the backward pass} \end{aligned} \]

Activation memory, not parameter memory, is usually the binding constraint. A network’s parameters are stored once; its activations are stored once per example in the batch, and every one of them must be kept alive from the forward pass until the corresponding backward step. Training also holds a gradient for every parameter, plus optimizer state, Adam keeps two moment estimates, so Adam triples the parameter memory relative to plain SGD.

The accounting for training is roughly \[\underbrace{4P}_{\text{weights}}+\underbrace{4P}_{\text{gradients}}+\underbrace{8P}_{\text{Adam moments}}+\underbrace{4\textstyle\sum_\ell Bd_\ell}_{\text{activations}}\ \text{bytes at float32},\] and the last term is the one that scales with batch size. This is why “reduce the batch size” is the standard response to an out-of-memory error, and why gradient checkpointing, recomputing activations in the backward pass instead of storing them, trades time for space.

count_params <- function(module) {
  ps <- module$parameters
  data.frame(
    parameter = names(ps),
    shape = vapply(ps, \(p) paste(dim(p), collapse = " x "), character(1)),
    n = vapply(ps, \(p) prod(dim(p)), numeric(1)),
    row.names = NULL)
}

mlp_demo <- nn_sequential(
  nn_linear(64, 256), nn_relu(),
  nn_linear(256, 128), nn_relu(),
  nn_linear(128, 2))

pc <- count_params(mlp_demo)
pc

c(total_parameters = sum(pc$n),
  weight_memory_MB_float32 = round(4 * sum(pc$n) / 1e6, 3),
  with_grads_and_adam_MB = round(16 * sum(pc$n) / 1e6, 3))

#>         total_parameters weight_memory_MB_float32   with_grads_and_adam_MB 
#>                49794.000                    0.199                    0.797

# Activation memory grows with the batch; parameter memory does not
widths <- c(64, 256, 128, 2)
batch_sizes <- c(32, 256, 2048)
acc <- expand.grid(batch = batch_sizes) |>
  mutate(activation_MB = 4 * batch * sum(widths[-1]) / 1e6,
         parameter_MB = 4 * (64*256 + 256 + 256*128 + 128 + 128*2 + 2) / 1e6,
         ratio = round(activation_MB / parameter_MB, 2))
acc |> mutate(across(where(is.numeric), \(z) round(z, 4)))

At a batch of 32 the parameters dominate; by 2,048 the activations do. For the convolutional networks of Part 2, §14.25 the crossover happens much earlier, because a convolutional layer has few parameters and enormous activations.

---

PART II: FROM PERCEPTRON TO NETWORK

14.4 The perceptron and its limits

The perceptron computes a weighted sum and thresholds it:

\[\hat y=\mathbb 1\big\{\mathbf w^\top\mathbf x+b>0\big\}.\]

Its decision boundary is the hyperplane \(\mathbf w^\top\mathbf x+b=0\), so it can represent exactly the linearly separable functions, and XOR is not one of them.

The biological analogy is history, not mechanism. The perceptron was inspired by a schematic account of a neuron: inputs weighted by synaptic strengths, summed, and fired above a threshold. That inspiration is real and worth knowing, and it explains essentially nothing about why deep networks work.

Biological neurons are spiking, stochastic, temporally coded, and locally plastic. Backpropagation requires a global, synchronous error signal propagated backward through the transpose of the forward weights, the “weight transport problem”, for which no biological mechanism is known. Throughout this chapter a network is treated as what it mathematically is: a composition of parameterized affine maps and pointwise nonlinearities, trained by gradient descent on a differentiable objective.

xor_x <- torch_tensor(matrix(c(0,0, 0,1, 1,0, 1,1), ncol = 2, byrow = TRUE))
xor_y <- torch_tensor(matrix(c(0, 1, 1, 0), ncol = 1))

fit_xor <- function(hidden = 0, epochs = 3000, lr = 0.1, seed = 1) {
  torch_manual_seed(seed)
  net <- if (hidden == 0) nn_linear(2, 1) else
    nn_sequential(nn_linear(2, hidden), nn_tanh(), nn_linear(hidden, 1))
  opt <- optim_adam(net$parameters, lr = lr)
  for (e in seq_len(epochs)) {
    opt$zero_grad()
    loss <- nnf_binary_cross_entropy_with_logits(net(xor_x), xor_y)
    loss$backward(); opt$step()
  }
  preds <- as.numeric(torch_sigmoid(net(xor_x)))
  list(net = net, preds = preds,
       accuracy = mean((preds > 0.5) == as.numeric(xor_y)))
}

r0 <- fit_xor(hidden = 0); r2 <- fit_xor(hidden = 2)
data.frame(
  model = c("linear (no hidden layer)", "one hidden layer, 2 units"),
  accuracy = c(r0$accuracy, r2$accuracy),
  pred_00 = round(c(r0$preds[1], r2$preds[1]), 3),
  pred_01 = round(c(r0$preds[2], r2$preds[2]), 3),
  pred_10 = round(c(r0$preds[3], r2$preds[3]), 3),
  pred_11 = round(c(r0$preds[4], r2$preds[4]), 3))

The linear model converges to predicting \(0.5\) everywhere, the best a hyperplane can do on XOR. Two hidden units suffice, because they can construct a nonlinear coordinate change in which the classes are separable.

14.4.1 NAND and functional completeness

XOR shows that one hidden layer is sometimes necessary. NAND shows that it is always sufficient for Boolean logic.

\[\mathrm{NAND}(x_1,x_2)=\neg(x_1\wedge x_2)=\mathbb 1\big\{-x_1-x_2+1.5>0\big\}\]

so a single perceptron with weights \((-1,-1)\) and bias \(1.5\) computes it exactly, no hidden layer required.

NAND is functionally complete. Every Boolean function of any number of variables can be written using NAND alone: \(\neg a=\mathrm{NAND}(a,a)\), \(a\wedge b=\mathrm{NAND}(\mathrm{NAND}(a,b),\mathrm{NAND}(a,b))\), \(a\vee b=\mathrm{NAND}(\neg a,\neg b)\), and \(\{\neg,\wedge,\vee\}\) generates the rest.

Since one perceptron computes NAND, a network of perceptrons can compute any Boolean function, a discrete precursor of the universal approximation result of §14.6, and it carries the same warning. The construction says a representation exists; it says nothing about how many units it needs. Composing NAND gates to build an \(n\)-input parity function requires \(\Theta(n)\) gates in depth \(O(\log n)\), but a depth-2 circuit needs exponentially many, the Boolean analogue of the depth–width separation in §14.7.

# A single perceptron computes NAND exactly, with weights fixed by hand
nand_x <- torch_tensor(matrix(c(0,0, 0,1, 1,0, 1,1), ncol = 2, byrow = TRUE))
w_nand <- torch_tensor(matrix(c(-1, -1), ncol = 1)); b_nand <- torch_tensor(1.5)
nand_out <- as.numeric((nand_x$matmul(w_nand) + b_nand) > 0)

# XOR built from NAND gates: XOR(a,b) = NAND(NAND(a,NAND(a,b)), NAND(b,NAND(a,b)))
nand <- function(a, b) as.numeric(!(a & b))
xor_from_nand <- function(a, b) { t <- nand(a, b); nand(nand(a, t), nand(b, t)) }

data.frame(
  x1 = c(0,0,1,1), x2 = c(0,1,0,1),
  NAND_perceptron = nand_out,
  NAND_truth = c(1,1,1,0),
  XOR_from_4_NANDs = mapply(xor_from_nand, c(0,0,1,1), c(0,1,0,1)),
  XOR_truth = c(0,1,1,0))

Four NAND gates reproduce XOR exactly, the same function the single-layer network of §14.4 could not represent. Depth, not width, is what bought it.

14.4.2 What the biological analogy is worth

The historical inspiration is real and worth stating once, precisely. The human brain has roughly \(10^{11}\) neurons and \(10^{15}\) synapses. A neuron integrates afferent input across its dendrites, and fires an action potential down its axon only when the aggregated membrane potential crosses a threshold. Synaptic strengths change with experience, and can be excitatory (\(w>0\)) or inhibitory (\(w<0\)). A perceptron’s weighted sum and threshold is a first-order rate-coding approximation of that: it models the frequency of firing and discards the timing of individual spikes.

Common misconception: “neural networks work because they simulate the brain.” The analogy explains where the vocabulary came from and almost nothing about why the method works.

Biological neurons are spiking, stochastic, temporally coded, and locally plastic, a synapse updates from signals available at that synapse. Backpropagation requires a global, synchronous error signal propagated backward through the transpose of the forward weights, and no biological mechanism is known that transports weights in this way. It is called the weight transport problem, and proposed biologically plausible alternatives — feedback alignment, predictive coding, remain research questions rather than settled accounts.

Treat a network as what it mathematically is: a composition of parameterized affine maps and pointwise nonlinearities, fitted by gradient descent on a differentiable objective. Every result in these five parts follows from that description, and none follows from neuroscience.

14.5 The multilayer perceptron

An \(L\)-layer network is a composition

\[\boxed{\;f(\mathbf x)=W^{(L)}\,\sigma\!\Big(W^{(L-1)}\sigma\big(\cdots\sigma(W^{(1)}\mathbf x+\mathbf b^{(1)})\cdots\big)+\mathbf b^{(L-1)}\Big)+\mathbf b^{(L)}\;}\]

with \(\sigma\) applied elementwise. The nonlinearity is what makes depth meaningful: without it the composition of affine maps \(W^{(L)}\cdots W^{(1)}\) is a single affine map, and a hundred layers collapse to one.

torch_manual_seed(7)
# A deep network with NO activation is exactly a linear map
deep_linear <- nn_sequential(nn_linear(5, 40), nn_linear(40, 40),
                           nn_linear(40, 40), nn_linear(40, 3))
xd <- torch_randn(100, 5)
Yd <- with_no_grad(deep_linear(xd))

# Recover the single equivalent affine map from the composed weights (layers 1 to 4)
Ws <- lapply(1:4, \(i) deep_linear[[i]]$weight)
bs <- lapply(1:4, \(i) deep_linear[[i]]$bias)

W_eq <- Ws[[4]]$matmul(Ws[[3]])$matmul(Ws[[2]])$matmul(Ws[[1]])
b_eq <- Ws[[4]]$matmul(Ws[[3]])$matmul(Ws[[2]])$matmul(bs[[1]]$unsqueeze(2))$squeeze() +
        Ws[[4]]$matmul(Ws[[3]])$matmul(bs[[2]]$unsqueeze(2))$squeeze() +
        Ws[[4]]$matmul(bs[[3]]$unsqueeze(2))$squeeze() + bs[[4]]
Y_eq <- xd$matmul(W_eq$t()) + b_eq

c(max_abs_difference = signif(as.numeric(torch_max(torch_abs(Yd - Y_eq))), 3),
  conclusion = "four linear layers = one linear layer")

#>                      max_abs_difference                              conclusion 
#>                              "1.19e-07" "four linear layers = one linear layer"

14.5.1 Layer notation

Writing the same network unit by unit fixes the indices, which is worth doing once because conventions differ between references.

\[\boxed{\;a_k^{(\ell)}=f\!\left(\underbrace{\sum_{i}w_{k,i}^{(\ell)}\,a_i^{(\ell-1)}+b_k^{(\ell)}}_{z_k^{(\ell)}\ \text{(pre-activation)}}\right)\;}\]

• \(a_i^{(\ell-1)}\) is the activation of unit \(i\) in the previous layer;
• \(w_{k,i}^{(\ell)}\) carries it to unit \(k\) of the current layer \(\ell\);
• \(b_k^{(\ell)}\) is unit \(k\)’s bias, and \(f\) the activation function.

\(W^{(\ell)}\) connects layer \(\ell-1\) to layer \(\ell\), and nothing else. Under the convention above, \(W^{(\ell)}\) has \(n_\ell\) rows (one per unit in the current layer) and \(n_{\ell-1}\) columns, so \(\mathbf z^{(\ell)}=W^{(\ell)}\mathbf a^{(\ell-1)}+\mathbf b^{(\ell)}\) is conformable. A statement that its columns index the next layer \((\ell+1)\) describes \(W^{(\ell+1)}\), and mixing the two makes every shape argument in a derivation wrong by one layer.

torch’s nn_linear(in_features, out_features) stores weight with shape \((\text{out},\text{in})\), matching the convention here, and computes \(\mathbf x W^\top+\mathbf b\) so that the batch dimension leads. Print dim() when in doubt; a silent transpose is a common source of shape errors that broadcast rather than error.

The bias is a learned intercept, not a residual. It shifts the decision boundary, exactly as the intercept does in linear regression (Chapter 3). It is not “information the model fails to explain”, that would be the residual \(y-\hat y\), which is a property of the fit, whereas \(b\) is a parameter the fit estimates.

net_s <- nn_sequential(nn_linear(7, 5), nn_tanh(), nn_linear(5, 3), nn_tanh(),
                       nn_linear(3, 2))
data.frame(
  layer = c("W(1): 7 -> 5", "W(2): 5 -> 3", "W(3): 3 -> 2"),
  weight_shape = c(paste(dim(net_s[[1]]$weight), collapse = " x "),
                   paste(dim(net_s[[3]]$weight), collapse = " x "),
                   paste(dim(net_s[[5]]$weight), collapse = " x ")),
  rows_are = "units in the CURRENT layer",
  columns_are = "units in the PREVIOUS layer")

# The batch dimension leads, and shapes chain through
xb <- torch_randn(16, 7)
c(input = paste(dim(xb), collapse = " x "),
  output = paste(dim(with_no_grad(net_s(xb))), collapse = " x "))

#>    input   output 
#> "16 x 7" "16 x 2"

# A fully connected 7-5-3-2 network, drawn to fix the notation above
layers_n <- c(7, 5, 3, 2)
nodes <- do.call(rbind, lapply(seq_along(layers_n), \(l)
  data.frame(layer = l, unit = seq_len(layers_n[l]),
             x = l, y = seq_len(layers_n[l]) - (layers_n[l] + 1)/2)))
edges <- do.call(rbind, lapply(seq_len(length(layers_n) - 1), function(l) {
  a <- nodes[nodes$layer == l, ]; b <- nodes[nodes$layer == l + 1, ]
  expand.grid(i = seq_len(nrow(a)), k = seq_len(nrow(b))) |>
    mutate(x = a$x[i], y = a$y[i], xend = b$x[k], yend = b$y[k], layer = l + 1)
}))
nodes$role <- c("Input", "Hidden", "Hidden", "Output")[nodes$layer]

ggplot() +
  geom_segment(data = edges, aes(x, y, xend = xend, yend = yend),
               color = "grey78", linewidth = 0.25) +
  geom_point(data = nodes, aes(x, y, color = role), size = 6) +
  scale_color_manual(values = c(Input = "#7FB069", Hidden = "#9EC5E8",
                                 Output = "#D8433B")) +
  scale_x_continuous(breaks = 1:4,
                     labels = c("l = 1\ninput", "l = 2\nhidden",
                                "l = 3\nhidden", "l = 4\noutput")) +
  labs(title = "A fully connected feed-forward network, 7-5-3-2",
       subtitle = expression("Each edge carries a weight "*w[list(k,i)]^(l)*
                             "; each non-input node adds a bias "*b[k]^(l)),
       x = NULL, y = NULL, color = NULL) +
  theme_dspa(10) +
  theme(axis.text.y = element_blank(), panel.grid = element_blank())

# --- Interactive equivalent ------------------------------------------------
plot_ly() |>
  add_segments(data = edges, x = ~x, y = ~y, xend = ~xend, yend = ~yend,
               line = list(color = "lightgray", width = 1), showlegend = FALSE) |>
  add_markers(data = nodes, x = ~x, y = ~y, color = ~role,
              marker = list(size = 22),
              text = ~paste0("a(", unit, ",", layer, ")"), hoverinfo = "text") |>
  layout(title = "Fully connected feed-forward network",
         xaxis = list(title = "Layer"), yaxis = list(title = ""))

14.5.2 Output layers and losses

The final layer’s activation and the loss must be chosen together, because each loss presumes a particular output parameterization.

\[ \begin{aligned} \textbf{Regression: }&\ \text{identity output} &\ \mathcal L&=\tfrac1n\textstyle\sum_i(y_i-\hat y_i)^2\\[1mm] \textbf{Binary: }&\ \hat p=\sigma(z)=\tfrac{1}{1+e^{-z}} &\ \mathcal L&=-\tfrac1n\textstyle\sum_i\big[y_i\log\hat p_i+(1-y_i)\log(1-\hat p_i)\big]\\[1mm] \textbf{Multiclass: }&\ \hat p_j=\mathrm{softmax}(z)_j=\dfrac{e^{z_j}}{\sum_{m}e^{z_m}} &\ \mathcal L&=-\tfrac1n\textstyle\sum_i\log\hat p_{i,y_i} \end{aligned} \]

Cross-entropy is the negative log-likelihood of the corresponding distribution, Bernoulli for binary, categorical for multiclass, so minimizing it is maximum likelihood estimation (Chapter 12). It is also a strictly proper scoring rule (Chapter 9, §9.8.2), which is why it is minimized by reporting the true probability and why accuracy is not a training objective.

Use the fused *_with_logits losses; never compose sigmoid or softmax with a separate log. Computing \(\log\sigma(z)\) in two steps overflows for large \(|z|\): \(\sigma(z)\) rounds to exactly 0 or 1, and \(\log 0=-\infty\) produces a NaN that propagates through the whole backward pass.

The fused form uses the numerically stable identity \[\log\sigma(z)=-\log\big(1+e^{-z}\big)=-\max(-z,0)-\log\big(1+e^{-|z|}\big),\] and the softmax analogue subtracts \(\max_j z_j\) before exponentiating. In torch these are nnf_binary_cross_entropy_with_logits() and nnf_cross_entropy(), and both take raw logits, not probabilities — passing an already-sigmoided value applies the nonlinearity twice.

z_big <- torch_tensor(c(-40, -10, 0, 10, 40))
y_tgt <- torch_tensor(c(0, 0, 1, 1, 1))

naive <- with_no_grad({
  p <- torch_sigmoid(z_big)
  -(y_tgt * torch_log(p) + (1 - y_tgt) * torch_log(1 - p))
})
fused <- with_no_grad(
  nnf_binary_cross_entropy_with_logits(z_big, y_tgt, reduction = "none"))

data.frame(logit = as.numeric(z_big), target = as.numeric(y_tgt),
           naive_two_step = signif(as.numeric(naive), 4),
           fused_with_logits = signif(as.numeric(fused), 4))

At \(|z|=40\) the two-step computation returns Inf or NaN while the fused form returns the correct finite value. This is the most common source of a NaN loss on the very first backward pass (§14.16).

logits <- torch_tensor(matrix(c(2.0, 1.0, 0.1,
                                1.0, 3.0, 0.2), nrow = 2, byrow = TRUE))
probs <- nnf_softmax(logits, dim = 2)
c(row_sums = paste(round(as.numeric(torch_sum(probs, dim = 2)), 6), collapse = ", "),
  note = "softmax outputs a probability vector per row")

#>                                       row_sums 
#>                                         "1, 1" 
#>                                           note 
#> "softmax outputs a probability vector per row"

# Softmax is shift-invariant: subtracting max_j z_j changes nothing but prevents overflow
shifted <- nnf_softmax(logits - torch_max(logits), dim = 2)
c(max_abs_difference = signif(as.numeric(torch_max(torch_abs(probs - shifted))), 3))

#> max_abs_difference 
#>           7.45e-09

14.6 Universal approximation

Universal approximation theorem (Cybenko 1989; Hornik 1991). Let \(\sigma\) be a non-polynomial continuous function. For any continuous \(f\) on a compact \(K\subset\mathbb R^d\) and any \(\varepsilon>0\), there exist \(N\), weights \(\mathbf w_i\), biases \(b_i\), and coefficients \(\alpha_i\) with \[\sup_{\mathbf x\in K}\left|f(\mathbf x)-\sum_{i=1}^{N}\alpha_i\,\sigma\big(\mathbf w_i^\top\mathbf x+b_i\big)\right|<\varepsilon.\]

Common misconception: “universal approximation means a neural network can learn any function.” The theorem is an existence result about representation. It makes no claim about any of the four things that determine whether a network is useful.

It does not bound \(N\). The required width can grow exponentially in the input dimension, approximating a general Lipschitz function on \([0,1]^d\) to accuracy \(\varepsilon\) needs \(N=\Omega(\varepsilon^{-d})\) units. The theorem is satisfied by networks too large to instantiate.

It says nothing about learnability. Weights that achieve the approximation exist; nothing says gradient descent finds them, and the objective is non-convex (Chapter 13, §13.14).

It says nothing about generalization. Approximation is measured on \(K\) with \(f\) known. A finite training sample constrains the function only where data exist.

It is not special to neural networks. Polynomials (Stone–Weierstrass), splines, Fourier series, and kernel machines are all universal approximators. Universality is a low bar; it is not why deep learning works.

# A one-hidden-layer network CAN approximate a wiggly 1-D function -- with
# enough units. Watch how many it takes.
set.seed(11); torch_manual_seed(11)
f_target <- function(x) sin(3 * x) + 0.4 * sin(11 * x)
xs <- seq(-3, 3, length.out = 400)
xt <- torch_tensor(matrix(xs, ncol = 1))
yt <- torch_tensor(matrix(f_target(xs), ncol = 1))

fit_width <- function(w, epochs = 2500) {
  torch_manual_seed(11)
  net <- nn_sequential(nn_linear(1, w), nn_tanh(), nn_linear(w, 1))
  opt <- optim_adam(net$parameters, lr = 0.02)
  for (e in seq_len(epochs)) {
    opt$zero_grad(); l <- nnf_mse_loss(net(xt), yt); l$backward(); opt$step()
  }
  list(pred = as.numeric(with_no_grad(net(xt))),
       mse = as.numeric(with_no_grad(nnf_mse_loss(net(xt), yt))))
}

widths_try <- c(2, 8, 32, 128)
uw <- lapply(widths_try, fit_width)
data.frame(hidden_units = widths_try,
           parameters = 3 * widths_try + 1,
           final_mse = signif(vapply(uw, \(z) z$mse, numeric(1)), 4))

bind_rows(lapply(seq_along(widths_try), \(i)
  data.frame(x = xs, y = uw[[i]]$pred,
             w = sprintf("%d hidden units", widths_try[i])))) |>
  mutate(w = factor(w, levels = sprintf("%d hidden units", widths_try))) |>
  ggplot(aes(x, y)) +
  geom_line(data = data.frame(x = xs, y = f_target(xs)), aes(x, y),
            color = "grey35", linetype = "dashed", linewidth = 0.7) +
  geom_line(color = "#3B7DD8", linewidth = 0.9) +
  facet_wrap(~ w) +
  labs(title = "Universal approximation is about existence, and existence needs width",
       subtitle = "Dashed: the target. One hidden layer, tanh activation, identical training budget",
       x = "x", y = "f(x)") +
  theme_dspa(10)

14.7 Depth versus width

If one hidden layer suffices in principle, why build deep networks?

Depth buys exponential efficiency. For ReLU networks, the number of linear regions a network can carve out of its input space grows polynomially in width but exponentially in depth (Montúfar et al., 2014). A network of depth \(L\) and width \(w\) on \(d\) inputs realizes on the order of \[\Big(\tfrac{w}{d}\Big)^{d(L-1)}w^{d}\quad\text{regions},\] so adding a layer multiplies the count while adding units merely raises a base.

There are functions computable by a depth-\(k\) network with polynomially many units that require exponentially many units at depth \(k-1\) (Telgarsky, 2016). Depth is not a convenience, it is a representational resource with no shallow substitute.

The countervailing fact is that depth makes optimization harder: gradients traverse more factors (§14.8), which is what initialization (§14.9), normalization (§14.10), and residual connections (Part 2, §14.28) exist to repair.

# Count linear regions empirically for 1-D ReLU networks of equal parameter count
count_regions <- function(depth, width, n_grid = 20000, seed = 3) {
  torch_manual_seed(seed)
  layers <- list(nn_linear(1, width), nn_relu())
  if (depth > 1) for (i in 2:depth) layers <- c(layers, list(nn_linear(width, width), nn_relu()))
  layers <- c(layers, list(nn_linear(width, 1)))
  net <- do.call(nn_sequential, layers)
  xg <- torch_tensor(matrix(seq(-4, 4, length.out = n_grid), ncol = 1))
  yg <- as.numeric(with_no_grad(net(xg)))
  slope <- diff(yg)
  # A new linear region begins wherever the slope changes appreciably
  sum(abs(diff(slope)) > 1e-9) + 1
}

data.frame(
  architecture = c("depth 1, width 24", "depth 2, width 12",
                   "depth 3, width 8", "depth 4, width 6"),
  depth = c(1, 2, 3, 4), width = c(24, 12, 8, 6),
  linear_regions = c(count_regions(1, 24), count_regions(2, 12),
                     count_regions(3, 8), count_regions(4, 6)))

Roughly matched parameter budgets, very different expressive capacity, the deeper networks partition the input into far more linear pieces.

---

PART III: MAKING TRAINING WORK

Sections 14.8–14.13 are, collectively, the answer to a single question: why does a network that is expressive enough often fail to train?

14.8 Activation functions

\[ \begin{aligned} \textbf{Sigmoid: }&\ \sigma(z)=\tfrac{1}{1+e^{-z}}, &\ \sigma'(z)&=\sigma(z)\big(1-\sigma(z)\big)\ \le\ \tfrac14\\ \textbf{Tanh: }&\ \tanh(z), &\ \tanh'(z)&=1-\tanh^2(z)\ \le\ 1\\ \textbf{ReLU: }&\ \max(0,z), &\ \tfrac{d}{dz}&=\mathbb 1\{z>0\}\\ \textbf{Leaky ReLU: }&\ \max(\alpha z,z), &\ \tfrac{d}{dz}&=\alpha\ \text{or}\ 1\\ \textbf{GELU: }&\ z\,\Phi(z), &\ \tfrac{d}{dz}&=\Phi(z)+z\phi(z)\\ \textbf{SiLU/Swish: }&\ z\,\sigma(z), &\ \tfrac{d}{dz}&=\sigma(z)\big(1+z(1-\sigma(z))\big) \end{aligned} \]

The derivative bounds are the whole story. Backpropagation through \(L\) layers multiplies \(L\) Jacobians (Chapter 12, §12.26.1), and each carries a factor of \(\sigma'\). With sigmoid, \(\sigma'\le\frac14\), so the gradient shrinks by at least \(4^{-L}\), at \(L=10\) that is \(10^{-6}\) before the weights are even considered. This single bound is why sigmoid activations made deep networks untrainable for two decades.

Common misconception: “ReLU solves the vanishing-gradient problem.” It removes the saturation factor, \(\sigma'=1\) on the positive side, so no shrinkage there, and introduces a different pathology.

A ReLU unit whose pre-activation is negative for every training example has gradient exactly zero, receives no update, and stays that way forever. It is a dead unit, and a large learning rate can kill a substantial fraction of a layer in a single step. Leaky ReLU (\(\alpha\approx0.01\)), GELU, and SiLU all keep a small negative-side gradient for exactly this reason.

ReLU also does not remove the weight contribution to gradient scaling. The product \(\prod_\ell W^{(\ell)\top}\) still explodes or vanishes unless the weights are scaled correctly, which is what §14.9 is for.

zs <- seq(-4, 4, length.out = 500)
act <- bind_rows(
  data.frame(z = zs, v = 1/(1+exp(-zs)), d = (1/(1+exp(-zs)))*(1-1/(1+exp(-zs))), a = "Sigmoid"),
  data.frame(z = zs, v = tanh(zs), d = 1 - tanh(zs)^2, a = "Tanh"),
  data.frame(z = zs, v = pmax(0, zs), d = as.numeric(zs > 0), a = "ReLU"),
  data.frame(z = zs, v = zs * pnorm(zs), d = pnorm(zs) + zs * dnorm(zs), a = "GELU"),
  data.frame(z = zs, v = zs/(1+exp(-zs)),
             d = (1/(1+exp(-zs)))*(1 + zs*(1-1/(1+exp(-zs)))), a = "SiLU"))

p_v <- ggplot(act, aes(z, v, color = a)) + geom_line(linewidth = 0.9) +
  scale_color_viridis_d(option = "turbo", end = 0.9) +
  labs(title = "Activation functions", x = NULL, y = expression(sigma(z)), color = NULL) +
  theme_dspa(9)
p_d <- ggplot(act, aes(z, d, color = a)) +
  geom_hline(yintercept = 0.25, linetype = "dotted", color = "grey40") +
  geom_line(linewidth = 0.9) +
  annotate("text", x = 3.4, y = 0.31, label = "sigmoid ceiling 1/4",
           size = 3, color = "grey35") +
  scale_color_viridis_d(option = "turbo", end = 0.9, guide = "none") +
  labs(title = "Their derivatives -- the factor multiplied at every layer",
       x = "z", y = expression(sigma*"'"*(z))) +
  theme_dspa(9)
p_v / p_d

# --- Interactive equivalent ------------------------------------------------
plot_ly(act, x = ~z, y = ~d, color = ~a, type = "scatter", mode = "lines") |>
  add_lines(x = range(zs), y = c(0.25, 0.25), name = "sigmoid ceiling",
            line = list(dash = "dot", color = "grey")) |>
  layout(title = "Activation derivatives",
         xaxis = list(title = "z"), yaxis = list(title = "derivative"))

# Dying ReLU, measured: fraction of units with zero gradient after training
count_dead <- function(lr, seed = 21, steps = 400, width = 256) {
  torch_manual_seed(seed); set.seed(seed)
  net <- nn_sequential(nn_linear(20, width), nn_relu(), nn_linear(width, 1))
  opt <- optim_sgd(net$parameters, lr = lr)
  X <- torch_randn(512, 20); Y <- torch_randn(512, 1)
  for (s in seq_len(steps)) {
    opt$zero_grad(); l <- nnf_mse_loss(net(X), Y); l$backward(); opt$step()
  }
  h <- with_no_grad(torch_relu(net[[1]](X)))
  as.numeric(torch_mean((torch_sum(h > 0, dim = 1) == 0)$to(dtype = torch_float())))
}
data.frame(learning_rate = c(0.01, 0.1, 0.5, 1.0),
           fraction_dead_units = vapply(c(0.01, 0.1, 0.5, 1.0), count_dead, numeric(1)))

The fraction of permanently dead units rises sharply with the learning rate. Those units are not slow, they are gone, and no subsequent training recovers them.

# Gradient magnitude reaching layer 1, as a function of depth and activation
grad_at_depth <- function(depth, act_fn, width = 48, seed = 5) {
  torch_manual_seed(seed)
  layers <- list(nn_linear(width, width), act_fn())
  if (depth > 1) for (i in 2:depth) layers <- c(layers, list(nn_linear(width, width), act_fn()))
  layers <- c(layers, list(nn_linear(width, 1)))
  net <- do.call(nn_sequential, layers)
  X <- torch_randn(64, width)
  out <- net(X)$sum(); out$backward()
  as.numeric(torch_norm(net[[1]]$weight$grad))
}
depths <- 1:12
acts <- list(Sigmoid = nn_sigmoid, Tanh = nn_tanh, ReLU = nn_relu, GELU = nn_gelu)
Zg <- t(vapply(names(acts), function(a)
  vapply(depths, \(d) grad_at_depth(d, acts[[a]]), numeric(1)), numeric(length(depths))))

plot_ly(x = depths, y = names(acts), z = log10(pmax(Zg, 1e-30)), type = "surface",
        colorscale = "Viridis",
        colorbar = list(title = "log10 gradient\nnorm at layer 1")) |>
  layout(title = "Gradient reaching the first layer, by depth and activation",
         scene = list(xaxis = list(title = "Hidden layers"),
                      yaxis = list(title = "Activation"),
                      zaxis = list(title = "log10 ||grad||")))

Rotate along the depth axis. The sigmoid surface falls away steeply, that is \(4^{-L}\) made visible. ReLU and GELU stay far flatter, which is why they made depth practical.

---

14.9 Initialization

Common misconception: “initialization is just where you start; training fixes it.” For a deep network, initialization determines whether training happens at all. The wrong scale makes activations vanish or explode geometrically with depth, and a network whose forward signal has already collapsed to zero by layer 20 has no gradient to descend.

The derivation. Consider layer \(\ell\) with \(n_{\text{in}}\) inputs, weights drawn i.i.d. with mean 0 and variance \(\operatorname{Var}(W)\), and inputs with variance \(\operatorname{Var}(a^{(\ell-1)})\). For the pre-activation \(z_i=\sum_{j=1}^{n_{\text{in}}}W_{ij}a_j\),

\[\operatorname{Var}(z)=n_{\text{in}}\operatorname{Var}(W)\operatorname{Var}(a^{(\ell-1)}).\]

To keep the variance constant across layers we need \(n_{\text{in}}\operatorname{Var}(W)=1\). Propagating gradients backward gives the same condition with \(n_{\text{out}}\). Two standard resolutions:

\[ \begin{aligned} \textbf{Xavier/Glorot }(\tanh,\ \text{sigmoid}):&\quad \operatorname{Var}(W)=\frac{2}{n_{\text{in}}+n_{\text{out}}}\\[2mm] \textbf{He/Kaiming }(\text{ReLU}):&\quad \operatorname{Var}(W)=\frac{2}{n_{\text{in}}} \end{aligned} \]

Why He carries the extra factor of 2. ReLU zeroes half its inputs in expectation for a symmetric pre-activation distribution, so \(\operatorname{Var}(\mathrm{ReLU}(z))=\tfrac12\operatorname{Var}(z)\). Doubling the weight variance exactly compensates for that halving.

propagate_variance <- function(gain, depth = 30, width = 256, act = nn_relu,
                               seed = 31) {
  torch_manual_seed(seed)
  h <- torch_randn(512, width)
  vars <- numeric(depth)
  a <- act()
  for (l in seq_len(depth)) {
    W <- torch_randn(width, width) * sqrt(gain / width)
    h <- a(h$matmul(W$t()))
    vars[l] <- as.numeric(torch_var(h))
  }
  vars
}

gains <- c(0.5, 1, 2, 4)     # gain = 2 is He initialization for ReLU
iv <- bind_rows(lapply(gains, \(g)
  data.frame(layer = 1:30, variance = pmax(propagate_variance(g), 1e-30),
             gain = sprintf("gain = %.1f%s", g, ifelse(g == 2, "  (He)", "")))))

ggplot(iv, aes(layer, variance, color = gain)) +
  geom_hline(yintercept = 1, linetype = "dashed", color = "grey40") +
  geom_line(linewidth = 1) +
  scale_y_log10() +
  scale_color_viridis_d(option = "plasma", end = 0.9) +
  labs(title = "Activation variance through 30 ReLU layers",
       subtitle = "Dashed: variance preserved. Only gain = 2 (He) holds it; the others vanish or explode geometrically",
       x = "Layer", y = "Activation variance (log scale)", color = NULL) +
  theme_dspa()

gain_grid <- seq(0.25, 4.5, length.out = 40)
depth_grid <- 1:30
Zi <- t(vapply(gain_grid, \(g) log10(pmax(propagate_variance(g, depth = 30), 1e-30)),
               numeric(30)))

plot_ly(x = depth_grid, y = gain_grid, z = Zi, type = "surface",
        colorscale = "RdBu",
        colorbar = list(title = "log10\nvariance")) |>
  add_surface(x = depth_grid, y = gain_grid,
              z = matrix(0, length(gain_grid), length(depth_grid)),
              opacity = 0.3, showscale = FALSE,
              colorscale = list(c(0, "black"), c(1, "black"))) |>
  layout(title = "Activation variance over depth and initialization gain; the flat plane is variance 1",
         scene = list(xaxis = list(title = "Layer depth"),
                      yaxis = list(title = "Initialization gain"),
                      zaxis = list(title = "log10 variance", range = c(-12, 8))))

The flat black plane marks preserved variance. The colored surface crosses it along a single ridge at gain 2, the He value, and falls away exponentially on both sides. Rotate to the far edge: by layer 30 the difference between gain 1 and gain 2 is many orders of magnitude.

# Does the initialization scale actually decide whether the network trains?
train_deep <- function(gain, depth = 20, width = 128, steps = 300, seed = 41) {
  torch_manual_seed(seed); set.seed(seed)
  X <- torch_randn(512, width); Y <- torch_randn(512, 1)
  layers <- list()
  for (l in seq_len(depth)) layers <- c(layers, list(nn_linear(width, width), nn_relu()))
  layers <- c(layers, list(nn_linear(width, 1)))
  net <- do.call(nn_sequential, layers)
  for (m in net$modules) if (inherits(m, "nn_linear")) {
    with_no_grad({ m$weight$normal_(0, sqrt(gain / width)); m$bias$zero_() })
  }
  opt <- optim_sgd(net$parameters, lr = 0.01)
  l0 <- as.numeric(with_no_grad(nnf_mse_loss(net(X), Y)))
  for (s in seq_len(steps)) {
    opt$zero_grad(); l <- nnf_mse_loss(net(X), Y); l$backward(); opt$step()
  }
  lf <- as.numeric(with_no_grad(nnf_mse_loss(net(X), Y)))
  c(initial_loss = l0, final_loss = lf, reduction = 1 - lf/l0)
}
as.data.frame(do.call(rbind, lapply(c(0.5, 1, 2, 4), function(g)
  c(gain = g, train_deep(g))))) |>
  mutate(across(everything(), \(z) signif(z, 4)))

At gain 0.5 the 20-layer network barely moves, the forward signal has already collapsed. The architecture is identical in every row; only the initial weight scale differs.

# torch's built-in initializers, and the fan values they use
lin <- nn_linear(64, 32)
fan_in <- 64; fan_out <- 32
data.frame(
  scheme = c("Xavier uniform", "Xavier normal", "He (Kaiming) normal, ReLU",
             "torch nn_linear default"),
  weight_sd_or_bound = signif(c(
    sqrt(6 / (fan_in + fan_out)),          # uniform bound
    sqrt(2 / (fan_in + fan_out)),          # normal sd
    sqrt(2 / fan_in),                      # He normal sd
    as.numeric(torch_std(lin$weight))), 4),
  note = c("bound of U(-b, b)", "sd of N(0, sd^2)", "sd of N(0, sd^2)",
           "empirical sd of the default"))

14.10 Normalization

Batch normalization standardizes each feature across the batch, then applies a learned affine map:

\[\hat z_i=\frac{z_i-\mu_{\mathcal B}}{\sqrt{\sigma^2_{\mathcal B}+\epsilon}}, \qquad y_i=\gamma\hat z_i+\beta,\]

with \(\mu_{\mathcal B},\sigma^2_{\mathcal B}\) computed over the batch dimension, and \(\gamma,\beta\) learned per feature.

Layer normalization standardizes across the feature dimension for each example independently, no batch statistics at all.

\[ \begin{aligned} \textbf{BatchNorm: }&\ \text{normalize over }(\text{batch}) \ \text{per feature} &&\text{batch-dependent; needs running statistics at test time}\\ \textbf{LayerNorm: }&\ \text{normalize over }(\text{features}) \ \text{per example} &&\text{batch-independent; identical at train and test} \end{aligned} \]

Common misconception: “batch normalization works by reducing internal covariate shift.” That was the original explanation, and the evidence against it is direct: deliberately injecting covariate shift after a BatchNorm layer does not degrade training, and networks with badly shifting internal distributions train fine when BatchNorm is present (Santurkar et al., 2018).

The better-supported account is that normalization smooths the loss landscape, it reduces the Lipschitz constant of the loss and of its gradient, which permits larger stable learning rates (Chapter 13, §13.5.1: the bound is \(\eta<2/L\), and normalization lowers \(L\)).

Two practical consequences follow, and both matter. BatchNorm couples the examples in a batch: an individual prediction depends on the other examples it happened to be batched with, which is why it degrades at small batch sizes and why it needs separate running statistics at test time. And it is the reason Transformers use LayerNorm instead (Part 3, §14.45) — sequence models have variable lengths and often tiny batches.

torch_manual_seed(51)
z <- torch_randn(8, 6) * 4 + 10        # badly scaled activations

bn <- nn_batch_norm1d(6); ln <- nn_layer_norm(6)
zb <- bn(z); zl <- ln(z)

data.frame(
  tensor = c("input", "after BatchNorm", "after LayerNorm"),
  mean_over_batch = signif(c(as.numeric(torch_mean(z)),
                             as.numeric(torch_mean(zb)),
                             as.numeric(torch_mean(zl))), 3),
  sd_per_feature = signif(c(as.numeric(torch_mean(torch_std(z, dim = 1))),
                            as.numeric(torch_mean(torch_std(zb, dim = 1))),
                            as.numeric(torch_mean(torch_std(zl, dim = 1)))), 3),
  sd_per_example = signif(c(as.numeric(torch_mean(torch_std(z, dim = 2))),
                            as.numeric(torch_mean(torch_std(zb, dim = 2))),
                            as.numeric(torch_mean(torch_std(zl, dim = 2)))), 3))

BatchNorm drives the per-feature standard deviation to 1; LayerNorm drives the per-example one. They normalize orthogonal directions of the same matrix.

# BatchNorm's batch statistics are noisy at small batch sizes
bn_noise <- function(bs, reps = 200, feat = 32, seed = 53) {
  torch_manual_seed(seed)
  full <- torch_randn(4096, feat)
  pop_var <- as.numeric(torch_var(full, dim = 1)$mean())
  est <- vapply(seq_len(reps), function(r) {
    idx <- sample(4096, bs)
    as.numeric(torch_var(full[idx, ], dim = 1)$mean())
  }, numeric(1))
  c(batch_size = bs, relative_sd = sd(est) / pop_var)
}
bnz <- as.data.frame(do.call(rbind, lapply(c(2, 4, 8, 16, 32, 128, 512), bn_noise)))
bnz |> mutate(relative_sd = round(relative_sd, 4))

ggplot(bnz, aes(batch_size, relative_sd)) +
  geom_line(linewidth = 1, color = "steelblue") + geom_point(size = 2.4) +
  scale_x_log10(breaks = bnz$batch_size) +
  labs(title = "BatchNorm's variance estimate is noisy at small batch sizes",
       subtitle = "Relative SD of the batch variance estimate. Below ~16 the normalization itself becomes a source of noise",
       x = "Batch size (log scale)", y = "Relative SD of batch variance") +
  theme_dspa()

14.11 Optimizers in practice

Chapter 13, §13.8 derived AdaGrad, RMSProp, and Adam and identified them as diagonal preconditioners. Two practical points that theory does not settle.

14.11.1 Adam is not always the answer

Adam converges faster early and often generalizes worse than well-tuned SGD with momentum on vision tasks. The adaptive scaling that speeds up early progress also shrinks the effective step in high-curvature directions, changing which minimum the trajectory settles into. In practice: SGD+momentum for convolutional vision models, Adam/AdamW for Transformers and anything with sparse or badly scaled gradients.

14.11.2 AdamW: weight decay is not \(L_2\) regularization

Common misconception: “adding an \(L_2\) penalty and setting weight_decay are the same thing.” For plain SGD they are equivalent. For Adam they are not, and the difference is large enough to matter.

Adding \(\frac\lambda2\|\theta\|^2\) to the loss puts \(\lambda\theta\) into the gradient, which then passes through Adam’s adaptive rescaling \(\hat m/(\sqrt{\hat v}+\epsilon)\). Parameters with historically large gradients get a smaller effective decay, and parameters with small gradients get a larger one, so the regularization strength varies per parameter in a way nobody intended.

AdamW (Loshchilov & Hutter, 2019) decouples the decay, applying it directly to the weights after the adaptive step: \[\theta_{t+1}=\theta_t-\eta\left(\frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon}+\lambda\theta_t\right)\] so every parameter decays at the same rate. This is the default for Transformer training, and it is a genuine algorithmic difference, not a reparameterization.

# L2-in-the-loss under Adam vs decoupled decay under AdamW, on identical setups
compare_decay <- function(method, lambda = 0.05, steps = 400, seed = 61) {
  torch_manual_seed(seed); set.seed(seed)
  X <- torch_randn(256, 20); Y <- torch_randn(256, 1)
  net <- nn_sequential(nn_linear(20, 64), nn_relu(), nn_linear(64, 1))
  opt <- switch(method,
    adam_l2   = optim_adam(net$parameters, lr = 0.01),
    adam_wd   = optim_adam(net$parameters, lr = 0.01, weight_decay = lambda),
    adamw     = optim_adamw(net$parameters, lr = 0.01, weight_decay = lambda))
  for (s in seq_len(steps)) {
    opt$zero_grad()
    l <- nnf_mse_loss(net(X), Y)
    if (method == "adam_l2") {
      pen <- Reduce(`+`, lapply(net$parameters, \(p) torch_sum(p^2)))
      l <- l + (lambda / 2) * pen
    }
    l$backward(); opt$step()
  }
  norms <- vapply(net$parameters, \(p) as.numeric(torch_norm(p)), numeric(1))
  c(final_mse = as.numeric(with_no_grad(nnf_mse_loss(net(X), Y))),
    total_weight_norm = sqrt(sum(norms^2)))
}
as.data.frame(do.call(rbind, lapply(
  c("adam_l2", "adam_wd", "adamw"), \(m) c(method = m, signif(compare_decay(m), 4)))))

The three rows use the same \(\lambda\), the same learning rate, the same seed, and the same architecture, and reach different weight norms, because the decay is being applied at different points in the update.

14.12 Learning-rate schedules

The learning rate is the most consequential hyperparameter, and holding it constant is rarely optimal.

\[ \begin{aligned} \textbf{Step decay: }&\ \eta_t=\eta_0\gamma^{\lfloor t/s\rfloor}\\ \textbf{Cosine annealing: }&\ \eta_t=\eta_{\min}+\tfrac12(\eta_0-\eta_{\min})\big(1+\cos(\pi t/T)\big)\\ \textbf{Linear warmup: }&\ \eta_t=\eta_0\cdot\min\!\big(1,\ t/T_{\text{warm}}\big) \end{aligned} \]

Warmup exists because of adaptive optimizers, not because of the model. Adam’s second-moment estimate \(\hat v_t\) is built from very few samples in the first steps, so \(1/\sqrt{\hat v_t}\) is both large and badly estimated, the effective step size early in training has enormous variance. Ramping \(\eta\) from near zero lets the moment estimates stabilize before full-size steps are taken. It is essential for Transformers, and largely unnecessary for SGD+momentum.

Tt <- 1000; eta0 <- 1e-3
sched <- data.frame(t = 1:Tt) |>
  mutate(Constant = eta0,
         `Step decay` = eta0 * 0.5^floor(t / 250),
         `Cosine` = 1e-5 + 0.5 * (eta0 - 1e-5) * (1 + cos(pi * t / Tt)),
         `Warmup + cosine` = ifelse(t < 100, eta0 * t / 100,
           1e-5 + 0.5*(eta0-1e-5)*(1+cos(pi*(t-100)/(Tt-100))))) |>
  pivot_longer(-t, names_to = "schedule", values_to = "lr")

ggplot(sched, aes(t, lr, color = schedule)) +
  geom_line(linewidth = 0.9) +
  scale_color_viridis_d(option = "turbo", end = 0.9) +
  labs(title = "Learning-rate schedules",
       subtitle = "Warmup ramps up so adaptive moment estimates can stabilize; cosine anneals for a fine final approach",
       x = "Step", y = "Learning rate", color = NULL) +
  theme_dspa()

# The learning rate and batch size interact; neither can be tuned alone
set.seed(71); torch_manual_seed(71)
n_lb <- 2048; p_lb <- 20
Xlb <- torch_randn(n_lb, p_lb)
beta_lb <- torch_randn(p_lb, 1)
Ylb <- Xlb$matmul(beta_lb) + torch_randn(n_lb, 1) * 0.5
Xval <- torch_randn(512, p_lb); Yval <- Xval$matmul(beta_lb) + torch_randn(512, 1) * 0.5

train_lb <- function(lr, bs, epochs = 12, seed = 73) {
  torch_manual_seed(seed)
  net <- nn_sequential(nn_linear(p_lb, 64), nn_relu(), nn_linear(64, 1))
  opt <- optim_sgd(net$parameters, lr = lr, momentum = 0.9)
  for (e in seq_len(epochs)) {
    perm <- sample(n_lb)
    for (s in seq(1, n_lb - bs + 1, by = bs)) {
      idx <- perm[s:(s + bs - 1)]
      opt$zero_grad()
      l <- nnf_mse_loss(net(Xlb[idx, ]), Ylb[idx, ])
      l$backward(); opt$step()
    }
  }
  v <- as.numeric(with_no_grad(nnf_mse_loss(net(Xval), Yval)))
  if (!is.finite(v)) 1e3 else min(v, 1e3)
}

lrs <- 10^seq(-3, 0, length.out = 12)
bss <- c(8, 16, 32, 64, 128, 256)
Zlb <- outer(bss, lrs, Vectorize(function(b, l) log10(train_lb(l, b))))

plot_ly(x = lrs, y = bss, z = Zlb, type = "surface",
        colorscale = "Inferno", reversescale = TRUE,
        colorbar = list(title = "log10\nval MSE")) |>
  layout(title = "Validation loss over learning rate and batch size",
         scene = list(xaxis = list(title = "Learning rate", type = "log"),
                      yaxis = list(title = "Batch size", type = "log"),
                      zaxis = list(title = "log10 validation MSE")))

The good region is a diagonal valley, not a rectangle: larger batches tolerate, and require, larger learning rates, because the gradient estimate is less noisy. Tuning \(\eta\) at one batch size and then changing the batch size moves you off the valley floor. The common heuristic is to scale \(\eta\) linearly with batch size, which is the direction this surface’s valley runs.

14.13 Regularization

\[ \begin{aligned} \textbf{Weight decay: }&\ \text{shrink }\theta\text{ toward }0\ \text{— see \S14.11.2 for the Adam caveat}\\ \textbf{Dropout: }&\ \text{zero each unit independently with probability }p\ \text{at \emph{training} time}\\ \textbf{Early stopping: }&\ \text{halt when validation loss stops improving}\\ \textbf{Augmentation: }&\ \text{apply label-preserving transformations to the inputs} \end{aligned} \]

Dropout behaves differently at training and test time, and the scaling is easy to get wrong. During training, units are zeroed with probability \(p\) and the survivors are scaled by \(1/(1-p)\) (“inverted dropout”) so the expected activation is unchanged. At test time dropout is off and no scaling is applied.

Forgetting to switch modes is a classic bug with a distinctive signature: validation loss higher than training loss and unstable across evaluations, because the model is still dropping units while being scored. In torch the switch is model$eval() and model$train(); in a hand-written loop it is easy to omit.

Conceptually dropout approximates training an exponentially large ensemble of subnetworks with shared weights, and averaging them at test time. That is also why it interacts poorly with BatchNorm, the two disagree about what the activation statistics should be, and why modern architectures often use one or the other rather than both.

torch_manual_seed(81)
drop <- nn_dropout(p = 0.5)
v <- torch_ones(1, 10)

drop$train()
train_out <- as.numeric(drop(v))
drop$eval()
eval_out <- as.numeric(drop(v))

data.frame(
  mode = c("train()", "eval()"),
  n_zeroed = c(sum(train_out == 0), sum(eval_out == 0)),
  value_of_survivors = c(unique(train_out[train_out != 0])[1], unique(eval_out)[1]),
  mean_output = c(mean(train_out), mean(eval_out)),
  note = c("survivors scaled by 1/(1-p) = 2", "identity: no dropping, no scaling"))

# A deliberately overparameterized network on a small sample
set.seed(91); torch_manual_seed(91)
n_r <- 200; p_r <- 30
Xr <- torch_randn(n_r, p_r)
br <- torch_zeros(p_r, 1); br[1:5, 1] <- 2
Yr <- Xr$matmul(br) + torch_randn(n_r, 1)
Xrv <- torch_randn(1000, p_r); Yrv <- Xrv$matmul(br) + torch_randn(1000, 1)

fit_reg <- function(kind, epochs = 400, seed = 93) {
  torch_manual_seed(seed)
  net <- if (kind == "dropout")
    nn_sequential(nn_linear(p_r, 512), nn_relu(), nn_dropout(0.5),
                  nn_linear(512, 512), nn_relu(), nn_dropout(0.5), nn_linear(512, 1))
  else
    nn_sequential(nn_linear(p_r, 512), nn_relu(),
                  nn_linear(512, 512), nn_relu(), nn_linear(512, 1))
  wd <- if (kind == "weight decay") 0.05 else 0
  opt <- optim_adamw(net$parameters, lr = 1e-3, weight_decay = wd)
  tr <- va <- numeric(epochs)
  for (e in seq_len(epochs)) {
    net$train(); opt$zero_grad()
    l <- nnf_mse_loss(net(Xr), Yr); l$backward(); opt$step()
    net$eval()
    with_no_grad({
      tr[e] <- as.numeric(nnf_mse_loss(net(Xr), Yr))
      va[e] <- as.numeric(nnf_mse_loss(net(Xrv), Yrv))
    })
  }
  data.frame(epoch = 1:epochs, train = tr, valid = va, method = kind)
}
regs <- bind_rows(lapply(c("none", "weight decay", "dropout"), fit_reg))

regs |> pivot_longer(c(train, valid), names_to = "split", values_to = "mse") |>
  ggplot(aes(epoch, mse, color = split)) +
  geom_line(linewidth = 0.8) +
  facet_wrap(~ method) +
  scale_y_log10() +
  scale_color_manual(values = c(train = "#7FB069", valid = "#D8433B")) +
  labs(title = "Regularization on an overparameterized network (786k parameters, 200 observations)",
       subtitle = "Unregularized training error goes to zero while validation error rises -- the gap is what regularization closes",
       x = "Epoch", y = "MSE (log scale)", color = NULL) +
  theme_dspa(10)

regs |> summarise(best_valid = min(valid), final_train = last(train), .by = method) |>
  mutate(across(where(is.numeric), \(z) signif(z, 4)))

# --- Interactive equivalent ------------------------------------------------
d <- filter(regs, method == "none")
plot_ly(d, x = ~epoch, y = ~train, type = "scatter", mode = "lines",
        name = "Training") |>
  add_lines(y = ~valid, name = "Validation") |>
  layout(title = "Unregularized training on a small sample",
         xaxis = list(title = "Epoch"),
         yaxis = list(title = "MSE", type = "log"))

Early stopping is regularization, and it must use a genuine validation split. Choosing the stopping epoch by watching the test set turns the test estimate into a selection-optimized one, the error of Chapter 9, §9.16, applied to a training curve. Split three ways: train, validation (for stopping and hyperparameters), test (touched once).

---

14.14 The training loop

14.14.1 Samples, batches, and epochs

Three words with precise meanings that are routinely confused.

Term	Definition	Governs
Sample	One element of the dataset — a patient, image, or record	The unit of the loss
Batch	A subset of \(b\) samples processed together to produce one parameter update	Gradient variance, memory, hardware utilization
Epoch	One complete pass over the training set — \(\lceil N/b\rceil\) updates	The unit in which training length is reported
Step / iteration	One parameter update, i.e. one batch	What a learning-rate schedule is indexed by

Common misconception: “larger batches give better fits and cost more compute.” Both halves are wrong in the way that matters.

A larger batch gives a lower-variance gradient estimate, the variance falls as \(\sigma^2/b\) (Chapter 13, §13.9) — but a lower-variance gradient is not a better fit. Large-batch training is well documented to generalize worse at a fixed epoch budget, because the gradient noise that small batches inject acts as an implicit regularizer and steers the trajectory toward flatter regions (Keskar et al., 2017). Compensating requires scaling the learning rate with the batch size (§14.12), the diagonal valley in that section’s surface.

And the cost is memory, not total compute. Per epoch, a larger batch does the same number of FLOPs in fewer, larger matrix multiplications, which is usually faster because it utilizes the hardware better. What grows linearly with \(b\) is the activation memory (§14.3), which is why “reduce the batch size” is the standard response to an out-of-memory error.

# Gradient variance falls as 1/b; per-epoch wall time usually falls too
grad_variance <- function(b, reps = 40, seed = 151) {
  torch_manual_seed(seed); set.seed(seed)
  net <- nn_sequential(nn_linear(20, 32), nn_relu(), nn_linear(32, 1))
  X <- torch_randn(4096, 20); Y <- torch_randn(4096, 1)
  gs <- vapply(seq_len(reps), function(r) {
    idx <- sample(4096, b)
    net$zero_grad()
    nnf_mse_loss(net(X[idx, ]), Y[idx, ])$backward()
    as.numeric(torch_norm(net[[1]]$weight$grad))
  }, numeric(1))
  t_epoch <- system.time({
    for (s in seq(1, 4096 - b + 1, by = b)) {
      net$zero_grad()
      nnf_mse_loss(net(X[s:(s+b-1), ]), Y[s:(s+b-1), ])$backward()
    }})[["elapsed"]]
  c(batch = b, sd_of_gradient_norm = sd(gs), seconds_per_epoch = t_epoch,
    activation_MB = 4 * b * 33 / 1e6)
}
as.data.frame(do.call(rbind, lapply(c(8, 32, 128, 512), grad_variance))) |>
  mutate(across(everything(), \(z) signif(z, 4)))

Gradient variability falls with \(b\) as predicted; wall time per epoch falls too, because larger matrix multiplications use the hardware better; and activation memory rises linearly. Those are three different quantities moving in three different directions, and conflating them is what produces the misconception.

14.14.2 The loop itself

Every deep-learning training loop has the same six steps. Writing it once by hand makes the framework versions legible.

train_manual <- function(net, X, Y, Xv, Yv, epochs = 200, lr = 1e-3,
                         batch = 64, seed = 101) {
  torch_manual_seed(seed); set.seed(seed)
  opt <- optim_adamw(net$parameters, lr = lr, weight_decay = 1e-4)
  n <- X$size(1)
  hist <- data.frame(epoch = integer(), train = numeric(), valid = numeric())
  for (e in seq_len(epochs)) {
    net$train()                                        # (0) training mode
    perm <- sample(n); tot <- 0
    for (s in seq(1, n - batch + 1, by = batch)) {
      idx <- perm[s:(s + batch - 1)]
      opt$zero_grad()                                  # (1) clear accumulated grads
      pred <- net(X[idx, ])                            # (2) forward
      loss <- nnf_mse_loss(pred, Y[idx, ])             # (3) loss
      loss$backward()                                  # (4) backward
      opt$step()                                       # (5) update
      tot <- tot + as.numeric(loss) * length(idx)
    }
    net$eval()                                         # (6) evaluation mode
    v <- with_no_grad(as.numeric(nnf_mse_loss(net(Xv), Yv)))
    hist <- rbind(hist, data.frame(epoch = e, train = tot / n, valid = v))
  }
  list(net = net, history = hist)
}

The numbered comments are the whole algorithm. Step (1) is the one most often omitted, §14.2 explains why its absence looks like a learning-rate problem. Steps (0) and (6) are the dropout and BatchNorm mode switches of §14.13.

# --- The same loop with luz, which supplies callbacks, metrics, and devices --
library(luz)
fitted <- nn_sequential(nn_linear(p, 64), nn_relu(),
                        nn_dropout(0.2), nn_linear(64, 1)) |>
  setup(loss = nnf_mse_loss, optimizer = optim_adamw,
        metrics = list(luz_metric_rmse())) |>
  set_opt_hparams(lr = 1e-3, weight_decay = 1e-4) |>
  fit(train_dl, epochs = 200, valid_data = valid_dl,
      callbacks = list(
        luz_callback_early_stopping(monitor = "valid_loss", patience = 20),
        luz_callback_lr_scheduler(lr_one_cycle, max_lr = 1e-2,
                                  epochs = 200, steps_per_epoch = length(train_dl))),
      verbose = FALSE)                                  # verbose = FALSE keeps output small

---

PART IV: A COMPLETE WORKED EXAMPLE

14.15 Case study: clinical tabular prediction

The point of this section is not that a neural network can be fitted to clinical data. It is whether it should be.

Common misconception: “deep learning is the strongest method available, so use it when accuracy matters.” On tabular data of moderate size, gradient-boosted trees usually match or beat neural networks, train in seconds rather than minutes, need almost no tuning, and handle mixed types and missingness natively (Grinsztajn et al., 2022).

Deep learning’s decisive advantages are in domains with structure a network can exploit, spatial locality in images (Part 2), sequential dependence in text and signals (Part 3), or where the sample is large enough for representation learning to pay for itself. A tabular clinical dataset with a few thousand rows is usually neither.

The only way to know is to run the baselines, on the same split, with the same metric. A neural network reported without a regularized logistic regression and a boosted-tree comparison is an uninterpretable number.

qol <- dspa_read("https://umich.instructure.com/files/481332/download?download_frd=1",
                 "Case06_QoL_Symptom_ChronicIllness.csv")

if (is.null(qol)) {
  set.seed(111); n_s <- 2000; p_s <- 24
  Xs <- matrix(rnorm(n_s * p_s), n_s, p_s)
  eta <- Xs[,1] * 0.9 - Xs[,2] * 0.7 + 0.5 * Xs[,3] * Xs[,4] + rnorm(n_s, sd = 0.8)
  qol <- as.data.frame(round(Xs * 2 + 5))
  names(qol) <- paste0("q", seq_len(p_s))
  qol$CHRONICDISEASESCORE <- as.numeric(eta > median(eta)) + rnorm(n_s, 1.5, 0.3)
  qol$CHARLSONSCORE <- rpois(n_s, 2)
  message("Note: QoL file unavailable; using a synthetic substitute.")
}

qol <- qol |> filter(CHRONICDISEASESCORE != -9, CHARLSONSCORE != -9)
cut_cds <- median(qol$CHRONICDISEASESCORE)
qol$cd <- factor(qol$CHRONICDISEASESCORE > cut_cds,
                 levels = c(FALSE, TRUE), labels = c("minor", "severe"))

# The feature set is defined ONCE, by name. CHRONICDISEASESCORE determines the
# outcome exactly and must be excluded (Chapter 11, Section 11.3).
DROP <- c("ID", "CHRONICDISEASESCORE", "INTERVIEWDATE", "cd")
FEATURES <- setdiff(names(qol), DROP)
FEATURES <- FEATURES[vapply(qol[FEATURES], \(v) is.numeric(v) && var(v, na.rm = TRUE) > 0,
                            logical(1))]
dat <- qol[complete.cases(qol[, c(FEATURES, "cd")]), c(FEATURES, "cd")]

c(observations = nrow(dat), features = length(FEATURES),
  positive_rate = round(mean(dat$cd == "severe"), 4))

#>  observations      features positive_rate 
#>     2190.0000       38.0000        0.4991

library(rsample)
set.seed(1234)
# THREE-way split: train fits, validation selects, test is touched once
sp1 <- initial_split(dat, prop = 0.7, strata = cd)
train_raw <- training(sp1)
sp2 <- initial_split(testing(sp1), prop = 0.5, strata = cd)
valid_raw <- training(sp2); test_raw <- testing(sp2)

# Standardization uses TRAINING statistics only (Chapter 5, Section 5.4.1)
mu <- vapply(train_raw[FEATURES], mean, numeric(1))
sdv <- vapply(train_raw[FEATURES], sd, numeric(1)); sdv[sdv == 0] <- 1
scale_with <- function(d) as.matrix(sweep(sweep(as.matrix(d[FEATURES]), 2, mu), 2, sdv, "/"))

Xtr <- scale_with(train_raw); ytr <- as.integer(train_raw$cd == "severe")
Xva <- scale_with(valid_raw); yva <- as.integer(valid_raw$cd == "severe")
Xte <- scale_with(test_raw);  yte <- as.integer(test_raw$cd  == "severe")

c(train = nrow(Xtr), validation = nrow(Xva), test = nrow(Xte),
  no_information_rate = round(max(mean(yte), 1 - mean(yte)), 4))

#>               train          validation                test no_information_rate 
#>           1532.0000            329.0000            329.0000              0.5015

14.15.1 The baselines, first

library(glmnet); library(ranger)

set.seed(121)
cv_lasso <- cv.glmnet(Xtr, ytr, family = "binomial", alpha = 1, nfolds = 10)
p_lasso <- as.numeric(predict(cv_lasso, newx = Xte, s = "lambda.1se", type = "response"))

rf <- ranger(x = Xtr, y = factor(ytr), probability = TRUE, num.trees = 500,
             seed = 121)
p_rf <- predict(rf, Xte)$predictions[, "1"]

gbm_ok <- requireNamespace("xgboost", quietly = TRUE)
p_gbm <- if (gbm_ok) {
  set.seed(121)
  
  # Create explicit DMatrix objects for training and validation
  dtrain <- xgboost::xgb.DMatrix(data = Xtr, label = ytr)
  dval   <- xgboost::xgb.DMatrix(data = Xva, label = yva)
  
  bst <- xgboost::xgb.train(
    params = list(
      objective = "binary:logistic",
      max_depth = 4,
      eta = 0.05,
      subsample = 0.8,
      colsample_bytree = 0.8
    ),
    data = dtrain,
    nrounds = 300,
    watchlist = list(val = dval),  # <-- Fixed from 'evals' to 'watchlist'
    early_stopping_rounds = 30,
    verbose = 0
  )
  
  predict(bst, Xte)
} else {
  rep(NA_real_, length(yte))
}

14.15.2 The network

torch_manual_seed(131); set.seed(131)

# Xtr is already numeric, so it defaults to Float. 
# ytr and yva are integers, so we must force them to Float32 for BCE loss.
Xtr_t <- torch_tensor(Xtr)
ytr_t <- torch_tensor(matrix(ytr, ncol = 1), dtype = torch_float32())

Xva_t <- torch_tensor(Xva)
yva_t <- torch_tensor(matrix(yva, ncol = 1), dtype = torch_float32())

Xte_t <- torch_tensor(Xte)

mlp <- nn_sequential(
  nn_linear(ncol(Xtr), 128), nn_batch_norm1d(128), nn_relu(), nn_dropout(0.3),
  nn_linear(128, 64),        nn_batch_norm1d(64),  nn_relu(), nn_dropout(0.3),
  nn_linear(64, 1))

opt <- optim_adamw(mlp$parameters, lr = 3e-3, weight_decay = 1e-3)
epochs <- 300; bs <- 64; n_tr <- nrow(Xtr)
hist <- data.frame(epoch = integer(), train = numeric(), valid = numeric())
best_val <- Inf; best_state <- NULL; patience <- 40; since <- 0

for (e in seq_len(epochs)) {
  mlp$train(); perm <- sample(n_tr); tot <- 0
  for (s in seq(1, n_tr - bs + 1, by = bs)) {
    idx <- perm[s:(s + bs - 1)]
    opt$zero_grad()
    l <- nnf_binary_cross_entropy_with_logits(mlp(Xtr_t[idx, ]), ytr_t[idx, ])
    l$backward(); opt$step()
    tot <- tot + as.numeric(l) * length(idx)
  }
  mlp$eval()
  v <- with_no_grad(as.numeric(
    nnf_binary_cross_entropy_with_logits(mlp(Xva_t), yva_t)))
  hist <- rbind(hist, data.frame(epoch = e, train = tot / n_tr, valid = v))
  
  # Early stopping on the VALIDATION set; the test set is never consulted
  if (v < best_val - 1e-5) {
    best_val <- v; since <- 0
    best_state <- lapply(mlp$state_dict(), \(t) t$clone())
  } else if ((since <- since + 1) >= patience) break
}
mlp$load_state_dict(best_state); mlp$eval()

c(epochs_run = nrow(hist), best_epoch = which.min(hist$valid),
  best_validation_loss = round(best_val, 4))

#>           epochs_run           best_epoch best_validation_loss 
#>               47.000                7.000                0.659

ggplot(pivot_longer(hist, c(train, valid), names_to = "split", values_to = "loss"),
       aes(epoch, loss, color = split)) +
  geom_vline(xintercept = which.min(hist$valid), linetype = "dashed",
             color = "grey35") +
  geom_line(linewidth = 0.8) +
  scale_color_manual(values = c(train = "#7FB069", valid = "#D8433B")) +
  labs(title = "Training curve with early stopping",
       subtitle = "Dashed: the selected epoch, chosen on validation loss. The test set plays no part in this decision",
       x = "Epoch", y = "Binary cross-entropy", color = NULL) +
  theme_dspa()

14.15.3 The comparison

library(pROC)
p_mlp <- as.numeric(with_no_grad(torch_sigmoid(mlp(Xte_t))))

score <- function(p, label) {
  if (all(is.na(p))) return(NULL)
  auc <- as.numeric(pROC::auc(pROC::roc(yte, p, quiet = TRUE)))
  brier <- mean((p - yte)^2)
  pc <- pmin(pmax(p, 1e-6), 1 - 1e-6)
  slope <- unname(coef(glm(yte ~ qlogis(pc), family = binomial()))[2])
  data.frame(model = label, AUC = round(auc, 4), Brier = round(brier, 4),
             calibration_slope = round(slope, 3),
             accuracy = round(mean((p > 0.5) == yte), 4))
}

results <- bind_rows(
  data.frame(model = "Majority class", AUC = 0.5,
             Brier = round(mean((mean(ytr) - yte)^2), 4),
             calibration_slope = NA,
             accuracy = round(max(mean(yte), 1 - mean(yte)), 4)),
  score(p_lasso, "LASSO logistic"),
  score(p_rf, "Random forest"),
  if (gbm_ok) score(p_gbm, "Gradient boosting") else NULL,
  score(p_mlp, "MLP (this section)"))
results |> arrange(desc(AUC))

Read the AUC column against the Brier score and the calibration slope together, as Chapter 9, §9.8 requires. AUC measures ranking only; a model can lead on AUC and produce probabilities that cannot be believed. A calibration slope below 1 means the predictions are too extreme, which is the characteristic failure of overparameterized networks and the subject of Part 5, §14.71.

# The cost side of the comparison, which accuracy tables omit
n_par_mlp <- sum(vapply(mlp$parameters, \(p) prod(dim(p)), numeric(1)))
data.frame(
  model = c("LASSO logistic", "Random forest", "MLP"),
  parameters_or_nodes = c(ncol(Xtr) + 1, "500 trees", n_par_mlp),
  tuning_burden = c("1 (lambda, by CV)", "2-3 (mtry, depth)",
                    "8+ (width, depth, lr, batch, dropout, decay, schedule, epochs)"),
  interpretable = c("yes: signed coefficients", "partially: importances",
                    "no, without post-hoc methods"))

14.16 Diagnostics

Training failures have distinctive signatures. Reading them saves hours.

Symptom	Likely cause	Check
Loss is NaN after a few steps	Learning rate too large; exploding gradients; \(\log(0)\)	Reduce \(\eta\); clip gradients; use *_with_logits losses
Loss flat from step 1	Dead ReLUs, vanishing gradient, or \(\eta\) far too small	Gradient norms per layer (§14.8)
Loss decreases then diverges	zero_grad() missing, or \(\eta>2/L\)	Inspect the loop; halve \(\eta\)
Training loss falls, validation rises	Overfitting	Regularize; early stop (§14.13)
Validation below training loss	Dropout still on at evaluation	model$eval()
Cannot overfit 10 examples	A bug, not a modelling problem	Overfit-a-batch test
Results change across identical runs	Unseeded, or GPU nondeterminism	§14.19

The single most useful diagnostic: try to overfit a tiny batch. Take ten examples, remove all regularization, and train until the loss reaches essentially zero. A correctly wired network must be able to memorize ten points. If it cannot, the problem is a bug, a shape mismatch, a detached graph, a missing zero_grad(), a loss applied to the wrong tensor, and no amount of hyperparameter tuning will fix it. This test takes seconds and separates implementation errors from modelling questions.

torch_manual_seed(141)
Xtiny <- Xtr_t[1:10, ]; ytiny <- ytr_t[1:10, ]
tiny <- nn_sequential(nn_linear(ncol(Xtr), 64), nn_relu(), nn_linear(64, 1))
opt_t <- optim_adam(tiny$parameters, lr = 0.05)
for (s in 1:500) {
  opt_t$zero_grad()
  l <- nnf_binary_cross_entropy_with_logits(tiny(Xtiny), ytiny)
  l$backward(); opt_t$step()
}
c(final_loss_on_10_examples = signif(as.numeric(l), 3),
  accuracy = mean((as.numeric(torch_sigmoid(tiny(Xtiny))) > 0.5) == as.numeric(ytiny)),
  verdict = "wiring is correct")

#> final_loss_on_10_examples                  accuracy                   verdict 
#>                       "0"                       "1"       "wiring is correct"

# Per-layer gradient norms: the diagnostic for a flat loss
torch_manual_seed(143)
probe <- nn_sequential(
  nn_linear(ncol(Xtr), 128), nn_sigmoid(),     # sigmoid chosen to show the problem
  nn_linear(128, 128), nn_sigmoid(),
  nn_linear(128, 128), nn_sigmoid(),
  nn_linear(128, 128), nn_sigmoid(),
  nn_linear(128, 1))
l <- nnf_binary_cross_entropy_with_logits(probe(Xtr_t), ytr_t)
l$backward()

gn <- data.frame(
  layer = seq_along(Filter(\(m) inherits(m, "nn_linear"), probe$modules)),
  grad_norm = vapply(Filter(\(m) inherits(m, "nn_linear"), probe$modules),
                     \(m) as.numeric(torch_norm(m$weight$grad)), numeric(1)))
gn$grad_norm <- gn$grad_norm[!is.na(gn$grad_norm)][seq_len(nrow(gn))]

ggplot(gn, aes(layer, pmax(grad_norm, 1e-16))) +
  geom_col(fill = "steelblue", width = 0.6) +
  scale_y_log10() +
  scale_x_continuous(breaks = gn$layer) +
  labs(title = "Gradient norm by layer, sigmoid activations",
       subtitle = "The signal decays by orders of magnitude toward the input -- early layers receive almost nothing",
       x = "Linear layer (1 = closest to input)", y = "Gradient norm (log scale)") +
  theme_dspa()

---

PART V: SYNTHESIS

14.17 The loss landscape

Neural network objectives are non-convex (Chapter 13, §13.14), and it is worth seeing what that means here rather than in a two-dimensional toy.

A trained network’s landscape can be visualized on a random two-dimensional slice through parameter space. Following Li et al. (2018), the two directions are filter-normalized, each random direction is rescaled layer by layer to match the corresponding weight’s norm, because without that normalization the apparent flatness is an artifact of scale invariance rather than a property of the loss.

filter_normalized_dir <- function(params, seed) {
  torch_manual_seed(seed)
  lapply(params, function(p) {
    d <- torch_randn_like(p)
    d * (torch_norm(p) / (torch_norm(d) + 1e-10))
  })
}

landscape <- function(net, X, Y, span = 1.0, n_grid = 21) {
  base <- lapply(net$parameters, \(p) p$clone()$detach())
  d1 <- filter_normalized_dir(base, 201)
  d2 <- filter_normalized_dir(base, 202)
  alphas <- seq(-span, span, length.out = n_grid)
  Z <- matrix(NA_real_, n_grid, n_grid)
  ps <- net$parameters
  for (i in seq_along(alphas)) for (j in seq_along(alphas)) {
    with_no_grad({
      for (k in seq_along(ps))
        ps[[k]]$copy_(base[[k]] + alphas[i] * d1[[k]] + alphas[j] * d2[[k]])
      Z[i, j] <- as.numeric(nnf_binary_cross_entropy_with_logits(net(X), Y))
    })
  }
  with_no_grad(for (k in seq_along(ps)) ps[[k]]$copy_(base[[k]]))   # restore
  list(alphas = alphas, Z = Z)
}

ls_res <- landscape(mlp, Xva_t, yva_t, span = 0.9, n_grid = 21)
c(loss_at_minimum = signif(ls_res$Z[11, 11], 4),
  loss_at_corner = signif(ls_res$Z[1, 1], 4),
  ratio = signif(ls_res$Z[1, 1] / ls_res$Z[11, 11], 3))

#> loss_at_minimum  loss_at_corner           ratio 
#>           0.659           2.015           3.060

plot_ly(x = ls_res$alphas, y = ls_res$alphas, z = ls_res$Z, type = "surface",
        colorscale = "Viridis",
        colorbar = list(title = "Validation\nloss")) |>
  add_trace(x = 0, y = 0, z = ls_res$Z[11, 11], type = "scatter3d",
            mode = "markers", name = "Trained solution",
            marker = list(size = 6, color = "red")) |>
  layout(title = "Filter-normalized loss surface around the trained solution",
         scene = list(xaxis = list(title = "Direction 1"),
                      yaxis = list(title = "Direction 2"),
                      zaxis = list(title = "Validation loss")))

The red marker is the solution early stopping selected. Rotate to see the basin: a wide, relatively flat region rather than a sharp needle. The relationship between basin flatness and generalization is contested, the measure is not reparameterization-invariant, which is exactly why the filter normalization is required, but the qualitative picture is that trained networks sit in broad valleys, not isolated pits.

14.18 Complexity summary

\(B\) = batch size, \(d_\ell\) = width of layer \(\ell\), \(L\) = depth, \(P\) = parameter count, \(E\) = epochs, \(N\) = training-set size.

Quantity	Cost	Note
Forward pass, one linear layer	\(O(Bd_{\text{in}}d_{\text{out}})\)	\(\approx 2Bd_{\text{in}}d_{\text{out}}\) FLOPs
Backward pass	\(\approx 2\times\) forward	Gradients w.r.t. both inputs and weights
Full training	\(O\!\big(ENL\bar d^{\,2}\big)\)	Three passes’ worth per step
Parameter memory	\(4P\) bytes (float32)	Stored once
Gradient memory	\(4P\) bytes	Stored once
Adam state	\(8P\) bytes	Two moment estimates — triples parameter memory
Activation memory	\(4B\sum_\ell d_\ell\) bytes	Scales with batch size; usually binding
Gradient checkpointing	\(O(\sqrt L)\) activations	Recompute in the backward pass: time for space
Reverse-mode AD gradient	\(O(1)\) forward passes	Independent of \(P\)
Finite-difference gradient	\(O(P)\) forward passes	Infeasible beyond toy sizes

Four consequences.

Reverse-mode AD is what makes any of this possible. A gradient for \(10^8\) parameters costs a constant multiple of one forward pass; by finite differences it would cost \(2\times10^8\) forward passes. That single asymptotic fact, from Chapter 13, §13.25, is the foundation of the field.

Adam’s memory cost is not incidental. Weights, gradients, and two moment estimates total \(16P\) bytes at float32, a \(10^9\)-parameter model needs 16 GB before a single activation is stored. This is why optimizer sharding and 8-bit optimizer states exist.

Activation memory scales with the batch and parameters do not. Reducing the batch size is the standard response to an out-of-memory error precisely because it is the only term that shrinks.

Depth costs linearly in compute and linearly in activation memory, but its representational return is exponential (§14.7). That asymmetry is why networks got deeper rather than wider.

14.19 Reproducibility

Common misconception: “set a seed and the run is reproducible.” On CPU with a fixed library version, largely true. On GPU it is not, and the reasons are structural rather than fixable by better practice.

cuDNN selects algorithms by benchmarked speed, and the fastest kernel can differ between runs on the same hardware. Many reductions use atomic accumulation, whose order is nondeterministic, and floating-point addition is not associative, so \((a+b)+c\ne a+(b+c)\) at the bit level. Multi-threaded CPU reductions have the same property.

What can be guaranteed: seed torch_manual_seed(), set.seed(), and the data-loader workers; pin library versions; and where determinism matters more than speed, request deterministic algorithms explicitly and accept the slowdown. What should be reported: the seed, the library versions, the hardware, and, most usefully, results across several seeds, since a single-seed result on a stochastic procedure is one draw from a distribution.

run_once <- function(seed) {
  torch_manual_seed(seed); set.seed(seed)
  net <- nn_sequential(nn_linear(ncol(Xtr), 32), nn_relu(), nn_linear(32, 1))
  opt <- optim_adam(net$parameters, lr = 0.01)
  for (s in 1:100) {
    opt$zero_grad()
    l <- nnf_binary_cross_entropy_with_logits(net(Xtr_t), ytr_t)
    l$backward(); opt$step()
  }
  net$eval()
  as.numeric(pROC::auc(pROC::roc(
    yva, as.numeric(with_no_grad(torch_sigmoid(net(Xva_t)))), quiet = TRUE)))
}
same_seed <- c(run_once(7), run_once(7))
diff_seeds <- vapply(1:8, run_once, numeric(1))

c(identical_across_same_seed = isTRUE(all.equal(same_seed[1], same_seed[2])),
  mean_auc_across_8_seeds = round(mean(diff_seeds), 4),
  sd_across_seeds = round(sd(diff_seeds), 4),
  range = paste(round(range(diff_seeds), 4), collapse = " to "))

#> identical_across_same_seed    mean_auc_across_8_seeds 
#>                     "TRUE"                   "0.5783" 
#>            sd_across_seeds                      range 
#>                   "0.0178"          "0.5626 to 0.612"

The spread across seeds is the honest uncertainty in a single reported number. Quoting one run’s AUC without it overstates precision.

---

14.20 Common pitfalls

#	Pitfall	Consequence	Fix
1	Omitting zero_grad()	Gradients accumulate; loss diverges after a few steps	Clear before every backward pass
2	Evaluating without model$eval()	Dropout and BatchNorm still in training mode	Switch modes; validation below training loss is the tell
3	Evaluating without with_no_grad()	Builds a graph; wastes memory and time	Wrap all inference
4	Reading universal approximation as learnability	Existence says nothing about optimization or generalization	Width may be exponential in \(d\)
5	Stacking linear layers without activations	The whole network collapses to one affine map	Nonlinearity between every pair
6	Default or ad-hoc initialization at depth	Activations vanish or explode geometrically	He for ReLU, Xavier for tanh
7	Treating ReLU as gradient-safe	Dying units, permanently	Monitor dead fractions; leaky/GELU/SiLU
8	Explaining BatchNorm by internal covariate shift	The mechanism is loss smoothing	Also: BatchNorm couples examples in a batch
9	BatchNorm with tiny batches	Batch statistics too noisy to normalize with	LayerNorm, GroupNorm, or larger batches
10	weight_decay in Adam read as \(L_2\)	Per-parameter decay varies with gradient history	AdamW decouples it
11	Tuning the learning rate independently of batch size	The good region is a diagonal valley	Scale \(\eta\) with batch size
12	No warmup with an adaptive optimizer	Early moment estimates give huge, erratic steps	Linear warmup over a few hundred steps
13	Early stopping on the test set	Turns the test estimate into a selection-optimized one	Three-way split
14	Reporting a network without baselines	An uninterpretable number	LASSO, random forest, boosting on the same split
15	Assuming deep learning wins on tabular data	Boosted trees usually match or beat it	Run them
16	Reading AUC as sufficient	Ranking says nothing about probabilities	Brier and calibration slope too
17	Hyperparameter tuning without a validation split	Optimism proportional to the search size	Nested protocol
18	Debugging a bug by tuning	A mis-wired network cannot be tuned into working	Overfit ten examples first
19	Ignoring per-layer gradient norms	Vanishing gradients look like slow learning	Log norms by layer
20	float64 for training	Doubles memory and traffic; buys nothing	float32, or mixed precision
21	Assuming a seed gives GPU determinism	cuDNN selection and atomics are nondeterministic	Report seeds, versions, and multi-seed spread
22	Reporting one run	One draw from a distribution	Several seeds; report the spread
23	Printing model objects or training logs into the document	Output balloons past 100 MB	verbose = 0; cache artifacts to disk
24	Reading the biological analogy as mechanism	Backprop has no biological counterpart	A network is a composition of affine maps and nonlinearities
25	Composing sigmoid or softmax with a separate log	NaN on the first backward pass	Fused *_with_logits losses; they take raw logits
26	Passing already-sigmoided values to a with_logits loss	Nonlinearity applied twice; silently wrong	Pass the raw output layer
27	Confusing \(W^{(\ell)}\)’s columns with the next layer	Every shape argument off by one layer	Rows = current units, columns = previous
28	Calling the bias a residual	It is a learned intercept, not a fit property	\(b\) is estimated; \(y-\hat y\) is observed
29	“Larger batches fit better and cost more compute”	Three quantities conflated	Variance falls as \(1/b\); memory rises; per-epoch time usually falls
30	Retired packages (mxnet, darch, deepnet, kerasR)	Unavailable or unmaintained	torch + luz

---

14.21 Practice problems

Problems continue in Part 2; those below use only Part 1 material.

Problem 1: Verify autograd against analytic gradients

Solution

# A two-layer network's gradient, by autograd and by hand
torch_manual_seed(301)
W1 <- torch_randn(4, 3, requires_grad = TRUE)
W2 <- torch_randn(1, 4, requires_grad = TRUE)
xin <- torch_randn(1, 3); ytgt <- torch_tensor(matrix(1.0))

h  <- torch_tanh(xin$matmul(W1$t()))
out <- h$matmul(W2$t())
loss <- ((out - ytgt)^2)$sum()
loss$backward()

# By hand: dL/dW2 = 2(out-y) h ;  dL/dW1 = [2(out-y) W2 * (1-h^2)]^T x
with_no_grad({
  delta <- 2 * (out - ytgt)
  gW2_manual <- delta$t()$matmul(h)
  gh <- delta$matmul(W2) * (1 - h^2)
  gW1_manual <- gh$t()$matmul(xin)
})
c(W2_max_abs_diff = signif(as.numeric(torch_max(torch_abs(W2$grad - gW2_manual))), 3),
  W1_max_abs_diff = signif(as.numeric(torch_max(torch_abs(W1$grad - gW1_manual))), 3))

#> W2_max_abs_diff W1_max_abs_diff 
#>               0               0

They agree to machine precision. Autograd is not an approximation, it applies the chain rule exactly, which is what distinguishes it from the finite differences of Chapter 13, §13.25.

Problem 2: Find the depth at which sigmoid stops training

Solution

depth_test <- function(depth, act, steps = 200, width = 64, seed = 311) {
  torch_manual_seed(seed); set.seed(seed)
  X <- torch_randn(256, width); Y <- torch_randn(256, 1)
  layers <- list()
  for (l in seq_len(depth)) layers <- c(layers, list(nn_linear(width, width), act()))
  layers <- c(layers, list(nn_linear(width, 1)))
  net <- do.call(nn_sequential, layers)
  opt <- optim_sgd(net$parameters, lr = 0.05)
  l0 <- as.numeric(with_no_grad(nnf_mse_loss(net(X), Y)))
  for (s in seq_len(steps)) {
    opt$zero_grad(); l <- nnf_mse_loss(net(X), Y); l$backward(); opt$step()
  }
  1 - as.numeric(with_no_grad(nnf_mse_loss(net(X), Y))) / l0
}
dps <- c(1, 2, 4, 8, 16, 24)
p2 <- bind_rows(lapply(c(Sigmoid = "nn_sigmoid", ReLU = "nn_relu"), \(a)
  data.frame(depth = dps, act = a,
             reduction = vapply(dps, \(d) depth_test(d, get(a)), numeric(1))))) |>
  mutate(act = ifelse(act == "nn_sigmoid", "Sigmoid", "ReLU"))
p2 |> pivot_wider(names_from = act, values_from = reduction) |>
  mutate(across(-depth, \(z) round(z, 4)))

ggplot(p2, aes(depth, reduction, color = act)) +
  geom_hline(yintercept = 0, color = "grey70") +
  geom_line(linewidth = 1) + geom_point(size = 2.4) +
  scale_x_log10(breaks = dps) +
  scale_color_manual(values = c(Sigmoid = "#D8433B", ReLU = "#3B7DD8")) +
  labs(title = "Fractional loss reduction in 200 steps, by depth",
       subtitle = "Identical budget and learning rate. Sigmoid stops making progress well before ReLU does",
       x = "Hidden layers (log scale)", y = "Fraction of initial loss removed",
       color = NULL) +
  theme_dspa()

Sigmoid’s progress collapses by roughly 8 layers, consistent with the \(4^{-L}\) bound of §14.8. ReLU keeps training far deeper.

Problem 3: Confirm the He factor of 2

Solution

# ReLU halves the variance of a symmetric input, so He doubles the weight variance
torch_manual_seed(321)
z <- torch_randn(200000)
c(var_before_relu = round(as.numeric(torch_var(z)), 4),
  var_after_relu = round(as.numeric(torch_var(torch_relu(z))), 4),
  ratio = round(as.numeric(torch_var(torch_relu(z))) / as.numeric(torch_var(z)), 4),
  theoretical_ratio = round((1 - 1/pi) / 2 + 0, 4))

#>   var_before_relu    var_after_relu             ratio theoretical_ratio 
#>            0.9985            0.3405            0.3410            0.3408

# Empirical check of the layer-to-layer variance ratio at each gain
one_layer_ratio <- function(gain, width = 512, seed = 323) {
  torch_manual_seed(seed)
  h <- torch_randn(4096, width)
  W <- torch_randn(width, width) * sqrt(gain / width)
  as.numeric(torch_var(torch_relu(h$matmul(W$t())))) / as.numeric(torch_var(h))
}
data.frame(gain = c(1, 2, 4),
           variance_ratio = round(vapply(c(1, 2, 4), one_layer_ratio, numeric(1)), 4),
           preserves_variance = abs(vapply(c(1,2,4), one_layer_ratio, numeric(1)) - 1) < 0.1)

Gain 2 is the value that returns a ratio of 1. The exact ReLU variance factor is \(\frac{1}{2}\big(1-\frac{1}{\pi}\big)+\frac{1}{2\pi}\) for a standard normal input, which is close enough to \(\tfrac12\) that the He factor of 2 is the right practical choice.

Problem 4: Reproduce the learning-rate/batch-size scaling rule

Solution

best_lr_for_batch <- function(bs) {
  lrs_try <- 10^seq(-3, -0.3, length.out = 12)
  v <- vapply(lrs_try, \(l) train_lb(l, bs), numeric(1))
  lrs_try[which.min(v)]
}
bs_try <- c(8, 16, 32, 64, 128, 256)
p4 <- data.frame(batch = bs_try,
                 best_lr = vapply(bs_try, best_lr_for_batch, numeric(1)))
p4$ratio_to_smallest <- round(p4$best_lr / p4$best_lr[1], 3)
p4$linear_rule <- round(p4$batch / p4$batch[1], 3)
p4

ggplot(p4, aes(batch, best_lr)) +
  geom_line(linewidth = 1, color = "steelblue") + geom_point(size = 2.6) +
  geom_line(aes(y = best_lr[1] * batch / batch[1]), linetype = "dashed",
            color = "firebrick") +
  scale_x_log10(breaks = bs_try) + scale_y_log10() +
  annotate("text", x = 100, y = max(p4$best_lr) * 1.4, size = 3.2,
           color = "firebrick", label = "linear scaling rule") +
  labs(title = "Best learning rate against batch size",
       subtitle = "Dashed: the linear scaling rule. On log-log axes a slope of 1 confirms it",
       x = "Batch size (log scale)", y = "Best learning rate (log scale)") +
  theme_dspa()

The empirical optimum tracks the linear rule reasonably at small and moderate batch sizes and flattens at large ones, the well-documented breakdown of linear scaling past a problem-dependent threshold.

Problem 5: Measure the cost of the overfit-a-batch test

Solution

# A deliberately broken network: the second layer is detached from the graph
broken_fit <- function(detach_it) {
  torch_manual_seed(331)
  L1 <- nn_linear(ncol(Xtr), 32); L2 <- nn_linear(32, 1)
  opt <- optim_adam(c(L1$parameters, L2$parameters), lr = 0.05)
  Xb <- Xtr_t[1:10, ]; yb <- ytr_t[1:10, ]
  for (s in 1:400) {
    opt$zero_grad()
    h <- torch_relu(L1(Xb))
    if (detach_it) h <- h$detach()          # the bug: gradient stops here
    l <- nnf_binary_cross_entropy_with_logits(L2(h), yb)
    l$backward(); opt$step()
  }
  as.numeric(l)
}
data.frame(
  version = c("correct", "first layer detached"),
  final_loss_on_10_examples = signif(c(broken_fit(FALSE), broken_fit(TRUE)), 4),
  can_memorize_10_points = c(broken_fit(FALSE) < 0.01, broken_fit(TRUE) < 0.01))

The broken version cannot drive the loss to zero on ten examples, which no amount of learning-rate tuning would have revealed. The test costs under a second and localizes the fault to the wiring.

Problem 6: Quantify seed variance

Solution

seeds <- 1:15
aucs <- vapply(seeds, run_once, numeric(1))
c(mean = round(mean(aucs), 4), sd = round(sd(aucs), 4),
  min = round(min(aucs), 4), max = round(max(aucs), 4),
  range_width = round(diff(range(aucs)), 4),
  ci95 = paste(round(mean(aucs) + c(-1.96, 1.96) * sd(aucs) / sqrt(length(aucs)), 4),
               collapse = " to "))

#>               mean                 sd                min                max 
#>           "0.5784"           "0.0183"           "0.5625"           "0.6204" 
#>        range_width               ci95 
#>            "0.058" "0.5691 to 0.5876"

ggplot(data.frame(auc = aucs), aes(auc)) +
  geom_histogram(bins = 10, fill = "steelblue", color = "white") +
  geom_vline(xintercept = mean(aucs), color = "firebrick", linewidth = 1) +
  labs(title = sprintf("Validation AUC across %d random seeds", length(seeds)),
       subtitle = "Identical architecture, data, and hyperparameters. Only the initialization and batch order differ",
       x = "AUC", y = "Runs") +
  theme_dspa()

The range across seeds is often comparable to the differences people report between architectures. A single-seed comparison cannot distinguish a real improvement from initialization luck; report the mean and spread over at least five seeds.

---

14.22 Checkpoint

1. Your loss is NaN after ten steps. Name three causes and how to distinguish them.
2. Your validation loss is consistently below your training loss. What is almost certainly wrong?
3. Universal approximation says one hidden layer suffices. Why build deep networks?
4. You switch from SGD to Adam and add weight_decay = 0.01. Is that \(L_2\) regularization?
5. Your 30-layer network’s loss does not move from step one. What do you check first?
6. Your MLP reaches 0.84 AUC on a clinical tabular dataset. Is that good?

Answers

1. Learning rate too large, exploding gradients, or a numerically invalid loss. Distinguish by inspection. If the loss is finite for a few steps and then blows up, \(\eta\) likely exceeds the stability bound \(2/L\) (Chapter 13, §13.5.1) — halve it and see whether the divergence point moves. If gradient norms grow geometrically with depth, it is exploding gradients, clip them, or check the initialization scale (§14.9). If the NaN appears on the very first backward pass, suspect \(\log(0)\) from a manually composed sigmoid-then-cross-entropy: use the fused *_with_logits form, which is numerically stable. Also check for NaN in the inputs, which propagates silently.
2. Dropout and BatchNorm are still in training mode during evaluation. Dropout zeroes half the units while you score, inflating training loss relative to a validation pass that should have them disabled, and the validation number is not comparable to anything. Call model$eval() before scoring and model$train() before the next epoch. Two lesser causes are worth ruling out: a validation set that is easier than the training set (small or badly stratified), and regularization applied to the training loss but not the validation loss, which makes the two numbers different quantities.
3. Because depth buys exponential representational efficiency, and the theorem never claimed otherwise. Universal approximation is an existence result with no bound on the required width, for a general Lipschitz function on \([0,1]^d\) it can be \(\Omega(\varepsilon^{-d})\) units. Depth changes the accounting: ReLU networks realize a number of linear regions that grows exponentially in depth and only polynomially in width, and there are functions requiring exponentially many units at depth \(k-1\) that depth \(k\) handles with polynomially many. The cost is optimization difficulty, which is what initialization, normalization, and residual connections exist to manage.
4. No, that is Adam with \(L_2\), not AdamW, and they differ. Adding \(\frac\lambda2\|\theta\|^2\) to the loss injects \(\lambda\theta\) into the gradient, which then passes through Adam’s adaptive rescaling \(\hat m/(\sqrt{\hat v}+\epsilon)\), so parameters with a history of large gradients receive a smaller effective decay than those with small ones, and the regularization strength varies per parameter in a way nobody chose. AdamW applies the decay directly to the weights after the adaptive step, so every parameter decays equally. For plain SGD the two are genuinely equivalent; for Adam they are not.
5. Per-layer gradient norms. A completely flat loss almost always means the signal is not reaching the early layers. Log \(\|\nabla W^{(\ell)}\|\) for every \(\ell\): a geometric decay toward the input is vanishing gradients, and the remedies are activation choice (§14.8), He initialization (§14.9), normalization (§14.10), and residual connections (Part 2, §14.28). If the norms are healthy but nothing moves, check that zero_grad() is being called, that the optimizer actually received the parameters, and that the learning rate is not absurdly small. And run the overfit-ten-examples test, if the network cannot memorize ten points, the problem is a bug, not a hyperparameter.
6. Unknown, until three other things are reported. First, the baselines: on tabular clinical data a LASSO logistic regression and a gradient-boosted tree frequently match or beat an MLP, train in seconds, and need almost no tuning, an AUC quoted without them is uninterpretable. Second, the calibration: AUC measures ranking only, so report the Brier score and the calibration slope (Chapter 9, §9.8); an overparameterized network is typically over-confident. Third, the variability: 0.84 from one seed is one draw, and the spread across seeds is often as large as the differences between architectures. Also confirm the number came from a test set used once, not from the split that chose the stopping epoch.

---

14.23 Summary

The computational substrate

• A tensor carries a device, a dtype, and a gradient history; the last of these is the mechanism of everything else.
• Reverse-mode AD computes a full gradient at constant cost in the parameter count, the fact that makes deep learning feasible.
• Gradients accumulate; zero_grad() before every backward pass.
• Activation memory scales with the batch and usually binds, and Adam’s two moment estimates triple parameter memory.

Representation

• A network without nonlinearities collapses to a single affine map.
• Universal approximation is existence, not learnability: no bound on width, no claim about optimization or generalization, and not unique to neural networks.
• Depth buys exponentially many linear regions; width buys polynomially many.

Making training work

• Activation derivative bounds determine trainable depth. Sigmoid’s \(\sigma'\le\frac14\) gives \(4^{-L}\); ReLU removes saturation and introduces dead units.
• Initialization decides whether training happens. He (\(2/n_{\text{in}}\)) for ReLU, Xavier for tanh, derived from variance preservation, with the factor of 2 compensating for ReLU’s halving.
• Normalization smooths the loss landscape; the internal-covariate-shift explanation is not supported. BatchNorm couples examples in a batch and fails at small batch sizes; LayerNorm does not.
• AdamW \(\ne\) Adam + \(L_2\). Decoupled decay applies equally to every parameter.
• The learning rate and batch size must be tuned together; the good region is a diagonal valley.
• Warmup exists for the optimizer, not the model, adaptive moment estimates need steps to stabilize.
• Dropout requires mode switching; validation loss below training loss is the signature of forgetting it.

Practice

• Overfit ten examples first. A network that cannot memorize them has a bug, and no tuning will fix a bug.
• Report baselines. On tabular data, boosted trees usually match or beat a network at a fraction of the cost and tuning burden.
• Report calibration, not only AUC, and report seed variance, since one run is one draw.
• GPU training is not bitwise reproducible; say so, and report across seeds.

---

14.24 Where Part 1 leads

Continue with	Content
Part 2: Convolutional Networks and Vision	Convolution as structured sparsity, receptive fields, residual connections and the gradient-flow argument that resolves §14.7’s depth-versus-optimization tension, transfer learning, CIFAR-10, interpretability
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 that closes Chapter 12’s table
Part 4: Generative and Representation Learning	Autoencoders and their exact relation to PCA, the ELBO derived, GANs, diffusion, self-supervised objectives
Part 5: Generalization, Uncertainty, and Practice	Double descent, why classical capacity bounds fail, calibration of deep networks (Chapter 9), ensembles, Bayesian hyperparameter optimization (Chapter 13), pruning and sparsity (Chapter 11)

Earlier chapters this part depended on

• Chapter 6: perceptrons and the backpropagation derivation
• Chapter 9: nested validation, calibration, honest reporting
• Chapter 10: the memory hierarchy and arithmetic intensity
• Chapter 11: regularization and selection bias
• Chapter 12: the vanishing-gradient product
• Chapter 13: gradient methods, adaptive 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] pROC_1.18.5     ranger_0.16.0   glmnet_4.1-8    Matrix_1.6-5   
#>  [5] rsample_1.2.1   luz_0.5.2       torch_0.13.0    plotly_4.12.1  
#>  [9] patchwork_1.3.0 tidyr_1.3.1     dplyr_1.1.4     ggplot2_4.0.1  
#> 
#> loaded via a namespace (and not attached):
#>  [1] shape_1.4.6.1      gtable_0.3.6       xfun_0.52          bslib_0.9.0       
#>  [5] htmlwidgets_1.6.4  processx_3.8.6     lattice_0.22-6     callr_3.7.6       
#>  [9] vctrs_0.6.5        tools_4.3.3        crosstalk_1.2.1    ps_1.9.0          
#> [13] generics_0.1.3     parallel_4.3.3     tibble_3.2.1       pkgconfig_2.0.3   
#> [17] data.table_1.16.4  RColorBrewer_1.1-3 S7_0.2.1           lifecycle_1.0.5   
#> [21] compiler_4.3.3     farver_2.1.2       progress_1.2.3     codetools_0.2-20  
#> [25] htmltools_0.5.8.1  sass_0.4.9         yaml_2.3.10        pillar_1.10.1     
#> [29] furrr_0.3.1        crayon_1.5.3       jquerylib_0.1.4    cachem_1.1.0      
#> [33] iterators_1.0.14   foreach_1.5.2      parallelly_1.37.1  tidyselect_1.2.1  
#> [37] digest_0.6.37      future_1.33.2      purrr_1.0.2        listenv_0.9.1     
#> [41] splines_4.3.3      labeling_0.4.3     fastmap_1.2.0      grid_4.3.3        
#> [45] cli_3.6.3          magrittr_2.0.3     survival_3.7-0     withr_3.0.2       
#> [49] prettyunits_1.2.0  scales_1.4.0       xgboost_1.7.11.1   bit64_4.0.5       
#> [53] rmarkdown_2.31     httr_1.4.7         globals_0.16.3     bit_4.0.5         
#> [57] otel_0.2.0         hms_1.1.3          evaluate_1.0.3     knitr_1.51        
#> [61] viridisLite_0.4.2  rlang_1.1.5        Rcpp_1.0.14        zeallot_0.1.0     
#> [65] glue_1.8.0         coro_1.0.4         rstudioapi_0.18.0  jsonlite_1.8.9    
#> [69] plyr_1.8.9         R6_2.6.1           fs_1.6.5

SOCR Resource | Visitor number