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.

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

How this part uses graphics

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

Heavy training is gated. HEAVY_EVAL defaults to FALSE, so 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.99080

A <- 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.738

The 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")))

DSPA3: Data Science and Predictive Analytics (UMich HS650)

Chapter 14, Part 3: Sequence Models - Recurrence and Attention

SOCR/MCAIM (Ivo Dinov)

August 31, 2026

Chapter 14 is published in five parts

How this part uses graphics