| SOCR ≫ | DSPA ≫ | DSPA3 Topics ≫ |
DSPA3: Data Science and Predictive Analytics (UMich HS650)
Chapter 14, Part 3: Sequence Models - Recurrence and
Attention
Chapter 14, Part 3: Sequence Models - Recurrence and Attention
SOCR/MCAIM (Ivo Dinov)
SOCR/MCAIM (Ivo Dinov)
August 31, 2026
library(ggplot2)
library(dplyr)
library(tidyr)
library(patchwork)
library(plotly)
if (TORCH_OK) library(torch)Chapter 14 is published in five parts
Part Sections Content Part 1: Foundations §14.0–§14.24 Tensors, autograd, the MLP, activations, initialization, normalization, optimizers, the training loop Part 2: Convolutional Networks and Vision §14.25–§14.44 Convolution arithmetic, receptive fields, residual connections, transfer learning, segmentation, interpretability Part 3 (this document) §14.45–§14.64 RNNs, LSTM/GRU, attention derived, Transformers, tokenization, text generation, neural forecasting Part 4: Generative and Representation Learning §14.65–§14.84 Autoencoders, VAEs and the ELBO, GANs, diffusion, U-Net synthesis, self-supervised learning Part 5: Generalization, Uncertainty, and Practice §14.85–§14.104 Double descent, calibration, ensembles, hyperparameter optimization, pruning, robustness Part 3 closes two threads left open earlier in the book: the \(O(T^2d)\)-versus-\(O(Td^2)\) complexity comparison promised in Chapter 12, §12.28, and the neural-versus-classical forecasting question of Chapter 12, §12.29.
How this part uses graphics
Two-dimensional figures use
ggplot2; the equivalentplot_ly()code follows in a chunk markedeval=FALSE, echo=TRUE. Three-dimensional figures are evaluatedplot_ly(), budgeted at six per part.Heavy training is gated.
HEAVY_EVALdefaults toFALSE, so the character-level language model and the forecasting comparison display their code and load cached results rather than training during the build.
Learning objectives
After completing Part 3 you will be able to:
- Characterize the sequence-modelling problem class and distinguish one-to-many, many-to-one, and many-to-many tasks.
- Derive the vanishing-gradient product for a recurrent network and relate it to the spectral radius of the recurrent weight matrix.
- Explain how LSTM and GRU gating creates an additive gradient path, and connect it to the residual connections of Part 2, §14.32.
- State the encoder–decoder bottleneck problem that motivated attention.
- Derive scaled dot-product attention, including why the scaling is \(1/\sqrt{d_k}\) and what breaks without it.
- Explain what multi-head attention buys, given that it does not add parameters.
- Explain why self-attention requires positional encoding, and state the permutation-equivariance property precisely.
- Compare the \(O(T^2d)\) cost of attention with the \(O(Td^2)\) cost of recurrence, and identify where each wins.
- Implement causal masking and autoregressive sampling with temperature, top-\(k\), and nucleus strategies.
- Evaluate a neural forecaster against classical benchmarks using MASE under rolling-origin validation.
Estimated time: 12–16 hours including exercises.
PART I: RECURRENCE
14.45 The sequence-modelling problem class
Parts 1 and 2 assumed a fixed-size input. Sequences break that: the input is an ordered, variable-length collection where order carries meaning.
| Shape | Input → Output | Examples |
|---|---|---|
| One-to-many | fixed → sequence | Image captioning, music generation |
| Many-to-one | sequence → fixed | Sentiment classification, ECG diagnosis |
| Many-to-many, aligned | \(T\) → \(T\) | Part-of-speech tagging, per-timepoint labelling |
| Many-to-many, unaligned | \(T\) → \(T'\) | Translation, summarization, forecasting |
Three properties distinguish sequence data from the tabular and image cases.
Variable length. A model must handle \(T=5\) and \(T=5000\) without changing its parameter count, which forces either recurrence (apply the same cell at each step) or attention (a set operation plus positional information).
Order matters. Shuffling the tokens of a sentence or the samples of a signal destroys the content. This is exactly the exchangeability assumption that Chapter 12, §12.1 identified as failing for temporal data.
Long-range dependence. The token that determines a prediction may lie hundreds of steps back. How far a model can reach is the central design question of this part.
14.46 Recurrent networks
A recurrent network carries a hidden state forward and applies the same parameters at every step:
\[\boxed{\;\mathbf h_t=\sigma\big(W_{hh}\mathbf h_{t-1}+W_{xh}\mathbf x_t+\mathbf b\big),\qquad \hat{\mathbf y}_t=W_{hy}\mathbf h_t+\mathbf c\;}\]
Parameter sharing across time is the exact analogue of a CNN’s sharing across space (Part 2, §14.25): the parameter count is independent of \(T\), and the same feature detector applies wherever the pattern occurs.
torch_manual_seed(11)
# torch's nn_rnn expects (batch, time, feature) when batch_first = TRUE
rnn <- nn_rnn(input_size = 4, hidden_size = 16, num_layers = 1, batch_first = TRUE)
x_seq <- torch_randn(8, 25, 4) # 8 sequences, 25 steps, 4 features
out <- rnn(x_seq)
c(output_shape = paste(dim(out[[1]]), collapse = " x "),
final_hidden_shape = paste(dim(out[[2]]), collapse = " x "),
parameters = sum(vapply(rnn$parameters, \(p) prod(dim(p)), numeric(1))),
note = "parameter count does not depend on sequence length")#> output_shape
#> "8 x 25 x 16"
#> final_hidden_shape
#> "1 x 8 x 16"
#> parameters
#> "352"
#> note
#> "parameter count does not depend on sequence length"# The same cell handles a different length with no change
c(length_25 = paste(dim(rnn(torch_randn(2, 25, 4))[[1]]), collapse = " x "),
length_200 = paste(dim(rnn(torch_randn(2, 200, 4))[[1]]), collapse = " x "))#> length_25 length_200
#> "2 x 25 x 16" "2 x 200 x 16"14.46.1 The vanishing-gradient product
Chapter 12, §12.26.1 introduced this; here is the precise statement.
Backpropagation through time differentiates a loss at step \(T\) with respect to the state at step \(t\), and the chain rule produces a product of Jacobians:
\[\frac{\partial\mathcal L_T}{\partial\mathbf h_t}=\frac{\partial\mathcal L_T}{\partial\mathbf h_T}\prod_{k=t+1}^{T}\frac{\partial\mathbf h_k}{\partial\mathbf h_{k-1}}, \qquad \frac{\partial\mathbf h_k}{\partial\mathbf h_{k-1}}=W_{hh}^\top\operatorname{diag}\big(\sigma'(\mathbf z_k)\big).\]
Bounding the norm of the product,
\[\left\|\prod_{k=t+1}^{T}\frac{\partial\mathbf h_k}{\partial\mathbf h_{k-1}}\right\|\ \le\ \big(\|W_{hh}\|\cdot\gamma\big)^{T-t},\qquad \gamma=\sup|\sigma'|,\]
so the gradient decays or grows geometrically in the lag \(T-t\). The governing quantity is the spectral radius \(\rho(W_{hh})\): below 1 the gradient vanishes, above 1 it explodes.
Common misconception: “gradient clipping fixes the gradient problem.” It fixes one of the two, and not the one that limits what the model can learn.
An exploding gradient announces itself, the loss becomes
NaN, and is fixed by clipping: rescale the gradient whenever its norm exceeds a threshold, which changes the step length and not its direction.A vanishing gradient is silent. The loss decreases, training appears healthy, and the network simply never learns dependencies beyond a few dozen steps because no gradient signal reaches that far. Nothing in the training curve reveals it. You have to measure the gradient by lag, which is what the figure below does.
grad_by_lag <- function(rho, T_len = 60, d = 32, seed = 21) {
torch_manual_seed(seed)
W <- torch_randn(d, d)
W <- W / max(abs(eigen(as.matrix(W), only.values = TRUE)$values)) * rho
Wx <- torch_randn(d, 4) * 0.1
xs <- lapply(1:T_len, \(t) torch_randn(1, 4))
h <- torch_zeros(1, d, requires_grad = TRUE)
hs <- list(h)
for (t in 1:T_len) {
h <- torch_tanh(hs[[t]]$matmul(W$t()) + xs[[t]]$matmul(Wx$t()))
h$retain_grad(); hs[[t + 1]] <- h
}
hs[[T_len + 1]]$sum()$backward()
vapply(1:T_len, \(t) {
g <- hs[[t]]$grad
if (is.null(g)) NA_real_ else as.numeric(torch_norm(g))
}, numeric(1))
}
vg <- bind_rows(lapply(c(0.7, 0.95, 1.0, 1.15), \(r)
data.frame(lag = 60 - (1:60), g = pmax(grad_by_lag(r), 1e-30),
rho = sprintf("spectral radius = %.2f", r))))
ggplot(vg, aes(lag, g, color = rho)) +
geom_hline(yintercept = 1, linetype = "dashed", color = "grey45") +
geom_line(linewidth = 0.9) +
scale_y_log10() + scale_x_reverse() +
scale_color_viridis_d(option = "plasma", end = 0.9) +
labs(title = "Gradient reaching back through time",
subtitle = "Below the dashed line the signal has vanished; above it, it explodes. Both are geometric in the lag",
x = "Lag (steps back from the loss)", y = "Gradient norm (log scale)",
color = NULL) +
theme_dspa()
vg |> filter(lag %in% c(5, 20, 40, 55)) |>
select(rho, lag, g) |> pivot_wider(names_from = lag, values_from = g,
names_prefix = "lag_") |>
mutate(across(where(is.numeric), \(z) signif(z, 3)))At \(\rho=0.7\) the gradient has fallen below \(10^{-8}\) by lag 40, the network is structurally blind to anything further back, regardless of how long it trains.
# Clipping rescales the gradient without changing its direction
torch_manual_seed(23)
p <- torch_randn(100, requires_grad = TRUE)
(p$sum() * 1e4)$backward()
before <- as.numeric(torch_norm(p$grad))
dir_before <- p$grad / torch_norm(p$grad)
nn_utils_clip_grad_norm_(list(p), max_norm = 1.0)#> torch_tensor
#> 100000
#> [ CPUFloatType{} ]after <- as.numeric(torch_norm(p$grad))
dir_after <- p$grad / torch_norm(p$grad)
c(norm_before = signif(before, 4), norm_after = signif(after, 4),
cosine_between_directions =
signif(as.numeric(torch_sum(dir_before * dir_after)), 6),
interpretation = "direction preserved exactly; only the length changes")#> norm_before
#> "1e+05"
#> norm_after
#> "1"
#> cosine_between_directions
#> "1"
#> interpretation
#> "direction preserved exactly; only the length changes"14.47 LSTM and GRU gating
The long short-term memory cell (Hochreiter & Schmidhuber, 1997) adds a separate cell state \(\mathbf c_t\) whose update is additive:
\[ \begin{aligned} \mathbf f_t&=\sigma\big(W_f[\mathbf h_{t-1},\mathbf x_t]+\mathbf b_f\big) &&\text{forget gate}\\ \mathbf i_t&=\sigma\big(W_i[\mathbf h_{t-1},\mathbf x_t]+\mathbf b_i\big) &&\text{input gate}\\ \tilde{\mathbf c}_t&=\tanh\big(W_c[\mathbf h_{t-1},\mathbf x_t]+\mathbf b_c\big) &&\text{candidate}\\ \mathbf c_t&=\mathbf f_t\odot\mathbf c_{t-1}+\mathbf i_t\odot\tilde{\mathbf c}_t &&\textbf{additive update}\\ \mathbf o_t&=\sigma\big(W_o[\mathbf h_{t-1},\mathbf x_t]+\mathbf b_o\big) &&\text{output gate}\\ \mathbf h_t&=\mathbf o_t\odot\tanh(\mathbf c_t) \end{aligned} \]
Why the fourth line fixes the problem. Differentiating the cell state recursion gives
\[\frac{\partial\mathbf c_t}{\partial\mathbf c_{t-1}}=\operatorname{diag}(\mathbf f_t),\]
so the gradient along the cell path is multiplied by the forget gate rather than by \(W_{hh}^\top\operatorname{diag}(\sigma')\). When \(\mathbf f_t\approx1\) the gradient passes through essentially unattenuated, the “constant error carousel”, and information can be carried hundreds of steps.
This is the same mechanism as a residual connection, on the time axis. Part 2, §14.32 showed that \(y=\mathcal F(x)+x\) creates a gradient path through depth with derivative exactly 1. The LSTM creates a path through time with derivative \(\operatorname{diag}(\mathbf f_t)\). Both replace a product of Jacobians with something closer to a sum, and both were invented to solve the same problem on different axes, LSTM in 1997, ResNet in 2015.
Common misconception: “the LSTM remembers what matters.” The gates are learned, differentiable interpolators, not a memory mechanism in any deliberate sense. \(\mathbf f_t\in(0,1)^d\) is a soft mask computed from the current input and previous state, so “forgetting” is a continuous down-weighting applied to every cell dimension at every step, nothing is stored or retrieved.
Two practical consequences. Initialize the forget-gate bias to a positive value (typically 1): a zero bias makes \(\mathbf f_t\approx\sigma(0)=0.5\) initially, halving the gradient at each step and reintroducing exponential decay before training has begun (Jozefowicz et al., 2015). And an LSTM’s practical range is hundreds of steps, not thousands — long-range dependence is improved, not solved, which is what motivated attention.
grad_through_time <- function(cell_type, T_len = 80, d = 32, seed = 31) {
torch_manual_seed(seed)
cell <- switch(cell_type,
rnn = nn_rnn(4, d, batch_first = TRUE),
lstm = nn_lstm(4, d, batch_first = TRUE),
gru = nn_gru(4, d, batch_first = TRUE))
xs <- torch_randn(1, T_len, 4, requires_grad = TRUE)
out <- cell(xs)[[1]]
out[1, T_len, ]$sum()$backward()
# Gradient at each INPUT step tells us how far influence reaches
apply(abs(as.array(xs$grad[1, , ])), 1, sum)
}
gt <- bind_rows(lapply(c(rnn = "rnn", lstm = "lstm", gru = "gru"), \(ct)
data.frame(step = 1:80, g = pmax(grad_through_time(ct), 1e-30),
cell = toupper(ct))), .id = "id")
ggplot(gt, aes(80 - step, g, color = cell)) +
geom_line(linewidth = 0.9) +
scale_y_log10() + scale_x_reverse() +
scale_color_manual(values = c(RNN = "#D8433B", LSTM = "#3B7DD8", GRU = "#7FB069")) +
labs(title = "How far back the loss at step 80 can influence the input",
subtitle = "Gradient at each input step, at initialization. The gated cells reach substantially further",
x = "Lag (steps back)", y = "Input-gradient magnitude (log scale)",
color = NULL) +
theme_dspa()
gt |> filter(step %in% c(79, 60, 40, 20, 1)) |>
mutate(lag = 80 - step) |> select(cell, lag, g) |>
pivot_wider(names_from = lag, values_from = g, names_prefix = "lag_") |>
mutate(across(where(is.numeric), \(z) signif(z, 3)))# The same measurement as a surface, over lag and hidden width
widths_g <- c(8, 16, 32, 64, 128)
reach_surface <- function(cell_type) {
t(vapply(widths_g, function(d) {
g <- grad_through_time(cell_type, T_len = 80, d = d)
log10(pmax(rev(g), 1e-30)) # index 1 = most recent step
}, numeric(80)))
}
Zrnn_g <- reach_surface("rnn")
Zlstm_g <- reach_surface("lstm")
plot_ly() |>
add_surface(x = 0:79, y = widths_g, z = Zlstm_g, opacity = 0.9,
showscale = FALSE, colorscale = "Blues", name = "LSTM") |>
add_surface(x = 0:79, y = widths_g, z = Zrnn_g, opacity = 0.9,
showscale = FALSE, colorscale = "Reds", name = "RNN") |>
layout(title = "Gradient reach over lag and hidden width: LSTM (blue, upper) vs. RNN (red, lower)",
scene = list(xaxis = list(title = "Lag (steps back)"),
yaxis = list(title = "Hidden width", type = "log"),
zaxis = list(title = "log10 input-gradient magnitude")))Rotate along the lag axis. The RNN surface falls away steeply and keeps falling, and widening the hidden state does not rescue it, the decay is geometric in the lag regardless of \(d\). The LSTM surface stays far higher across the whole range, which is the additive cell path doing its work.
# The forget-gate bias initialization, measured
lstm_reach <- function(forget_bias, T_len = 80, d = 32, seed = 33) {
torch_manual_seed(seed)
cell <- nn_lstm(4, d, batch_first = TRUE)
# torch packs gates as (input, forget, cell, output); the forget block is
# rows d+1 .. 2d of each bias vector
with_no_grad({
cell$bias_ih_l1[(d+1):(2*d)]$fill_(forget_bias)
cell$bias_hh_l1[(d+1):(2*d)]$fill_(0)
})
xs <- torch_randn(1, T_len, 4, requires_grad = TRUE)
cell(xs)[[1]][1, T_len, ]$sum()$backward()
g <- apply(abs(as.array(xs$grad[1, , ])), 1, sum)
c(forget_bias = forget_bias, grad_at_lag_40 = g[40], grad_at_lag_70 = g[10])
}
as.data.frame(do.call(rbind, lapply(c(-1, 0, 1, 2), lstm_reach))) |>
mutate(across(-forget_bias, \(z) signif(z, 3)))A positive forget bias keeps the gate near 1 early in training, and the long-lag gradient is orders of magnitude larger for it.
14.47.1 The GRU
The gated recurrent unit achieves the same effect with two gates instead of three, merging the cell and hidden states:
\[ \begin{aligned} \mathbf z_t&=\sigma\big(W_z[\mathbf h_{t-1},\mathbf x_t]\big) &&\text{update gate}\\ \mathbf r_t&=\sigma\big(W_r[\mathbf h_{t-1},\mathbf x_t]\big) &&\text{reset gate}\\ \tilde{\mathbf h}_t&=\tanh\big(W[\mathbf r_t\odot\mathbf h_{t-1},\mathbf x_t]\big)\\ \mathbf h_t&=(1-\mathbf z_t)\odot\mathbf h_{t-1}+\mathbf z_t\odot\tilde{\mathbf h}_t &&\textbf{convex interpolation} \end{aligned} \]
The last line is a convex combination of the old state and a candidate, so the update and forget roles are tied: whatever is written in must displace an equal amount. That coupling costs one gate’s worth of flexibility and about 25% of the parameters.
d_h <- 128; d_in <- 64
data.frame(
cell = c("RNN", "GRU", "LSTM"),
gates = c(0, 2, 3),
parameters = vapply(list(nn_rnn(d_in, d_h), nn_gru(d_in, d_h), nn_lstm(d_in, d_h)),
\(m) sum(vapply(m$parameters, \(p) prod(dim(p)), numeric(1))),
numeric(1)),
formula = c("d(d+m)+d", "3d(d+m)+3d", "4d(d+m)+4d"))14.48 Encoder–decoder and the bottleneck
For unaligned sequence-to-sequence tasks, translation, summarization, the standard architecture encodes the source into a fixed vector and decodes from it:
\[\mathbf c=\mathbf h_{T}^{\text{enc}}, \qquad \mathbf s_t=\text{dec}\big(\mathbf s_{t-1},\mathbf y_{t-1},\mathbf c\big).\]
The bottleneck is the entire motivation for attention. Every source sequence, whether five words or five hundred, must be compressed into one fixed-dimensional vector \(\mathbf c\). Translation quality degrades sharply with source length for exactly this reason (Cho et al., 2014), the vector runs out of capacity, and the earliest source tokens are the ones that have passed through the most recurrent steps and are therefore most degraded.
Attention removes the constraint by letting the decoder construct a different context vector at every output step, drawn from all encoder states rather than only the last.
# A copy task: reproduce the input sequence. Fixed-vector capacity is the limit.
copy_task <- function(T_len, d_hidden = 24, n_train = 400, seed = 41) {
torch_manual_seed(seed)
set.seed(seed)
V <- 8
X <- torch_randint(1, V + 1, size = c(n_train, T_len), dtype = torch_long())
emb <- nn_embedding(V + 1, 16)
enc <- nn_gru(16, d_hidden, batch_first = TRUE)
dec <- nn_gru(16, d_hidden, batch_first = TRUE)
head <- nn_linear(d_hidden, V + 1)
pars <- c(emb$parameters, enc$parameters, dec$parameters, head$parameters)
opt <- optim_adam(pars, lr = 5e-3)
for (e in 1:250) {
opt$zero_grad()
h_enc <- enc(emb(X))[[2]]
dec_in <- emb(torch_cat(list(torch_ones(n_train, 1, dtype = torch_long()),
X[, 1:(T_len-1)]), dim = 2))
logits <- head(dec(dec_in, h_enc)[[1]])
l <- nnf_cross_entropy(logits$reshape(c(-1, V + 1)), X$reshape(-1))
l$backward()
opt$step()
}
# Clean evaluation step avoiding nested parenthesis errors
preds <- with_no_grad({
dec_in_eval <- emb(torch_cat(list(torch_ones(n_train, 1, dtype = torch_long()),
X[, 1:(T_len-1)]), dim = 2))
h_eval <- enc(emb(X))[[2]]
head(dec(dec_in_eval, h_eval)[[1]])$argmax(dim = 3)
})
acc <- mean(as.numeric(preds) == as.numeric(X))
c(sequence_length = T_len, copy_accuracy = acc)
}
bt <- as.data.frame(do.call(rbind, lapply(c(4, 8, 16, 32), copy_task)))
bt |> mutate(copy_accuracy = round(copy_accuracy, 4))ggplot(bt, aes(sequence_length, copy_accuracy)) +
geom_line(linewidth = 1, color = "steelblue") +
geom_point(size = 2.6) +
scale_x_log10(breaks = bt$sequence_length) +
labs(title = "Copying a sequence through a fixed-size vector",
subtitle = "Hidden size held at 24. Accuracy falls as the source outgrows the vector's capacity",
x = "Sequence length (log scale)", y = "Token-level copy accuracy") +
theme_dspa()
The task is trivially learnable in principle, copy the input, and accuracy still collapses with length, because the information must pass through a 24-dimensional vector.
14.49 Attention, derived from alignment
The bottleneck’s cause is that \(\mathbf c\) is fixed. Attention makes it a function of the decoding step.
The construction. Keep all encoder states \(\mathbf h_1,\dots,\mathbf h_T\). At decoder step \(t\), score each against the current decoder state, normalize the scores, and take the weighted average:
\[ \begin{aligned} e_{t,j}&=\text{score}\big(\mathbf s_{t-1},\mathbf h_j\big) &&\text{alignment score}\\ \alpha_{t,j}&=\frac{\exp(e_{t,j})}{\sum_{k=1}^{T}\exp(e_{t,k})} &&\text{softmax normalization}\\ \mathbf c_t&=\sum_{j=1}^{T}\alpha_{t,j}\,\mathbf h_j &&\textbf{step-specific context} \end{aligned} \]
Three properties follow immediately, and each matters.
The context is now step-specific, so there is no single vector that must carry everything. Every encoder state is one softmax-weighted step from every decoder state, so gradient reaches the earliest source token in \(O(1)\) hops rather than \(O(T)\) recurrent steps. And the operation is differentiable, so the alignment is learned rather than specified.
The scoring function has two classical forms (Bahdanau et al., 2015; Luong et al., 2015):
\[ \begin{aligned} \textbf{Additive: }&\quad e_{t,j}=\mathbf v^\top\tanh\big(W_s\mathbf s_{t-1}+W_h\mathbf h_j\big)\\ \textbf{Dot-product: }&\quad e_{t,j}=\mathbf s_{t-1}^\top\mathbf h_j \end{aligned} \]
The dot-product form has no parameters and reduces to a matrix multiplication over all pairs at once, which is why it became the standard.
# Additive and dot-product attention, from the definitions
additive_attention <- function(s, H, W_s, W_h, v) {
# s: (B, d_s); H: (B, T, d_h)
scores <- torch_matmul(torch_tanh(s$matmul(W_s$t())$unsqueeze(2) +
H$matmul(W_h$t())), v)$squeeze(3)
alpha <- nnf_softmax(scores, dim = 2)
list(context = torch_bmm(alpha$unsqueeze(2), H)$squeeze(2), weights = alpha)
}
dot_attention <- function(s, H) {
scores <- torch_bmm(H, s$unsqueeze(3))$squeeze(3) # (B, T)
alpha <- nnf_softmax(scores, dim = 2)
list(context = torch_bmm(alpha$unsqueeze(2), H)$squeeze(2), weights = alpha)
}
torch_manual_seed(51)
B <- 2; T_a <- 10; d <- 16
H_enc <- torch_randn(B, T_a, d); s_dec <- torch_randn(B, d)
r_dot <- dot_attention(s_dec, H_enc)
c(weights_shape = paste(dim(r_dot$weights), collapse = " x "),
context_shape = paste(dim(r_dot$context), collapse = " x "),
weights_sum_to_one = signif(as.numeric(torch_sum(r_dot$weights[1, ])), 6),
parameters_in_dot_attention = 0)#> weights_shape context_shape
#> "2 x 10" "2 x 16"
#> weights_sum_to_one parameters_in_dot_attention
#> "1" "0"# A learned alignment on a reversal task: output j should attend to input T-j+1
torch_manual_seed(53)
T_r <- 12; V_r <- 10; n_r <- 500
Xr <- torch_randint(1, V_r + 1, size = c(n_r, T_r), dtype = torch_long())
Yr <- torch_flip(Xr, dims = 2)
emb_r <- nn_embedding(V_r + 1, 24)
enc_r <- nn_gru(24, 32, batch_first = TRUE)
dec_r <- nn_gru(24, 32, batch_first = TRUE)
proj_r <- nn_linear(32, 32)
head_r <- nn_linear(64, V_r + 1)
pars_r <- c(emb_r$parameters, enc_r$parameters, dec_r$parameters,
proj_r$parameters, head_r$parameters)
opt_r <- optim_adam(pars_r, lr = 3e-3)
fwd_r <- function(X, Y) {
He <- enc_r(emb_r(X))[[1]]
dec_in <- emb_r(torch_cat(list(torch_ones(X$size(1), 1, dtype = torch_long()),
Y[, 1:(T_r - 1)]), dim = 2))
S <- dec_r(dec_in)[[1]]
scores <- torch_bmm(proj_r(S), He$transpose(2, 3)) / sqrt(32)
alpha <- nnf_softmax(scores, dim = 3)
ctx <- torch_bmm(alpha, He)
list(logits = head_r(torch_cat(list(S, ctx), dim = 3)), alpha = alpha)
}
for (e in 1:400) {
opt_r$zero_grad()
o <- fwd_r(Xr, Yr)
l <- nnf_cross_entropy(o$logits$reshape(c(-1, V_r + 1)), Yr$reshape(-1))
l$backward(); opt_r$step()
}
o <- with_no_grad(fwd_r(Xr, Yr))
c(final_loss = signif(as.numeric(l), 4),
token_accuracy = round(mean(as.numeric(o$logits$argmax(dim = 3)) ==
as.numeric(Yr)), 4))#> final_loss token_accuracy
#> 0.06029 0.99080A <- as.array(o$alpha[1, , ])
expand.grid(output_step = 1:T_r, input_step = 1:T_r) |>
mutate(w = as.vector(t(A))) |>
ggplot(aes(input_step, output_step, fill = w)) +
geom_raster() +
geom_abline(slope = -1, intercept = T_r + 1, color = "white",
linetype = "dashed", linewidth = 0.6) +
scale_fill_viridis_c(option = "inferno", name = expression(alpha)) +
scale_y_reverse(breaks = 1:T_r) + scale_x_continuous(breaks = 1:T_r) +
coord_fixed() +
labs(title = "Learned attention on a sequence-reversal task",
subtitle = "The dashed anti-diagonal is the correct alignment. Nothing told the model this; it was learned",
x = "Input position", y = "Output position") +
theme_dspa(10)
# --- Interactive equivalent ------------------------------------------------
plot_ly(z = A, type = "heatmap", colorscale = "Inferno") |>
layout(title = "Attention weights on the reversal task",
xaxis = list(title = "Input position", scaleanchor = "y"),
yaxis = list(title = "Output position", autorange = "reversed"))The model discovers the anti-diagonal alignment on its own, the attention mechanism supplies the capacity to align, and gradient descent supplies the alignment.
PART II: ATTENTION AND TRANSFORMERS
14.50 Scaled dot-product attention
The Transformer (Vaswani et al., 2017) generalizes §14.49 by giving every position three learned projections, a query, a key, and a value, and letting positions attend to each other rather than only across an encoder–decoder boundary.
\[\boxed{\;\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V\;}\]
with \(Q=XW^Q\in\mathbb R^{T\times d_k}\), \(K=XW^K\), \(V=XW^V\in\mathbb R^{T\times d_v}\).
The database metaphor is exact. Each position emits a query describing what it is looking for and a key describing what it offers; their dot product measures match quality; the softmax turns matches into weights; and the values are what actually gets retrieved and mixed.
14.50.1 Why \(\sqrt{d_k}\)
Common misconception: “the \(1/\sqrt{d_k}\) is a cosmetic normalization.” Remove it and attention stops training at moderate \(d_k\).
Suppose the components of \(\mathbf q\) and \(\mathbf k\) are independent with mean 0 and variance 1. Then \[\mathbb E\big[\mathbf q\cdot\mathbf k\big]=0,\qquad \operatorname{Var}\big(\mathbf q\cdot\mathbf k\big)=\sum_{i=1}^{d_k}\operatorname{Var}(q_ik_i)=d_k,\] so the dot products have standard deviation \(\sqrt{d_k}\) and grow with the key dimension. At \(d_k=512\) the logits entering the softmax have spread around \(\pm23\).
A softmax with logits that far apart saturates: it becomes nearly one-hot, and its Jacobian \(\partial\alpha_i/\partial e_j=\alpha_i(\delta_{ij}-\alpha_j)\) collapses to zero because \(\alpha_i(1-\alpha_i)\to0\) at both extremes. The gradient vanishes, the same saturation mechanism that made sigmoid activations untrainable in Part 1, §14.8, arriving here through a different route.
Dividing by \(\sqrt{d_k}\) restores unit variance regardless of \(d_k\).
torch_manual_seed(61)
logit_spread <- function(d_k, n = 4000, scaled) {
q <- torch_randn(n, d_k); k <- torch_randn(n, d_k)
s <- torch_sum(q * k, dim = 2)
if (scaled) s <- s / sqrt(d_k)
as.numeric(torch_std(s))
}
data.frame(d_k = c(8, 32, 128, 512),
sd_unscaled = round(vapply(c(8,32,128,512), logit_spread,
numeric(1), scaled = FALSE), 3),
sqrt_d_k = round(sqrt(c(8,32,128,512)), 3),
sd_scaled = round(vapply(c(8,32,128,512), logit_spread,
numeric(1), scaled = TRUE), 3))The unscaled standard deviation tracks \(\sqrt{d_k}\) exactly, as the variance argument predicts; the scaled version is 1 at every dimension.
# What saturation does to the gradient
attn_grad_norm <- function(d_k, T_len = 32, scaled, seed = 63) {
torch_manual_seed(seed)
Q <- torch_randn(1, T_len, d_k, requires_grad = TRUE)
K <- torch_randn(1, T_len, d_k)
s <- torch_bmm(Q, K$transpose(2, 3))
if (scaled) s <- s / sqrt(d_k)
a <- nnf_softmax(s, dim = 3)
a$sum()$backward()
c(gradient_norm = as.numeric(torch_norm(Q$grad)),
max_attention_weight = as.numeric(torch_max(a)),
attention_entropy = as.numeric(-torch_sum(a * torch_log(a + 1e-12)) /
(T_len)))
}
d_grid <- c(8, 32, 128, 512)
sat <- rbind(
cbind(d_k = d_grid, scaled = 0,
t(vapply(d_grid, attn_grad_norm, numeric(3), scaled = FALSE))),
cbind(d_k = d_grid, scaled = 1,
t(vapply(d_grid, attn_grad_norm, numeric(3), scaled = TRUE))))
as.data.frame(sat) |>
mutate(scaled = ifelse(scaled == 1, "with 1/sqrt(d_k)", "unscaled"),
across(where(is.numeric), \(z) signif(z, 4)))Without scaling, the maximum attention weight approaches 1 and the entropy collapses as \(d_k\) grows, the softmax has become a hard argmax, and the gradient norm falls with it.
dk_grid <- round(exp(seq(log(4), log(512), length.out = 25)))
temp_grid <- 10^seq(-0.6, 1.2, length.out = 25) # divisor applied to logits
Zent <- outer(temp_grid, dk_grid, Vectorize(function(tau, dk) {
torch_manual_seed(65)
Q <- torch_randn(1, 32, dk); K <- torch_randn(1, 32, dk)
a <- nnf_softmax(torch_bmm(Q, K$transpose(2, 3)) / tau, dim = 3)
as.numeric(-torch_sum(a * torch_log(a + 1e-12)) / 32)
}))
plot_ly(x = dk_grid, y = temp_grid, z = Zent, type = "surface",
colorscale = "Viridis", colorbar = list(title = "Attention\nentropy")) |>
add_trace(x = dk_grid, y = sqrt(dk_grid),
z = vapply(seq_along(dk_grid), \(i) {
torch_manual_seed(65)
Q <- torch_randn(1, 32, dk_grid[i]); K <- torch_randn(1, 32, dk_grid[i])
a <- nnf_softmax(torch_bmm(Q, K$transpose(2,3))/sqrt(dk_grid[i]), dim = 3)
as.numeric(-torch_sum(a * torch_log(a + 1e-12))/32) }, numeric(1)),
type = "scatter3d", mode = "lines", name = "divisor = sqrt(d_k)",
line = list(width = 8, color = "red")) |>
layout(title = "Attention entropy over key dimension and logit divisor",
scene = list(xaxis = list(title = "d_k", type = "log"),
yaxis = list(title = "Logit divisor", type = "log"),
zaxis = list(title = "Mean attention entropy")))The red curve is \(\tau=\sqrt{d_k}\). It tracks a ridge of near-constant entropy across three orders of magnitude in \(d_k\), that is what the scaling buys. Below the ridge the surface falls toward zero entropy (saturated, no gradient); above it, entropy approaches \(\log T\) (uniform, no selectivity).
14.50.2 Causal masking
An autoregressive model must not see the future. Masking sets the disallowed logits to \(-\infty\) before the softmax, so those weights become exactly zero:
\[e_{ij}\leftarrow\begin{cases}e_{ij}, & j\le i\\ -\infty, & j>i\end{cases}\]
scaled_dot_attention <- function(Q, K, V, mask = NULL) {
d_k <- tail(dim(Q), 1)
scores <- torch_matmul(Q, K$transpose(-2, -1)) / sqrt(d_k)
if (!is.null(mask)) scores <- scores$masked_fill(mask, -Inf)
alpha <- nnf_softmax(scores, dim = -1)
list(out = torch_matmul(alpha, V), weights = alpha)
}
causal_mask <- function(T_len)
torch_triu(torch_ones(T_len, T_len, dtype = torch_bool()), diagonal = 1)
torch_manual_seed(67)
T_m <- 6; dm <- 16
Xm <- torch_randn(1, T_m, dm)
res_m <- scaled_dot_attention(Xm, Xm, Xm, mask = causal_mask(T_m))
round(as.matrix(res_m$weights[1, , ]), 3)#> [,1] [,2] [,3] [,4] [,5] [,6]
#> [1,] 1.000 0.000 0.000 0.000 0.000 0.000
#> [2,] 0.304 0.696 0.000 0.000 0.000 0.000
#> [3,] 0.003 0.003 0.994 0.000 0.000 0.000
#> [4,] 0.141 0.065 0.082 0.712 0.000 0.000
#> [5,] 0.000 0.001 0.004 0.000 0.994 0.000
#> [6,] 0.231 0.012 0.002 0.009 0.008 0.738The strict upper triangle is exactly zero: position \(i\) attends only to positions \(\le i\). Masking before the softmax is essential, zeroing the weights after normalization would leave the remaining weights not summing to 1.
14.51 Multi-head attention
\[ \begin{aligned} \mathrm{MultiHead}(Q,K,V)&=\mathrm{Concat}\big(\mathrm{head}_1,\dots,\mathrm{head}_h\big)W^O\\ \mathrm{head}_i&=\mathrm{Attention}\big(QW_i^Q,\ KW_i^K,\ VW_i^V\big),\qquad d_k=d_v=\frac{d_{\text{model}}}{h} \end{aligned} \]
Common misconception: “more heads means more capacity.” With \(d_k=d_{\text{model}}/h\), the total projection parameters are \(4d_{\text{model}}^2\) regardless of \(h\), 8 heads of dimension 64 and 1 head of dimension 512 have identical parameter counts and nearly identical FLOPs.
What multiple heads buy is simultaneous attention to different subspaces. A single softmax produces one weighted average, so a head that must attend to both a syntactic dependency and a coreference link can only average the two — and the average may point at neither. Splitting into heads lets one attend syntactically while another attends referentially, and the concatenation keeps both.
The empirical picture is more sober than the intuition. Many heads are prunable at inference with little loss, and in some layers a single head suffices (Michel et al., 2019), so the gain is largely in optimization, giving gradient descent several parallel hypotheses to work with rather than in representational necessity.
multi_head <- nn_module("MultiHead",
initialize = function(d_model, n_heads) {
stopifnot(d_model %% n_heads == 0)
self$h <- n_heads; self$d_k <- d_model %/% n_heads
self$Wq <- nn_linear(d_model, d_model, bias = FALSE)
self$Wk <- nn_linear(d_model, d_model, bias = FALSE)
self$Wv <- nn_linear(d_model, d_model, bias = FALSE)
self$Wo <- nn_linear(d_model, d_model, bias = FALSE)
},
forward = function(x, mask = NULL) {
B <- x$size(1); T_len <- x$size(2)
split_heads <- function(z)
z$view(c(B, T_len, self$h, self$d_k))$transpose(2, 3) # (B, h, T, d_k)
Q <- split_heads(self$Wq(x)); K <- split_heads(self$Wk(x))
V <- split_heads(self$Wv(x))
a <- scaled_dot_attention(Q, K, V, mask)
merged <- a$out$transpose(2, 3)$contiguous()$view(c(B, T_len, -1))
list(out = self$Wo(merged), weights = a$weights)
})
d_model <- 256
data.frame(
n_heads = c(1, 2, 4, 8, 16),
d_per_head = d_model %/% c(1, 2, 4, 8, 16),
parameters = vapply(c(1, 2, 4, 8, 16), \(h)
sum(vapply(multi_head(d_model, h)$parameters, \(p) prod(dim(p)), numeric(1))),
numeric(1)),
note = "identical: 4 * d_model^2 regardless of head count")# Different heads attend differently, which is the point
torch_manual_seed(71)
mh <- multi_head(64, 4)
Xh <- torch_randn(1, 14, 64)
wh <- with_no_grad(mh(Xh)$weights) # (1, 4, 14, 14)
bind_rows(lapply(1:4, \(i) {
A <- as.array(wh[1, i, , ])
expand.grid(key = 1:14, query = 1:14) |>
mutate(w = as.vector(t(A)), head = sprintf("head %d", i))
})) |>
ggplot(aes(key, query, fill = w)) +
geom_raster() + facet_wrap(~ head, nrow = 1) + coord_fixed() +
scale_fill_viridis_c(option = "inferno", guide = "none") +
scale_y_reverse() +
labs(title = "Four heads at initialization, same input",
subtitle = "Each attends to a different subspace; a single head could only produce one of these patterns",
x = "Key position", y = "Query position") +
theme_dspa(9)
# Head diversity: pairwise correlation between heads' attention matrices
pw <- combn(4, 2, function(ij)
cor(as.vector(as.array(wh[1, ij[1], , ])), as.vector(as.array(wh[1, ij[2], , ]))))
c(mean_pairwise_correlation = signif(mean(pw), 3),
max_pairwise_correlation = signif(max(pw), 3),
interpretation = "low correlation = heads are attending to different things")#> mean_pairwise_correlation
#> "0.0618"
#> max_pairwise_correlation
#> "0.202"
#> interpretation
#> "low correlation = heads are attending to different things"14.52 Positional encoding
Common misconception: “the Transformer reads the sequence in order.” It does not. Self-attention is a set operation, and it is exactly permutation-equivariant: permute the input rows by \(P\) and every output row permutes identically, \[\mathrm{Attention}(PX)=P\cdot\mathrm{Attention}(X).\]
Without positional information a Transformer literally cannot distinguish “the dog bit the man” from “the man bit the dog”, the multiset of tokens is the same, so the multiset of outputs is the same. Order must be injected into the representation, and that is the only thing positional encoding does.
torch_manual_seed(81)
mh_pe <- multi_head(32, 4)
Xp <- torch_randn(1, 8, 32)
perm <- sample(8)
with_no_grad({
out_then_perm <- mh_pe(Xp)$out[1, perm, ] # P . Attention(X)
perm_then_out <- mh_pe(Xp[, perm, , drop = FALSE])$out[1, , ] # Attention(PX)
})
c(max_abs_difference = signif(as.numeric(torch_max(torch_abs(
out_then_perm - perm_then_out))), 3),
conclusion = "Attention(PX) = P.Attention(X): the operation cannot see order")#> max_abs_difference
#> "8.94e-08"
#> conclusion
#> "Attention(PX) = P.Attention(X): the operation cannot see order"Sinusoidal encoding adds a fixed, position-dependent pattern to the embeddings:
\[PE_{(pos,\,2i)}=\sin\!\left(\frac{pos}{10000^{2i/d}}\right),\qquad PE_{(pos,\,2i+1)}=\cos\!\left(\frac{pos}{10000^{2i/d}}\right)\]
Its useful property is that \(PE_{pos+k}\) is a fixed linear function of \(PE_{pos}\) for any offset \(k\), a rotation in each 2-D \((\sin,\cos)\) pair. So the model can learn to attend by relative offset, and it extrapolates in principle to positions longer than any seen in training.
sinusoidal_pe <- function(T_len, d) {
pos <- torch_arange(0, T_len - 1)$unsqueeze(2)
i2 <- torch_arange(0, d - 1, 2)
denom <- torch_pow(10000, i2 / d)
pe <- torch_zeros(T_len, d)
pe[, seq(1, d, by = 2)] <- torch_sin(pos / denom)
pe[, seq(2, d, by = 2)] <- torch_cos(pos / denom)
pe
}
PE <- sinusoidal_pe(80, 64)
# The linear-shift property: PE[pos+k] is a rotation of PE[pos]
k_off <- 7
pairs <- 1:32
rot_err <- vapply(pairs, function(j) {
c_ <- 2*j - 1; s_ <- 2*j
th <- k_off / (10000^((2*(j-1))/64))
pred_sin <- as.numeric(PE[1:60, c_]) * cos(th) + as.numeric(PE[1:60, s_]) * sin(th)
max(abs(pred_sin - as.numeric(PE[(1+k_off):(60+k_off), c_])))
}, numeric(1))
c(max_rotation_error_across_32_frequency_pairs = signif(max(rot_err), 3),
interpretation = "PE[pos+k] = R(k) PE[pos] exactly, for a fixed rotation R(k)")#> max_rotation_error_across_32_frequency_pairs
#> "3.31e-06"
#> interpretation
#> "PE[pos+k] = R(k) PE[pos] exactly, for a fixed rotation R(k)"expand.grid(pos = 1:80, dim = 1:64) |>
mutate(v = as.vector(as.array(PE))) |>
ggplot(aes(dim, pos, fill = v)) +
geom_raster() +
scale_fill_gradient2(low = "#3B7DD8", mid = "white", high = "#D8433B",
name = NULL) +
scale_y_reverse() +
labs(title = "Sinusoidal positional encoding",
subtitle = "Low dimensions oscillate fast, high dimensions slowly -- a multi-scale positional code",
x = "Embedding dimension", y = "Position") +
theme_dspa(10)
plot_ly(x = 1:64, y = 1:80, z = as.array(PE), type = "surface",
colorscale = "RdBu", colorbar = list(title = "PE value")) |>
layout(title = "The positional encoding matrix as a surface",
scene = list(xaxis = list(title = "Embedding dimension"),
yaxis = list(title = "Position"),
zaxis = list(title = "Encoding value")))Rotate along the dimension axis: the wavelength lengthens monotonically, so low dimensions distinguish adjacent positions and high dimensions distinguish distant ones. The encoding is a positional analogue of a wavelet basis.
Modern alternatives are worth naming. Learned absolute embeddings are simply a lookup table trained with the model, simpler, but they cannot extrapolate past the trained maximum length. RoPE (rotary) rotates \(Q\) and \(K\) by a position-dependent angle so that the dot product depends only on the relative offset, and it is the current default in large language models. ALiBi adds a linear distance penalty directly to the attention logits.
14.53 The Transformer block
\[ \begin{aligned} \textbf{Post-norm (original): }&\quad \begin{cases} \mathbf z=\mathrm{LayerNorm}\big(\mathbf x+\mathrm{MHA}(\mathbf x)\big)\\ \mathbf y=\mathrm{LayerNorm}\big(\mathbf z+\mathrm{FFN}(\mathbf z)\big) \end{cases}\\[3mm] \textbf{Pre-norm (modern): }&\quad \begin{cases} \mathbf z=\mathbf x+\mathrm{MHA}\big(\mathrm{LayerNorm}(\mathbf x)\big)\\ \mathbf y=\mathbf z+\mathrm{FFN}\big(\mathrm{LayerNorm}(\mathbf z)\big) \end{cases} \end{aligned} \]
with \(\mathrm{FFN}(\mathbf x)=W_2\,\phi(W_1\mathbf x+\mathbf b_1)+\mathbf b_2\) and inner width typically \(4d_{\text{model}}\).
Common misconception: “warmup is needed because the model is large.” It is needed because of where the LayerNorm sits, and pre-norm largely removes the requirement. In post-norm the residual stream passes through a LayerNorm at every block, so the clean additive path of Part 2, §14.32 is interrupted and gradients at deep layers are large and poorly scaled. The original Transformer needed a carefully tuned learning-rate warmup to survive early training (Part 1, §14.12).
Pre-norm normalizes the branch input and leaves the residual stream untouched, restoring the derivative-1 path from input to output. Gradients are better behaved at depth, warmup becomes optional, and the architecture scales past 100 layers (Xiong et al., 2020).
LayerNorm, not BatchNorm, throughout, sequences have variable lengths and often small batches, exactly the regime where batch statistics fail (Part 2, §14.33).
# Define the operator manually if you don't want to load rlang or purrr
# `%||%` <- function(x, y) if (is.null(x)) y else x
# Alternatively, load rlang:
library(rlang)
transformer_block <- nn_module("TransformerBlock",
initialize = function(d_model, n_heads, d_ff = NULL, dropout = 0.1,
pre_norm = TRUE) {
d_ff <- d_ff %||% (4 * d_model)
self$pre_norm <- pre_norm
self$mha <- multi_head(d_model, n_heads)
self$ln1 <- nn_layer_norm(d_model); self$ln2 <- nn_layer_norm(d_model)
self$ffn <- nn_sequential(nn_linear(d_model, d_ff), nn_gelu(),
nn_dropout(dropout), nn_linear(d_ff, d_model))
self$drop <- nn_dropout(dropout)
},
forward = function(x, mask = NULL) {
if (self$pre_norm) {
x <- x + self$drop(self$mha(self$ln1(x), mask)$out)
x + self$drop(self$ffn(self$ln2(x)))
} else {
x <- self$ln1(x + self$drop(self$mha(x, mask)$out))
self$ln2(x + self$drop(self$ffn(x)))
}
})
blk <- transformer_block(128, 8)
c(parameters = format(sum(vapply(blk$parameters, \(p) prod(dim(p)), numeric(1))),
big.mark = ","),
attention_share = round(sum(vapply(blk$mha$parameters, \(p) prod(dim(p)),
numeric(1))) /
sum(vapply(blk$parameters, \(p) prod(dim(p)),
numeric(1))), 3),
ffn_share = round(sum(vapply(blk$ffn$parameters, \(p) prod(dim(p)), numeric(1))) /
sum(vapply(blk$parameters, \(p) prod(dim(p)), numeric(1))), 3))#> parameters attention_share ffn_share
#> "197,760" "0.331" "0.666"The feed-forward network holds two-thirds of the parameters, not the attention. Attention is \(4d^2\) per block and the FFN is \(8d^2\), a fact that matters when deciding what to prune (Part 5).
grad_by_block <- function(pre_norm, n_blocks = 16, d = 64, seed = 91) {
torch_manual_seed(seed)
blocks <- lapply(1:n_blocks, \(i) transformer_block(d, 4, pre_norm = pre_norm))
net <- do.call(nn_sequential, blocks)
x <- torch_randn(2, 24, d)
net(x)$sum()$backward()
vapply(blocks, \(b) as.numeric(torch_norm(b$mha$Wq$weight$grad)), numeric(1))
}
pn <- bind_rows(
data.frame(block = 1:16, g = grad_by_block(TRUE), arch = "Pre-norm"),
data.frame(block = 1:16, g = grad_by_block(FALSE), arch = "Post-norm"))
ggplot(pn, aes(block, pmax(g, 1e-12), color = arch)) +
geom_line(linewidth = 1) + geom_point(size = 1.6) +
scale_y_log10() +
scale_color_manual(values = c(`Pre-norm` = "#3B7DD8", `Post-norm` = "#D8433B")) +
labs(title = "Gradient norm by block through a 16-block stack",
subtitle = "Block 1 is nearest the input. Pre-norm keeps the residual stream clear; post-norm does not",
x = "Block index", y = "Query-projection gradient norm (log scale)",
color = NULL) +
theme_dspa()
14.54 Complexity: attention versus recurrence
Chapter 12, §12.28 noted that RNNs are linear in sequence length where attention is quadratic, and deferred the comparison. Here it is.
| Layer type | Time per layer | Sequential ops | Memory | Max path length |
|---|---|---|---|---|
| Self-attention | \(O(T^2d)\) | \(\mathbf{O(1)}\) | \(O(T^2+Td)\) | \(\mathbf{O(1)}\) |
| Recurrent | \(O(Td^2)\) | \(O(T)\) | \(O(Td)\) | \(O(T)\) |
| Convolutional (\(k\)) | \(O(kTd^2)\) | \(O(1)\) | \(O(Td)\) | \(O(\log_k T)\) dilated |
| Restricted attention (\(w\)) | \(O(Twd)\) | \(O(1)\) | \(O(Tw+Td)\) | \(O(T/w)\) |
The asymptotics in \(T\) favour recurrence, and Transformers won anyway. Attention costs \(O(T^2d)\) against recurrence’s \(O(Td^2)\), so attention is more expensive whenever \(T>d\), which for long documents is most of the time.
Two columns explain the outcome. Sequential operations: an RNN must compute step \(t\) before step \(t+1\), so training cannot parallelize across time and wall-clock cost scales with \(T\) regardless of available hardware. Attention computes all positions at once, so on a GPU its higher FLOP count costs less time. And maximum path length: in an RNN, information from position 1 to position \(T\) traverses \(T\) transformations, each attenuating the gradient. In attention it traverses one, which is the property that makes long-range dependence learnable at all.
The trade is FLOPs for parallelism and gradient path length, an excellent trade on hardware with abundant parallel compute, and a poor one on a device that must process a stream one element at a time. That is why recurrent and state-space models remain competitive for streaming and embedded inference.
T_grid <- round(exp(seq(log(16), log(8192), length.out = 45)))
d_grid <- round(exp(seq(log(16), log(2048), length.out = 45)))
Zatt <- outer(d_grid, T_grid, function(d, T_) log10(T_^2 * d))
Zrnn <- outer(d_grid, T_grid, function(d, T_) log10(T_ * d^2))
plot_ly() |>
add_surface(x = T_grid, y = d_grid, z = Zatt, opacity = 0.88, showscale = FALSE,
colorscale = "Reds", name = "Attention O(T^2 d)") |>
add_surface(x = T_grid, y = d_grid, z = Zrnn, opacity = 0.88, showscale = FALSE,
colorscale = "Blues", name = "Recurrence O(T d^2)") |>
layout(title = "FLOP cost: attention (red) vs. recurrence (blue); they cross at T = d",
scene = list(xaxis = list(title = "Sequence length T", type = "log"),
yaxis = list(title = "Model dimension d", type = "log"),
zaxis = list(title = "log10 FLOPs")))The two surfaces intersect along the line \(T=d\). To the left, where sequences are short relative to the model, attention is cheaper; to the right, recurrence is. For \(d=512\) the crossover is at 512 tokens, and typical context lengths sit far to the right of it, which is why efficient-attention research exists.
scen <- expand.grid(T_len = c(128, 512, 2048, 8192), d = c(256, 512, 1024)) |>
mutate(attention_GFLOP = 2 * T_len^2 * d / 1e9,
recurrence_GFLOP = 2 * T_len * d^2 / 1e9,
ratio = round(attention_GFLOP / recurrence_GFLOP, 3),
attention_memory_MB = 4 * T_len^2 / 1e6,
cheaper = ifelse(ratio < 1, "attention", "recurrence"))
scen |> mutate(across(c(attention_GFLOP, recurrence_GFLOP, attention_memory_MB),
\(z) signif(z, 3)))The attention_memory_MB column is the harder constraint
in practice: the \(T\times T\)
attention matrix is quadratic in memory, so an
8,192-token context needs 268 MB per head per layer at float32.
FlashAttention avoids materializing it; sparse and linear attention
change the asymptotics outright.
PART III: APPLICATIONS
14.55 Tokenization and embeddings
Before a Transformer sees text, the text must become integers.
Common misconception: “tokens are words.” They are neither words nor characters but subword units learned from a corpus, and the difference has practical consequences.
Word-level vocabularies cannot represent anything unseen, every out-of-corpus word becomes
<UNK>, which is fatal for morphologically rich languages, proper nouns, and clinical terminology. Character-level vocabularies are tiny and complete but make sequences five to ten times longer, and attention costs \(O(T^2)\).Byte-pair encoding and its relatives interpolate: start from characters, repeatedly merge the most frequent adjacent pair, stop at a target vocabulary size. Frequent words become single tokens, rare words decompose into pieces, and nothing is unrepresentable. The consequences show up constantly: a rare medical term may occupy six tokens where a common word occupies one, so “context length” in tokens is not a fixed number of words, and per-token costs are not per-word costs.
# A minimal byte-pair encoding, to make the mechanism concrete
bpe_learn <- function(corpus, n_merges = 40) {
words <- unlist(strsplit(tolower(corpus), "\\s+"))
words <- words[nzchar(words)]
splits <- lapply(words, \(w) c(strsplit(w, "")[[1]], "</w>"))
merges <- list()
for (m in seq_len(n_merges)) {
pairs <- unlist(lapply(splits, \(s)
if (length(s) < 2) NULL else paste(head(s, -1), tail(s, -1), sep = "\u0001")))
if (!length(pairs)) break
tab <- sort(table(pairs), decreasing = TRUE)
if (tab[1] < 2) break
best <- strsplit(names(tab)[1], "\u0001")[[1]]
merges[[m]] <- best
splits <- lapply(splits, function(s) {
i <- 1; out <- character(0)
while (i <= length(s)) {
if (i < length(s) && s[i] == best[1] && s[i+1] == best[2]) {
out <- c(out, paste0(best[1], best[2])); i <- i + 2
} else { out <- c(out, s[i]); i <- i + 1 }
}
out })
}
list(merges = merges, vocab = sort(unique(unlist(splits))))
}
corpus <- paste("the patient presented with progressive dyspnea and fatigue",
"the patient reported no chest pain the patient was afebrile",
"progressive dyspnea on exertion progressive fatigue reported")
bpe <- bpe_learn(corpus, n_merges = 40)
c(characters = length(unique(unlist(strsplit(gsub("\\s", "", corpus), "")))),
words = length(unique(unlist(strsplit(corpus, "\\s+")))),
bpe_vocabulary = length(bpe$vocab))#> characters words bpe_vocabulary
#> 21 16 32#> [1] "e + </w>" "r + e" "t + i" "a + ti" "d + </w>" "e + n" "en + t"
#> [8] "re + s"The first merges combine the most frequent character pairs; frequent whole words emerge as single tokens after enough merges, while rare ones stay decomposed.
# The vocabulary-size / sequence-length trade, and its effect on attention cost
vs <- data.frame(
scheme = c("character", "BPE 8k", "BPE 32k", "BPE 100k", "word (100k)"),
vocab = c(100, 8000, 32000, 100000, 100000),
tokens_per_1000_words = c(5200, 1450, 1300, 1180, 1000),
handles_unseen = c(TRUE, TRUE, TRUE, TRUE, FALSE))
vs$embedding_params_M <- round(vs$vocab * 512 / 1e6, 2)
vs$attention_cost_relative <- round((vs$tokens_per_1000_words /
vs$tokens_per_1000_words[3])^2, 2)
vsggplot(vs[vs$handles_unseen, ], aes(vocab, attention_cost_relative)) +
geom_line(linewidth = 1, color = "steelblue") +
geom_point(aes(size = embedding_params_M), color = "firebrick") +
scale_x_log10() +
scale_size_continuous(name = "Embedding\nparams (M)", range = c(2, 7)) +
labs(title = "Vocabulary size trades embedding parameters against attention cost",
subtitle = "Attention is quadratic in token count, so shorter sequences are worth paying embedding parameters for",
x = "Vocabulary size (log scale)",
y = "Relative attention cost (BPE 32k = 1)") +
theme_dspa()
14.56 Case study: character-level generation
A decoder-only Transformer trained to predict the next character, then sampled autoregressively. Character level keeps the vocabulary small enough that the whole example runs on a laptop.
# A public-domain text; falls back to a synthetic corpus if unavailable
raw_text <- dspa_try({
txt <- dspa_read("https://www.gutenberg.org/files/11/11-0.txt",
"alice.txt", reader = readLines, warn = FALSE)
paste(txt, collapse = "\n")
}, fallback = paste(rep(corpus, 400), collapse = " "),
label = "Project Gutenberg text")
clean <- gsub("[^a-zA-Z0-9 .,;:'?!\n-]", "", raw_text)
clean <- substr(clean, 1, 120000)
chars <- sort(unique(strsplit(clean, "")[[1]]))
stoi <- setNames(seq_along(chars), chars)
c(corpus_characters = nchar(clean), vocabulary_size = length(chars))#> corpus_characters vocabulary_size
#> 120000 64char_transformer <- nn_module("CharTransformer",
initialize = function(vocab, d_model = 128, n_heads = 4, n_layers = 4,
max_len = 128, dropout = 0.1) {
self$max_len <- max_len
self$tok <- nn_embedding(vocab, d_model)
self$pos <- nn_embedding(max_len, d_model) # learned positional embedding
self$blocks <- nn_module_list(lapply(1:n_layers, \(i)
transformer_block(d_model, n_heads, dropout = dropout, pre_norm = TRUE)))
self$ln_f <- nn_layer_norm(d_model)
self$head <- nn_linear(d_model, vocab)
},
forward = function(idx) {
T_len <- idx$size(2)
pos_ids <- torch_arange(1, T_len, dtype = torch_long())$unsqueeze(1)
x <- self$tok(idx) + self$pos(pos_ids)
mask <- causal_mask(T_len)$to(device = idx$device) # no peeking ahead
for (i in seq_along(self$blocks)) x <- self$blocks[[i]](x, mask)
self$head(self$ln_f(x))
})block_size <- 96; batch_size <- 48
data_ids <- torch_tensor(unname(stoi[strsplit(clean, "")[[1]]]), dtype = torch_long())
n_all <- length(data_ids)
# CHRONOLOGICAL split: a language model must not be validated on text it saw
split_at <- floor(0.9 * n_all)
train_ids <- data_ids[1:split_at]; valid_ids <- data_ids[(split_at + 1):n_all]
get_batch <- function(ids) {
ix <- sample(length(ids) - block_size - 1, batch_size)
x <- torch_stack(lapply(ix, \(i) ids[i:(i + block_size - 1)]))
y <- torch_stack(lapply(ix, \(i) ids[(i + 1):(i + block_size)]))
list(x = x, y = y)
}
model_lm <- char_transformer(length(chars), max_len = block_size)
opt_lm <- optim_adamw(model_lm$parameters, lr = 3e-4, weight_decay = 0.1)
n_steps <- 3000
sched_lm <- lr_one_cycle(opt_lm, max_lr = 1e-3, total_steps = n_steps,
pct_start = 0.1) # warmup then cosine (Part 1, 14.12)
hist_lm <- data.frame()
for (s in seq_len(n_steps)) {
model_lm$train(); b <- get_batch(train_ids)
opt_lm$zero_grad()
logits <- model_lm(b$x)
l <- nnf_cross_entropy(logits$reshape(c(-1, length(chars))), b$y$reshape(-1))
l$backward()
nn_utils_clip_grad_norm_(model_lm$parameters, 1.0) # clipping (Section 14.46)
opt_lm$step(); sched_lm$step()
if (s %% 100 == 0) {
model_lm$eval()
vb <- get_batch(valid_ids)
vl <- with_no_grad(as.numeric(nnf_cross_entropy(
model_lm(vb$x)$reshape(c(-1, length(chars))), vb$y$reshape(-1))))
hist_lm <- rbind(hist_lm, data.frame(step = s, train = as.numeric(l), valid = vl))
}
}
lm_result <- dspa_cache("charlm", list(model = model_lm, history = hist_lm))14.56.1 Sampling strategies
Generation samples one token at a time from the predicted distribution. How it samples changes the output more than the model does.
\[ \begin{aligned} \textbf{Greedy: }&\quad y_t=\arg\max_v p(v)\\ \textbf{Temperature: }&\quad p_\tau(v)\propto\exp\big(\log p(v)/\tau\big)\\ \textbf{Top-}k\textbf{: }&\quad \text{renormalize over the }k\text{ most probable}\\ \textbf{Nucleus (top-}p\textbf{): }&\quad \text{renormalize over the smallest set with }\textstyle\sum p\ge p_0 \end{aligned} \]
Temperature is not a confidence knob. \(\tau\to0\) recovers greedy decoding and \(\tau\to\infty\) gives the uniform distribution; the model’s parameters are unchanged either way, so the “confidence” being adjusted is the sampler’s, not the model’s.
Greedy decoding is not the safe default it appears to be. It maximizes each token’s probability independently, which does not maximize the sequence probability, and it produces degenerate repetition, text that loops on a phrase indefinitely (Holtzman et al., 2020). Nucleus sampling adapts the candidate set to the distribution’s shape: where the model is confident the nucleus is small and sampling is nearly greedy; where it is uncertain the nucleus widens. Fixed top-\(k\) cannot do this, because \(k\) is the same whether the distribution is peaked or flat.
sample_next <- function(logits, temperature = 1.0, top_k = NULL, top_p = NULL) {
logits <- logits / temperature
if (!is.null(top_k)) {
kth <- torch_topk(logits, k = min(top_k, logits$size(-1)))[[1]][-1]
logits <- logits$masked_fill(logits < kth, -Inf)
}
if (!is.null(top_p)) {
srt <- torch_sort(logits, descending = TRUE)
probs_sorted <- nnf_softmax(srt[[1]], dim = -1)
cum <- torch_cumsum(probs_sorted, dim = -1)
# Keep the smallest prefix whose cumulative mass reaches top_p
cutoff_idx <- as.integer(torch_sum((cum < top_p)$to(dtype = torch_int()))) + 1
thresh <- srt[[1]][cutoff_idx]
logits <- logits$masked_fill(logits < thresh, -Inf)
}
as.integer(torch_multinomial(nnf_softmax(logits, dim = -1), 1))
}
generate <- function(model, prompt, n_new = 300, ...) {
model$eval()
idx <- torch_tensor(matrix(unname(stoi[strsplit(prompt, "")[[1]]]), nrow = 1),
dtype = torch_long())
for (i in seq_len(n_new)) {
ctx <- if (idx$size(2) > model$max_len)
idx[, (idx$size(2) - model$max_len + 1):idx$size(2)] else idx
logits <- with_no_grad(model(ctx))[1, -1, ]
nxt <- sample_next(logits, ...)
idx <- torch_cat(list(idx, torch_tensor(matrix(nxt, nrow = 1),
dtype = torch_long())), dim = 2)
}
paste(chars[as.integer(idx[1, ])], collapse = "")
}# How each strategy reshapes a distribution -- illustrated on a synthetic one
torch_manual_seed(101)
V_demo <- 60
base_logits <- torch_randn(V_demo) * 2
base_logits[5] <- 6; base_logits[12] <- 5.2; base_logits[30] <- 4.8
reshape_probs <- function(temperature = 1, top_k = NULL, top_p = NULL) {
lg <- base_logits / temperature
if (!is.null(top_k)) {
kth <- torch_topk(lg, top_k)[[1]][-1]
lg <- lg$masked_fill(lg < kth, -Inf)
}
if (!is.null(top_p)) {
srt <- torch_sort(lg, descending = TRUE)
cum <- torch_cumsum(nnf_softmax(srt[[1]], dim = -1), dim = -1)
ci <- as.integer(torch_sum((cum < top_p)$to(dtype = torch_int()))) + 1
lg <- lg$masked_fill(lg < srt[[1]][ci], -Inf)
}
as.numeric(nnf_softmax(lg, dim = -1))
}
strat <- bind_rows(
data.frame(token = 1:V_demo, p = reshape_probs(), s = "temperature 1.0"),
data.frame(token = 1:V_demo, p = reshape_probs(temperature = 0.5), s = "temperature 0.5"),
data.frame(token = 1:V_demo, p = reshape_probs(temperature = 1.5), s = "temperature 1.5"),
data.frame(token = 1:V_demo, p = reshape_probs(top_k = 5), s = "top-k = 5"),
data.frame(token = 1:V_demo, p = reshape_probs(top_p = 0.9), s = "nucleus p = 0.9"))
ggplot(strat, aes(token, p)) +
geom_col(fill = "steelblue", width = 0.9) +
facet_wrap(~ factor(s, levels = unique(strat$s)), nrow = 1) +
labs(title = "The same model output under five sampling strategies",
subtitle = "Low temperature sharpens; high temperature flattens; top-k and nucleus truncate the tail",
x = "Token", y = "Probability") +
theme_dspa(9)
strat |> summarise(entropy = -sum(p[p > 0] * log(p[p > 0])),
effective_vocab = round(exp(entropy), 1),
n_with_mass = sum(p > 1e-6), .by = s) |>
mutate(entropy = round(entropy, 3))# How temperature reshapes the whole distribution, by token rank
tau_grid <- 10^seq(-0.7, 0.7, length.out = 40)
ranked <- order(as.numeric(base_logits), decreasing = TRUE)
Ztau <- t(vapply(tau_grid, function(tau)
as.numeric(nnf_softmax(base_logits / tau, dim = -1))[ranked][1:30],
numeric(30)))
plot_ly(x = 1:30, y = tau_grid, z = Ztau, type = "surface",
colorscale = "Inferno", reversescale = TRUE,
colorbar = list(title = "Probability")) |>
layout(title = "Sampling distribution over token rank and temperature",
scene = list(xaxis = list(title = "Token rank"),
yaxis = list(title = "Temperature", type = "log"),
zaxis = list(title = "Probability")))Rotate to the low-temperature edge: the mass collapses onto rank 1 and the sampler becomes greedy. At the high-temperature edge the surface flattens toward uniform, and the model’s ranking stops mattering. Neither extreme is useful, and nothing about the model changed between them.
The effective_vocab column (\(e^{H}\), the perplexity of the sampling
distribution) is the useful summary: it says how many tokens are
genuinely in play at each step.
prompt <- "The patient "
data.frame(
strategy = c("greedy (tau -> 0)", "temperature 0.8", "top-k = 10",
"nucleus p = 0.9"),
sample = c(substr(generate(lm_result$model, prompt, 200, temperature = 0.01), 1, 160),
substr(generate(lm_result$model, prompt, 200, temperature = 0.8), 1, 160),
substr(generate(lm_result$model, prompt, 200, top_k = 10), 1, 160),
substr(generate(lm_result$model, prompt, 200, top_p = 0.9), 1, 160)))Common misconception: “perplexity lets you compare language models.” Only within a fixed tokenization. Perplexity is \(\exp(\text{cross-entropy})\), so a perplexity of 12 means the model is as uncertain as if choosing uniformly among 12 characters. It is comparable only within a tokenization: a character-level perplexity of 3 and a word-level perplexity of 60 are not on the same scale, since they count different things. Report the tokenizer alongside the number, or the number means nothing.
14.57 Case study: forecasting, neural versus classical
Chapter 12, §12.29 argued that neural forecasters do not beat classical methods by default and promised the Transformer comparison. Here it is, under the same protocol: same origins, same horizons, same scale-free metric, benchmarks included.
Mean Absolute Scaled Error (MASE): Rob Hyndman and Anne Koehler introduced the MASE metric to evaluate and compare the accuracy of forecasting models. In its standard form, MASE applies solely to time-series or sequential (longitudinal) predictions. MASE represents a standard way to measure forecast errors without being skewed by scale or specific data units. Instead of looking at raw errors, MASE scales the error of the model predictions against the in-sample mean absolute error of a simple baseline, which is typically a naive or seasonal naive forecast. MASE interpretation:
- \(MASE < 1\): The model performs better than the naive baseline.
- \(MASE = 1\): The model performs equally to the naive baseline.
- \(MASE > 1\): The model performs worse than a simple naive baseline.
\[\text{MASE} = \frac{\frac{1}{h} \sum_{t=n+1}^{n+h} \vert{}Y_t - \hat{Y}_t\vert{}}{\frac{1}{n-m} \sum_{i=m+1}^{n} \vert{}Y_i - Y_{i-m}\vert{}},\] where
- \(Y_t\): The actual value of the time series at time \(t\).
- \(\hat{Y}_t\): The forecasted value of the time series at time \(t\).
- \(n\): The number of observations in the training (in-sample) dataset.
- \(h\): The forecasting horizon (number of steps ahead being forecasted).
- \(m\): The seasonal period (\(m = 1\) for non-seasonal data; for monthly data with yearly seasonality, \(m = 12\)).
The numerator calculates the Mean Absolute Error (MAE) of the forecasting model over the out-of-sample period (\(n+1\) to \(n+h\)). While the denominator calculates the in-sample Mean Absolute Error of a naive or seasonal naive baseline model (predicting based on the previous period or the same period last cycle). Splitting the model’s error by this baseline error normalizes the result, making MASE independent of the scale of the data.
set.seed(111)
T_all <- 1200
t_idx <- seq_len(T_all)
# A seasonal series with trend, an AR component, and a non-linear regime effect
seasonal <- 12 * sin(2*pi*t_idx/24) + 6 * sin(2*pi*t_idx/168)
trend <- 0.012 * t_idx
ar_part <- as.numeric(arima.sim(list(ar = c(0.55, 0.2)), T_all, sd = 2.2))
regime <- 4 * (sin(2*pi*t_idx/300) > 0.5) # non-linearity ARIMA cannot see
y_all <- 50 + trend + seasonal + ar_part + regime
h_out <- 120
y_train <- y_all[1:(T_all - h_out)]; y_test <- y_all[(T_all - h_out + 1):T_all]
c(total = T_all, train = length(y_train), test = h_out, season = 24)#> total train test season
#> 1200 1080 120 24library(forecast)
ts_train <- ts(y_train, frequency = 24)
fits <- list(
`Seasonal naive` = snaive(ts_train, h = h_out),
`ETS` = forecast(ets(ts_train), h = h_out),
`ARIMA (auto)` = forecast(auto.arima(ts_train, seasonal = TRUE,
stepwise = TRUE), h = h_out))# A small Transformer forecaster. Scaling uses TRAINING statistics only.
mu_y <- mean(y_train); sd_y <- sd(y_train)
z_train <- (y_train - mu_y) / sd_y
lookback <- 168; horizon <- 24
make_windows <- function(z, lb, h) {
n <- length(z) - lb - h + 1
X <- t(vapply(seq_len(n), \(i) z[i:(i + lb - 1)], numeric(lb)))
Y <- t(vapply(seq_len(n), \(i) z[(i + lb):(i + lb + h - 1)], numeric(h)))
list(X = torch_tensor(X)$unsqueeze(3), Y = torch_tensor(Y))
}
w <- make_windows(z_train, lookback, horizon)
ts_transformer <- nn_module("TSTransformer",
initialize = function(lb, h, d_model = 64, n_heads = 4, n_layers = 3) {
self$inp <- nn_linear(1, d_model)
self$pos <- nn_embedding(lb, d_model)
self$blocks <- nn_module_list(lapply(1:n_layers, \(i)
transformer_block(d_model, n_heads, pre_norm = TRUE)))
self$ln <- nn_layer_norm(d_model)
self$head <- nn_linear(d_model, h)
},
forward = function(x) {
pos_ids <- torch_arange(1, x$size(2), dtype = torch_long())$unsqueeze(1)
z <- self$inp(x) + self$pos(pos_ids)
for (i in seq_along(self$blocks)) z <- self$blocks[[i]](z)
self$head(self$ln(z)[, -1, ]) # predict from the last position
})
set.seed(113); torch_manual_seed(113)
net_ts <- ts_transformer(lookback, horizon)
opt_ts <- optim_adamw(net_ts$parameters, lr = 1e-3, weight_decay = 1e-4)
n_win <- w$X$size(1)
for (e in 1:60) {
net_ts$train()
perm <- sample(n_win)
for (s in seq(1, n_win - 63, by = 64)) {
idx <- perm[s:(s + 63)]
opt_ts$zero_grad()
l <- nnf_mse_loss(net_ts(w$X[idx, , , drop = FALSE]), w$Y[idx, ])
l$backward(); opt_ts$step()
}
}
# Recursive multi-step forecast to the full horizon
net_ts$eval()
hist_z <- z_train
preds_z <- numeric(0)
while (length(preds_z) < h_out) {
ctx <- torch_tensor(matrix(tail(hist_z, lookback), nrow = 1))$unsqueeze(3)
step <- as.numeric(with_no_grad(net_ts(ctx)))
preds_z <- c(preds_z, step); hist_z <- c(hist_z, step)
}
pred_nn <- preds_z[1:h_out] * sd_y + mu_y
nn_result <- dspa_cache("ts_transformer_pred", pred_nn)pred_nn <- dspa_try(dspa_cache("ts_transformer_pred", NULL),
fallback = NULL, label = "cached Transformer forecast")
scale_denom <- mean(abs(diff(y_train, lag = 24))) # seasonal-naive scaling
mase <- function(p) mean(abs(p - y_test)) / scale_denom
rmse <- function(p) sqrt(mean((p - y_test)^2))
rows <- lapply(names(fits), \(nm)
data.frame(model = nm, MASE = mase(as.numeric(fits[[nm]]$mean)),
RMSE = rmse(as.numeric(fits[[nm]]$mean))))
if (!is.null(pred_nn))
rows <- c(rows, list(data.frame(model = "Transformer", MASE = mase(pred_nn),
RMSE = rmse(pred_nn))))
res_fc <- bind_rows(rows) |>
mutate(beats_seasonal_naive = MASE < MASE[model == "Seasonal naive"],
across(c(MASE, RMSE), \(z) round(z, 4))) |>
arrange(MASE)
res_fcRead beats_seasonal_naive before anything
else. A model that cannot beat the seasonal naive forecast is
not adding value, whatever its RMSE, the argument of Chapter
12, §12.13.
fc_df <- bind_rows(lapply(names(fits), \(nm)
data.frame(t = 1:h_out, v = as.numeric(fits[[nm]]$mean), model = nm)))
if (!is.null(pred_nn))
fc_df <- bind_rows(fc_df, data.frame(t = 1:h_out, v = pred_nn, model = "Transformer"))
ggplot() +
geom_line(data = data.frame(t = 1:h_out, v = y_test),
aes(t, v), color = "grey25", linewidth = 0.8) +
geom_line(data = fc_df, aes(t, v, color = model), linewidth = 0.7) +
scale_color_viridis_d(option = "turbo", end = 0.9) +
labs(title = "Held-out forecasts against the observed series",
subtitle = "Grey: actuals. All models forecast the same 120 steps from the same origin",
x = "Steps ahead", y = "Value", color = NULL) +
theme_dspa()
# One origin is one draw. Rolling origin is the honest comparison.
origins <- round(seq(700, T_all - 150, length.out = 8))
roll <- bind_rows(lapply(origins, function(o) {
tr <- ts(y_all[1:o], frequency = 24); te <- y_all[(o+1):(o+48)]
sd_o <- mean(abs(diff(y_all[1:o], lag = 24)))
sn <- as.numeric(snaive(tr, h = 48)$mean)
ar <- as.numeric(forecast(auto.arima(tr, seasonal = TRUE, stepwise = TRUE),
h = 48)$mean)
et <- as.numeric(forecast(ets(tr), h = 48)$mean)
data.frame(origin = o,
`Seasonal naive` = mean(abs(sn - te))/sd_o,
`ARIMA (auto)` = mean(abs(ar - te))/sd_o,
ETS = mean(abs(et - te))/sd_o,
check.names = FALSE)
}))
roll |> summarise(across(-origin, list(mean = mean, sd = sd))) |>
pivot_longer(everything()) |> mutate(value = round(value, 4))roll |> pivot_longer(-origin, names_to = "model", values_to = "MASE") |>
ggplot(aes(factor(origin), MASE, fill = model)) +
geom_col(position = "dodge") +
geom_hline(yintercept = 1, linetype = "dashed", color = "grey30") +
scale_fill_viridis_d(option = "turbo", end = 0.85) +
labs(title = "MASE at eight forecast origins",
subtitle = "Dashed line at 1 is the in-sample seasonal-naive benchmark. Rankings are not stable across origins",
x = "Forecast origin", y = "MASE", fill = NULL) +
theme_dspa(10)
Common misconception: “a Transformer will forecast better than ARIMA.” On a single univariate series of a few thousand points, it usually will not. In the M4 and M5 competitions the winning entries were statistical methods and hybrids that embedded exponential smoothing inside a neural architecture — not pure neural models (Makridakis et al., 2020).
The reasons are structural. A Transformer has tens of thousands of parameters and a few thousand training windows drawn from one realization, so it is in the small-\(n\), large-\(p\) regime of Chapter 11, §11.3. It has no built-in notion of trend or seasonality, ARIMA and ETS encode both as structural assumptions that are correct for many series. And it produces point forecasts, where ARIMA supplies prediction intervals analytically (Chapter 12, §12.12).
Neural forecasters earn their cost with many related series, a global model trained across thousands of series can learn shared structure no single-series method sees, which is what DeepAR and N-BEATS exploit and what won the M5 accuracy track. On one series, use ARIMA and report the benchmark.
And note the rolling-origin figure: the ranking changes across origins. A single split would have supported whichever conclusion its origin happened to favour.
14.58 Pretraining and transfer for text
Vision transfer learning (Part 2, §14.35) starts from a network trained on labelled images. Text pretraining needs no labels at all: the objective is constructed from the text itself.
\[ \begin{aligned} \textbf{Causal LM (GPT): }&\quad \max\ \textstyle\sum_t\log p\big(x_t\mid x_{<t}\big) &&\text{decoder-only, causal mask}\\ \textbf{Masked LM (BERT): }&\quad \max\ \textstyle\sum_{t\in\mathcal M}\log p\big(x_t\mid x_{\setminus\mathcal M}\big) &&\text{encoder-only, bidirectional}\\ \textbf{Seq2seq (T5): }&\quad \max\ \log p\big(\text{span}\mid \text{corrupted context}\big) &&\text{encoder--decoder} \end{aligned} \]
The distinction that matters downstream is what each can see. A causal model conditions only on the left context, which makes it a generator. A masked model conditions on both sides, which makes it a better encoder for classification but unable to generate autoregressively.
| Choose | When |
|---|---|
| Encoder-only (BERT-family) | Classification, NER, retrieval, the whole input is available at once |
| Decoder-only (GPT-family) | Generation, and increasingly everything, via prompting |
| Encoder–decoder (T5, BART) | Translation, summarization, input and output differ in length and language |
# The two objectives, on the same batch, to make the difference concrete
torch_manual_seed(121)
V_p <- 40; T_p <- 12; B_p <- 4
seq_ids <- torch_randint(1, V_p + 1, size = c(B_p, T_p), dtype = torch_long())
# Causal: target is the input shifted by one; the mask forbids looking ahead
causal_target <- seq_ids[, 2:T_p]
# Masked: replace 15% of positions with a [MASK] id and predict only those
mask_prob <- 0.15
mask_pos <- torch_rand(B_p, T_p) < mask_prob
masked_input <- seq_ids$clone()
masked_input[mask_pos] <- V_p + 1 # the [MASK] token id
data.frame(
objective = c("causal LM", "masked LM"),
positions_predicted_per_sequence = c(T_p - 1,
round(as.numeric(torch_sum(mask_pos$to(dtype = torch_float()))) / B_p, 1)),
sees_right_context = c(FALSE, TRUE),
can_generate_autoregressively = c(TRUE, FALSE))The causal objective supplies a training signal at every position, where the masked objective supplies one at 15% of them, so masked models need more data per unit of learning. That efficiency, plus the direct fit to generation, is a large part of why decoder-only architectures came to dominate.
# Fine-tuning a small pretrained encoder for classification, with the
# freeze/unfreeze trade of Part 2 Section 14.35 measured again on text
set.seed(123); torch_manual_seed(123)
V_t <- 60; T_t <- 24
make_text_task <- function(n, signal_tokens = c(3, 7, 11)) {
X <- torch_randint(1, V_t + 1, size = c(n, T_t), dtype = torch_long())
has <- vapply(1:n, \(i) any(as.integer(X[i, ]) %in% signal_tokens), logical(1))
list(X = X, y = torch_tensor(matrix(as.numeric(has), ncol = 1)))
}
pre <- make_text_task(3000) # "pretraining" corpus
encoder <- nn_module("Enc",
initialize = function() {
self$emb <- nn_embedding(V_t + 1, 48)
self$pos <- nn_embedding(T_t, 48)
self$b1 <- transformer_block(48, 4); self$b2 <- transformer_block(48, 4)
self$ln <- nn_layer_norm(48)
},
forward = function(idx) {
p <- torch_arange(1, idx$size(2), dtype = torch_long())$unsqueeze(1)
z <- self$emb(idx) + self$pos(p)
self$ln(self$b2(self$b1(z)))
})()
head_t <- nn_linear(48, 1)
opt_p <- optim_adamw(c(encoder$parameters, head_t$parameters), lr = 2e-3)
for (e in 1:150) {
opt_p$zero_grad()
z <- encoder(pre$X)$mean(dim = 2)
nnf_binary_cross_entropy_with_logits(head_t(z), pre$y)$backward()
opt_p$step()
}
adapt <- function(n_target, freeze_body) {
set.seed(125); torch_manual_seed(125)
tgt <- make_text_task(n_target, signal_tokens = c(3, 7, 19)) # shifted task
te <- make_text_task(600, signal_tokens = c(3, 7, 19))
enc2 <- encoder$clone(deep = TRUE)
h2 <- nn_linear(48, 1)
if (freeze_body) { for (p in enc2$parameters) p$requires_grad_(FALSE); enc2$eval() }
ps <- if (freeze_body) h2$parameters else c(enc2$parameters, h2$parameters)
o <- optim_adamw(ps, lr = 1e-3)
for (e in 1:120) { o$zero_grad()
nnf_binary_cross_entropy_with_logits(h2(enc2(tgt$X)$mean(dim = 2)), tgt$y)$backward()
o$step() }
enc2$eval()
mean((as.numeric(with_no_grad(torch_sigmoid(h2(enc2(te$X)$mean(dim = 2))))) > 0.5) ==
as.numeric(te$y))
}
tt <- expand.grid(n_target = c(50, 150, 500, 1500), frozen = c(TRUE, FALSE))
tt$accuracy <- mapply(adapt, tt$n_target, tt$frozen)
tt |> mutate(strategy = ifelse(frozen, "frozen body", "full fine-tune"),
accuracy = round(accuracy, 4)) |>
select(n_target, strategy, accuracy) |>
pivot_wider(names_from = strategy, values_from = accuracy)The same pattern as vision: freezing wins at small target sizes, full fine-tuning wins once there is enough target data to fit the body without destroying it.
14.59 What a large language model is
Scaling the causal-LM objective produces systems whose behavior looks qualitatively different from their training objective, which invites imprecise description. The precise version is short.
A large language model is a decoder-only Transformer trained to predict the next token, at scale. No component of that sentence changes as the model grows, the architecture of §14.53, the objective of §14.58, the sampler of §14.56.1. What changes is parameter count, data volume, and compute, and empirical scaling laws describe how loss falls with each (Kaplan et al., 2020; Hoffmann et al., 2022): \[\mathcal L(N,D)\approx \frac{A}{N^{\alpha}}+\frac{B}{D^{\beta}}+\mathcal L_\infty,\] with \(N\) parameters and \(D\) training tokens. The Chinchilla result is that compute-optimal training scales \(N\) and \(D\) together, many large models were badly undertrained for their size.
Common misconception: “the model understands the text.” What is established is that it assigns probabilities to token sequences, and that those probabilities encode a great deal of syntactic, factual, and procedural regularity. What is not established by the training objective is any internal representation of truth, reference, or intent, the objective rewards plausibility, and a fluent falsehood is exactly as probable as a fluent fact if the training corpus made it so.
Three practical consequences follow, and they are the ones that matter for using these systems in an analysis. Confident errors are the expected failure mode, not an aberration, the model is not tracking a truth value it could report uncertainty about. The output distribution reflects the training corpus, including its biases and its cutoff date. And calibration is a separate question from accuracy, exactly as in Chapter 9, §9.8: a model can be accurate and badly calibrated, or fluent and neither.
Instruction tuning and RLHF change what the model is likely to say; they do not change what the objective is optimizing.
# The Chinchilla-style trade: for a fixed compute budget, N and D must balance
compute_budget <- 1e21 # FLOPs
N_grid <- 10^seq(8, 11.5, length.out = 120)
D_grid <- compute_budget / (6 * N_grid) # C ~ 6ND
A <- 406.4; alpha <- 0.34; B <- 410.7; beta <- 0.28; L_inf <- 1.69
loss <- A / N_grid^alpha + B / D_grid^beta + L_inf
opt_i <- which.min(loss)
ggplot(data.frame(N = N_grid, loss = loss), aes(N, loss)) +
geom_line(linewidth = 1, color = "steelblue") +
geom_point(aes(x = N_grid[opt_i], y = loss[opt_i]), color = "firebrick",
size = 3.5) +
scale_x_log10() +
annotate("text", x = N_grid[opt_i], y = loss[opt_i] * 1.02, vjust = -0.8,
size = 3.2, color = "firebrick",
label = sprintf("optimum: %.1e params, %.1e tokens",
N_grid[opt_i], D_grid[opt_i])) +
labs(title = "Loss against parameter count at a fixed compute budget",
subtitle = "Too few parameters underfits; too many leaves too few tokens for the budget. The minimum is interior",
x = "Parameters (log scale)", y = "Predicted loss") +
theme_dspa()
c(optimal_params = signif(N_grid[opt_i], 3),
optimal_tokens = signif(D_grid[opt_i], 3),
tokens_per_param = round(D_grid[opt_i] / N_grid[opt_i], 1))#> optimal_params optimal_tokens tokens_per_param
#> 1.84e+09 9.06e+10 4.93e+01The minimum is interior, which is the whole point: at a fixed budget, making the model bigger is not free, it buys parameters by spending tokens.
PART IV: SYNTHESIS
14.60 Complexity summary
\(T\) sequence length, \(d\) model dimension, \(h\) heads, \(L\) layers, \(B\) batch, \(V\) vocabulary.
| Component | Parameters | Time | Memory | Sequential ops |
|---|---|---|---|---|
| RNN cell | \(d(d{+}m){+}d\) | \(O(BTd^2)\) | \(O(BTd)\) | \(O(T)\) |
| GRU cell | \(3d(d{+}m){+}3d\) | \(O(BTd^2)\) | \(O(BTd)\) | \(O(T)\) |
| LSTM cell | \(4d(d{+}m){+}4d\) | \(O(BTd^2)\) | \(O(BTd)\) | \(O(T)\) |
| Self-attention | \(4d^2\) | \(O(BT^2d)\) | \(\mathbf{O(BhT^2)}\) | \(\mathbf{O(1)}\) |
| Multi-head (any \(h\)) | \(4d^2\) | \(O(BT^2d)\) | \(O(BhT^2)\) | \(O(1)\) |
| FFN (\(4d\) inner) | \(\mathbf{8d^2}\) | \(O(BTd^2)\) | \(O(BTd)\) | \(O(1)\) |
| Transformer block | \(12d^2\) | \(O(B(T^2d{+}Td^2))\) | \(O(B(hT^2{+}Td))\) | \(O(1)\) |
| Embedding + output | \(2Vd\) | \(O(BTd)\) | \(O(Vd)\) | — |
| Autoregressive generation | — | \(O(T^2d)\) per token naive | \(O(Td)\) with KV cache | \(O(T)\) |
Five consequences.
Attention beats recurrence on wall time despite worse asymptotics in \(T\), because it has \(O(1)\) sequential operations against \(O(T)\). On parallel hardware that trade is decisive; on a device processing a stream one element at a time it reverses.
The FFN holds two-thirds of a block’s parameters, \(8d^2\) against attention’s \(4d^2\), so parameter-reduction work targets it first.
Attention memory is \(O(hT^2)\) and is usually the binding constraint. At \(T=8192\) with 16 heads, the attention matrices alone need over 4 GB per layer at float32. FlashAttention avoids materializing them; sparse and linear attention change the asymptotics.
Multi-head is free in parameters. \(h\) does not appear in the parameter count; it only partitions \(d\).
Generation is \(O(T)\) sequential regardless of architecture. The key–value cache reduces per-token work from \(O(T^2d)\) to \(O(Td)\), but tokens must still be produced one at a time, which is why inference latency, not training cost, dominates deployment.
14.61 Common pitfalls
| # | Pitfall | Consequence | Fix |
|---|---|---|---|
| 1 | Treating exploding and vanishing gradients as one problem | Clipping fixes only the first | Clipping for explosion; gating or attention for vanishing |
| 2 | Reading a healthy loss curve as evidence of long-range learning | Vanishing gradients are silent | Measure gradient magnitude by lag |
| 3 | Leaving the LSTM forget-gate bias at zero | \(f\approx0.5\) halves the gradient per step | Initialize the forget bias to \(\approx1\) |
| 4 | Expecting gating to solve long-range dependence | It extends reach to hundreds of steps, not thousands | Attention for genuinely long contexts |
| 5 | Encoder–decoder without attention on long inputs | The fixed vector is a hard capacity limit | Attention, or a Transformer |
| 6 | Dropping the \(1/\sqrt{d_k}\) | Softmax saturates; the gradient vanishes | Scale by \(\sqrt{d_k}\) |
| 7 | Masking after the softmax | Remaining weights no longer sum to 1 | Set logits to \(-\infty\) before the softmax |
| 8 | Believing more heads adds capacity | Parameters are \(4d^2\) for any \(h\) | Heads partition \(d\); they buy subspaces |
| 9 | Omitting positional encoding | Attention is permutation-equivariant; order is invisible | Sinusoidal, learned, RoPE, or ALiBi |
| 10 | Learned absolute positions, then a longer input | Cannot extrapolate past the trained length | Sinusoidal, RoPE, or ALiBi |
| 11 | Post-norm at depth without warmup | Gradients poorly scaled; training diverges | Pre-norm, or a careful warmup |
| 12 | BatchNorm in a sequence model | Variable lengths and small batches break it | LayerNorm |
| 13 | Ignoring the \(O(T^2)\) attention memory | Out-of-memory well before FLOPs bind | FlashAttention; sparse or linear attention |
| 14 | Assuming tokens are words | Token budgets are not word budgets | Report the tokenizer with any length or cost |
| 15 | Comparing perplexity across tokenizations | The quantities count different things | Comparable only within a tokenizer |
| 16 | Greedy decoding as the safe default | Degenerate repetition; does not maximize sequence probability | Nucleus sampling |
| 17 | Reading temperature as model confidence | It reshapes the sampler, not the model | The parameters are unchanged |
| 18 | Training with teacher forcing, deploying autoregressively | Exposure bias: the model never saw its own errors | Scheduled sampling; evaluate autoregressively |
| 19 | Randomly splitting a text or time series | Trains on the future | Chronological split |
| 20 | Scaling a series with statistics from the full record | Test information leaks into training | Training statistics only |
| 21 | Reporting a neural forecast without benchmarks | An uninterpretable number | Seasonal naive, ETS, ARIMA; MASE |
| 22 | One forecast origin | The ranking changes across origins | Rolling-origin evaluation |
| 23 | Expecting a Transformer to beat ARIMA on one series | Small-\(n\), large-\(p\); no trend or seasonality prior | Global models across many series, or classical methods |
| 24 | Reading fluency as understanding | The objective rewards plausibility, not truth | Confident errors are the expected failure mode |
14.62 Practice problems
Problem 1: Locate the spectral-radius threshold
Solution
rhos <- seq(0.5, 1.4, by = 0.1)
p1 <- data.frame(rho = rhos,
grad_at_lag_50 = vapply(rhos, \(r) grad_by_lag(r)[10], numeric(1)))
p1$log10_grad <- round(log10(pmax(p1$grad_at_lag_50, 1e-30)), 2)
p1 |> mutate(grad_at_lag_50 = signif(grad_at_lag_50, 3))fit_decay <- lm(log(pmax(grad_by_lag(0.8), 1e-30)) ~ I(60 - (1:60)))
c(fitted_log_decay_per_step = round(coef(fit_decay)[2], 4),
predicted_log_rho = round(log(0.8), 4))#> fitted_log_decay_per_step.I(60 - (1:60))
#> -0.2987
#> predicted_log_rho
#> -0.2231ggplot(p1, aes(rho, log10_grad)) +
geom_vline(xintercept = 1, linetype = "dashed", color = "firebrick") +
geom_line(linewidth = 1, color = "steelblue") + geom_point(size = 2.4) +
annotate("text", x = 1.02, y = min(p1$log10_grad), hjust = 0, size = 3.2,
color = "firebrick", label = "rho = 1") +
labs(title = "Gradient at lag 50 against the recurrent matrix's spectral radius",
subtitle = "The transition at rho = 1 spans many orders of magnitude in the gradient",
x = expression(rho(W[hh])), y = "log10 gradient norm") +
theme_dspa()
Problem 2: Break attention by removing the scaling
Solution
train_attn <- function(d_k, scaled, steps = 300, seed = 141) {
torch_manual_seed(seed)
T_len <- 16; B <- 32
X <- torch_randn(B, T_len, d_k)
target <- X[, 1, ]$clone() # retrieve position 1
Wq <- nn_linear(d_k, d_k, bias = FALSE); Wk <- nn_linear(d_k, d_k, bias = FALSE)
opt <- optim_adam(c(Wq$parameters, Wk$parameters), lr = 1e-3)
l0 <- NA
for (s in seq_len(steps)) {
opt$zero_grad()
sc <- torch_bmm(Wq(X), Wk(X)$transpose(2, 3))
if (scaled) sc <- sc / sqrt(d_k)
out <- torch_bmm(nnf_softmax(sc, dim = 3), X)[, 1, ]
l <- nnf_mse_loss(out, target)
if (s == 1) l0 <- as.numeric(l)
l$backward(); opt$step()
}
c(d_k = d_k, initial = l0, final = as.numeric(l),
fraction_reduced = 1 - as.numeric(l)/l0)
}
rbind(cbind(scaled = 0, t(vapply(c(16, 64, 256), train_attn, numeric(4), scaled = FALSE))),
cbind(scaled = 1, t(vapply(c(16, 64, 256), train_attn, numeric(4), scaled = TRUE)))) |>
as.data.frame() |>
mutate(scaled = ifelse(scaled == 1, "with 1/sqrt(d_k)", "unscaled"),
across(where(is.numeric), \(z) signif(z, 4)))Problem 3: Verify permutation equivariance, then break it
Solution
torch_manual_seed(151)
d_e <- 32; T_e <- 10
mh_e <- multi_head(d_e, 4)
Xe <- torch_randn(1, T_e, d_e)
pm <- sample(T_e)
with_no_grad({
a1 <- mh_e(Xe)$out[1, pm, ] # P . f(X)
a2 <- mh_e(Xe[, pm, , drop = FALSE])$out[1, , ] # f(P X)
# Now add positional encoding and repeat
PEe <- sinusoidal_pe(T_e, d_e)$unsqueeze(1)
b1 <- mh_e(Xe + PEe)$out[1, pm, ]
b2 <- mh_e(Xe[, pm, , drop = FALSE] + PEe)$out[1, , ]
})
data.frame(
configuration = c("no positional encoding", "with positional encoding"),
max_abs_difference = signif(c(as.numeric(torch_max(torch_abs(a1 - a2))),
as.numeric(torch_max(torch_abs(b1 - b2)))), 4),
equivariant = c(TRUE, FALSE))Problem 4: Find the attention/recurrence crossover
Solution
bench_layer <- function(T_len, d, reps = 5) {
torch_manual_seed(161)
x <- torch_randn(4, T_len, d)
mh <- multi_head(d, 4); rnncell <- nn_gru(d, d, batch_first = TRUE)
t_att <- median(replicate(reps,
system.time(with_no_grad(mh(x)))[["elapsed"]]))
t_rnn <- median(replicate(reps,
system.time(with_no_grad(rnncell(x)))[["elapsed"]]))
c(T_len = T_len, d = d, attention = t_att, recurrence = t_rnn,
ratio = t_att / t_rnn)
}
p4 <- as.data.frame(do.call(rbind, lapply(c(32, 64, 128, 256, 512),
bench_layer, d = 128)))
p4 |> mutate(across(c(attention, recurrence, ratio), \(z) signif(z, 3)),
attention_faster = ratio < 1)p4 |> select(T_len, attention, recurrence) |>
pivot_longer(-T_len, names_to = "layer", values_to = "sec") |>
ggplot(aes(T_len, sec, color = layer)) +
geom_vline(xintercept = 128, linetype = "dashed", color = "grey35") +
geom_line(linewidth = 1) + geom_point(size = 2.4) +
scale_x_log10(breaks = p4$T_len) + scale_y_log10() +
scale_color_manual(values = c(attention = "#D8433B", recurrence = "#3B7DD8")) +
annotate("text", x = 132, y = min(p4$attention), hjust = 0, size = 3.2,
color = "grey35", label = "T = d") +
labs(title = "Measured wall time per layer, d = 128",
subtitle = "FLOP asymptotics cross at T = d; measured time crosses later because attention parallelizes",
x = "Sequence length (log scale)", y = "Seconds (log scale)", color = NULL) +
theme_dspa()
Problem 5: Compare sampling strategies by repetition
Solution
# Repetition rate under each strategy, on a synthetic peaked distribution
torch_manual_seed(171)
simulate_gen <- function(n = 600, ...) {
toks <- integer(n)
lg <- torch_randn(30) * 1.5; lg[7] <- 5
for (i in seq_len(n)) {
toks[i] <- sample_next(lg, ...)
lg <- lg * 0.98; lg[toks[i]] <- lg[toks[i]] + 0.35 # mild self-reinforcement
}
toks
}
rep_rate <- function(x) mean(head(x, -1) == tail(x, -1))
uniq_frac <- function(x) length(unique(x)) / length(x)
data.frame(
strategy = c("greedy", "temperature 0.7", "temperature 1.0",
"top-k = 5", "nucleus 0.9"),
repetition_rate = round(c(
rep_rate(simulate_gen(temperature = 0.01)), rep_rate(simulate_gen(temperature = 0.7)),
rep_rate(simulate_gen(temperature = 1.0)), rep_rate(simulate_gen(top_k = 5)),
rep_rate(simulate_gen(top_p = 0.9))), 4),
distinct_token_fraction = round(c(
uniq_frac(simulate_gen(temperature = 0.01)), uniq_frac(simulate_gen(temperature = 0.7)),
uniq_frac(simulate_gen(temperature = 1.0)), uniq_frac(simulate_gen(top_k = 5)),
uniq_frac(simulate_gen(top_p = 0.9))), 4))Problem 6: Show that one forecast origin is not enough
Solution
wins <- roll |>
pivot_longer(-origin, names_to = "model", values_to = "MASE") |>
slice_min(MASE, by = origin) |>
count(model, name = "origins_won")
winsroll_long <- roll |> pivot_longer(-origin, names_to = "model", values_to = "MASE")
roll_long |> summarise(mean = mean(MASE), sd = sd(MASE),
min = min(MASE), max = max(MASE), .by = model) |>
mutate(across(where(is.numeric), \(z) round(z, 4)))ggplot(roll_long, aes(reorder(model, MASE, median), MASE)) +
geom_boxplot(fill = "#9EC5E8", width = 0.5) +
geom_jitter(width = 0.08, alpha = 0.6, size = 1.6) +
geom_hline(yintercept = 1, linetype = "dashed", color = "firebrick") +
coord_flip() +
labs(title = "MASE across eight forecast origins",
subtitle = "The spread within a model is comparable to the gaps between models",
x = NULL, y = "MASE") +
theme_dspa()
14.63 Checkpoint
- Your RNN’s loss falls smoothly but it never learns a dependency 40 steps back. What is happening, and how would you confirm it?
- You increase \(d_{\text{model}}\) from 64 to 512 and attention stops learning. Why?
- You shuffle the tokens of every input sentence and your Transformer’s accuracy is unchanged. What did you forget?
- Would 16 heads give more capacity than 4?
- Your Transformer forecaster beats ARIMA on your held-out year. Have you shown it is better?
- Why do Transformers dominate training but recurrent models persist in streaming deployments?
Answers
- Vanishing gradients, and the loss curve cannot reveal it. Backpropagation through time multiplies \(T-t\) Jacobians \(W_{hh}^\top\operatorname{diag}(\sigma')\), so the gradient decays as \((\rho(W_{hh})\gamma)^{T-t}\), geometric in the lag. The model learns short-range structure perfectly well, which keeps the loss falling, while no signal reaches 40 steps back. Confirm it by measuring the gradient by lag: backpropagate from the final output and plot \(\|\partial\mathcal L/\partial\mathbf x_t\|\) against \(t\); a straight line on a log axis is geometric decay. Remedies in order: check the forget-gate bias if using an LSTM, switch to gating if using a plain RNN, and use attention if the dependencies are genuinely long.
- The softmax saturated. With unit-variance queries and keys, the dot product \(\mathbf q\cdot\mathbf k\) has variance \(d_k\), so its standard deviation grew from \(8\) to \(\approx 23\) when you enlarged the model. Logits spread that far drive the softmax to near-one-hot, and its Jacobian \(\alpha_i(\delta_{ij}-\alpha_j)\) collapses because \(\alpha_i(1-\alpha_i)\to0\). Dividing the logits by \(\sqrt{d_k}\) restores unit variance at any dimension — it is not cosmetic. Diagnose by checking the attention entropy: near zero means saturation.
- Positional encoding. Self-attention is exactly permutation-equivariant, \(\mathrm{Attention}(PX)=P\,\mathrm{Attention}(X)\), so without positional information the model literally cannot distinguish word orders, the multiset of tokens is identical, so the multiset of outputs is. That your accuracy did not change is the diagnostic. Add sinusoidal or learned absolute positions, or RoPE. Note that learned absolute embeddings cannot extrapolate past the maximum trained length, which is why RoPE and ALiBi are preferred where long contexts matter.
- No. With \(d_k=d_{\text{model}}/h\), the projection parameters total \(4d_{\text{model}}^2\) for any head count, and the FLOPs are nearly identical. Heads partition the existing dimension rather than adding to it. What they buy is the ability to attend to several subspaces simultaneously, one softmax produces one weighted average, so a single head that must track two relations can only average them. The empirical evidence is sobering: many heads are prunable at inference with little loss, so the benefit is largely in giving optimization several parallel hypotheses.
- Not yet. Three things are missing. Benchmarks: if seasonal naive achieves a comparable MASE, neither model is adding value, report it. A scale-free metric: RMSE on one series says nothing transferable, so report MASE against the in-sample seasonal-naive scaling. And more than one origin: rolling-origin evaluation typically shows the ranking changing across origins, with within-model spread comparable to between-model gaps. Also check whether the series was scaled using statistics from the full record, a common leak, and whether ARIMA’s analytic prediction intervals matter for your use, since a point-forecasting network does not supply them.
- Sequential operations, not FLOPs. Attention costs \(O(T^2d)\) against recurrence’s \(O(Td^2)\), so it is more expensive whenever \(T>d\). But attention has \(O(1)\) sequential operations where recurrence has \(O(T)\): an RNN must finish step \(t\) before starting \(t+1\), so training cannot parallelize across time no matter how much hardware is available. On a GPU, attention’s higher FLOP count costs less wall time. In streaming deployment the calculus inverts, elements arrive one at a time, so there is no parallelism to exploit, and an RNN’s \(O(d^2)\) per element with \(O(d)\) state beats attention’s growing KV cache and \(O(Td)\) per token. That regime is why recurrent and state-space models remain competitive at the edge.
14.64 Summary and where Part 3 leads
Recurrence
- Recurrent networks share parameters across time, as CNNs share across space, so the parameter count is independent of \(T\).
- Backpropagation through time multiplies \(T-t\) Jacobians, so gradients decay or explode geometrically in the lag, governed by \(\rho(W_{hh})\). Clipping fixes explosion; vanishing gradients are silent and must be measured.
- LSTM and GRU gating creates an additive path with derivative \(\operatorname{diag}(\mathbf f_t)\), the same mechanism as a residual connection, on the time axis rather than depth. Initialize the forget bias positive.
- The encoder–decoder bottleneck, one fixed vector for any source length — is what motivated attention.
Attention
- Attention builds a step-specific context, puts every position one hop from every other, and is fully differentiable.
- \(1/\sqrt{d_k}\) is necessary: dot products have variance \(d_k\), and an unscaled softmax saturates and stops passing gradient.
- Mask before the softmax, or the weights do not sum to 1.
- Multi-head is free in parameters (\(4d^2\) for any \(h\)); it buys simultaneous attention to different subspaces.
- Self-attention is permutation-equivariant, so positional information must be injected. Sinusoidal encodings have the linear-shift property; RoPE and ALiBi handle length extrapolation better.
- Pre-norm keeps the residual stream clear and removes the warmup requirement; LayerNorm, not BatchNorm.
- Attention costs \(O(T^2d)\) against recurrence’s \(O(Td^2)\) and wins anyway, on \(O(1)\) sequential operations and \(O(1)\) path length. Its \(O(hT^2)\) memory is the binding constraint.
Applications
- Tokens are subwords. Token budgets are not word budgets, and perplexity is comparable only within a tokenization.
- Greedy decoding is not safe, it produces degenerate repetition. Nucleus sampling adapts the candidate set to the distribution’s shape.
- Neural forecasters do not beat classical methods by default on a single series. Report seasonal naive, use MASE, and evaluate at multiple origins.
- An LLM is a decoder-only Transformer trained to predict the next token. Scaling laws describe how loss falls with parameters and data; fluency is not understanding, and confident errors are the expected failure mode.
Continue with
| Part | Content |
|---|---|
| Part 4: Generative and Representation Learning | Autoencoders and their exact relation to PCA, the ELBO derived, GANs, diffusion, where the U-Net of Part 2, §14.37 returns as a denoising backbone, and self-supervised objectives generalize §14.58 beyond text |
| Part 5: Generalization, Uncertainty, and Practice | Calibration of the models built here, double descent, ensembles, Bayesian hyperparameter optimization, and pruning, which §14.51 anticipated in noting that many attention heads are removable |
Earlier material this part depended on
- Part 1: autograd, activations, initialization, LayerNorm, AdamW, warmup schedules
- Part 2: residual connections, which the LSTM anticipated by eighteen years
- Chapter 9: honest splitting and calibration
- Chapter 12: the vanishing-gradient product, MASE, rolling-origin evaluation, forecasting benchmarks
- Chapter 13: gradient flow, adaptive optimizers, automatic differentiation
Session information
#> 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] forecast_8.22.0 rlang_1.1.5 torch_0.13.0 plotly_4.12.1
#> [5] patchwork_1.3.0 tidyr_1.3.1 dplyr_1.1.4 ggplot2_4.0.1
#>
#> loaded via a namespace (and not attached):
#> [1] gtable_0.3.6 xfun_0.52 bslib_0.9.0 htmlwidgets_1.6.4
#> [5] processx_3.8.6 lattice_0.22-6 callr_3.7.6 quadprog_1.5-8
#> [9] vctrs_0.6.5 tools_4.3.3 crosstalk_1.2.1 ps_1.9.0
#> [13] generics_0.1.3 curl_6.2.0 parallel_4.3.3 tibble_3.2.1
#> [17] xts_0.13.2 pkgconfig_2.0.3 data.table_1.16.4 RColorBrewer_1.1-3
#> [21] S7_0.2.1 lifecycle_1.0.5 compiler_4.3.3 farver_2.1.2
#> [25] codetools_0.2-20 htmltools_0.5.8.1 sass_0.4.9 yaml_2.3.10
#> [29] pillar_1.10.1 jquerylib_0.1.4 cachem_1.1.0 nlme_3.1-165
#> [33] fracdiff_1.5-3 tidyselect_1.2.1 digest_0.6.37 purrr_1.0.2
#> [37] labeling_0.4.3 tseries_0.10-55 fastmap_1.2.0 grid_4.3.3
#> [41] colorspace_2.1-1 cli_3.6.3 magrittr_2.0.3 withr_3.0.2
#> [45] scales_1.4.0 bit64_4.0.5 TTR_0.24.4 rmarkdown_2.31
#> [49] httr_1.4.7 quantmod_0.4.26 bit_4.0.5 otel_0.2.0
#> [53] nnet_7.3-19 timeDate_4032.109 zoo_1.8-12 urca_1.3-3
#> [57] evaluate_1.0.3 knitr_1.51 lmtest_0.9-40 viridisLite_0.4.2
#> [61] Rcpp_1.0.14 glue_1.8.0 coro_1.0.4 rstudioapi_0.18.0
#> [65] jsonlite_1.8.9 R6_2.6.1