---
title: "The anatomy lab: locating sentiment in a language model"
format: html
engine: knitr
knitr:
  opts_chunk:
    collapse: true
    comment: "#>"
vignette: >
  %\VignetteIndexEntry{The anatomy lab: locating sentiment in a language model}
  %\VignetteEngine{quarto::html}
  %\VignetteEncoding{UTF-8}
---

```{r}
#| label: setup
#| include: false
# The heavy steps need a local GGUF model and glmnet; a checker (CRAN,
# r-universe) has neither. We gate those chunks on availability and fall back to
# an illustrative curve, so this vignette always renders with a figure. Point
# RELM_DEMO_MODEL (or RELM_TEST_MODEL_QWEN) at a GGUF file to render the
# figure from real activations.
model_path <- Sys.getenv("RELM_DEMO_MODEL", Sys.getenv("RELM_TEST_MODEL_QWEN", ""))
have_model <- nzchar(model_path) && file.exists(model_path)
have_glmnet <- requireNamespace("glmnet", quietly = TRUE)
run_real <- have_model && have_glmnet
```

`llm_trace()` observes the residual stream during a forward pass and hands it back
as a plain data frame. That turns a classic mechanistic-interpretability question
--- *where in the network does a concept become linearly readable?* --- into an
ordinary R modelling exercise. Here the concept is **sentiment**.

## Grouped probes in five lines

Prepare `labels`, a source ID for each prompt, and reserved source IDs before
examining evaluation results. Related documents, paraphrases and matched
contrasts share a source ID. With a loaded model `m`:

```{r}
#| eval: false
tr <- llm_trace(m, prompts, layers = 6:18, positions = "last")
fit <- llm_probe(labels ~ activations(layer = 6:18), tr,
                 groups = source_ids, test_groups = reserved_ids, seed = 42)
summary(fit)
plot(fit)
```

`predict(fit, new_trace)` returns positive-class probabilities using the layer
chosen on development CV and its saved preprocessing. `fit$audit` records each
prompt's group, partition and fold. Without `test_groups`, scores are explicitly
exploratory and have no inferential intervals. A spilled trace uses an external
label vector, because the in-memory proxy has zero rows.

Held-out intervals resample whole groups and are conditional on the fitted
models, development data, selection and split. They do not measure training
uncertainty or provide simultaneous coverage across layers. Small, undefined
or degenerate intervals are withheld with a reason in `fit$metrics`. The selected
layer is fixed before evaluation; repeatedly revising the analysis after seeing
the holdout consumes its independent status. Capture and fitting budgets still
apply: spilling does not make the statistical solver out of core.

An optional ggplot2 rendering uses the ordinary result table, without adding a
package dependency to relm:

```{r}
#| eval: false
ggplot2::ggplot(fit$metrics, ggplot2::aes(layer, test_score)) +
  ggplot2::geom_line() + ggplot2::geom_point() +
  ggplot2::geom_errorbar(ggplot2::aes(ymin = lower, ymax = upper), na.rm = TRUE) +
  ggplot2::labs(y = "Held-out AUC", subtitle = "Conditional pointwise intervals; development-selected layer fixed")
```

The repository's `tests/demos/demo-probe-evaluation.R` runs a fixed paired-source
example with a simple lexical baseline and a paired-label control. Its synthetic
labels deliberately reveal why good probe scores alone do not establish unique
representational information or causal use.

## Historical exploratory anatomy analysis

The longer analysis below predates the grouped probe API. Its selected-OOF
bootstrap does not account for selecting lambda or searching layers. Treat its
figures as exploratory, not independent final evaluation of the selected probe.

The recipe is deliberately classical: trace a fixed contrast set, take a concept
direction with `prcomp()`, fit one cross-validated linear probe per layer, and
plot decodability against depth. Then we *steer* along that direction and check
the effect on held-out prompts.

## AUC needs no dependency

Decodability is the probe's AUC. AUC is the rank-based Mann--Whitney statistic ---
twelve lines of base R, with average ranks giving exact tie handling --- so we
inline it rather than take a package dependency.

```{r}
#| label: auc
demo_auc <- function(scores, labels) {
  pos <- as.numeric(scores)[labels]
  neg <- as.numeric(scores)[!labels]
  n_pos <- length(pos)
  n_neg <- length(neg)
  r <- rank(c(pos, neg)) # average ranks => correct tie handling
  (sum(r[seq_along(pos)]) - n_pos * (n_pos + 1) / 2) / (n_pos * n_neg)
}

# Stratified percentile-bootstrap CI (resample within class, fixed seed).
demo_auc_ci <- function(scores, labels, B = 2000L, seed = 1L) {
  scores <- as.numeric(scores)
  pos <- scores[labels]
  neg <- scores[!labels]
  est <- demo_auc(scores, labels)
  set.seed(seed)
  boot <- vapply(seq_len(B), function(i) {
    demo_auc(c(sample(pos, replace = TRUE), sample(neg, replace = TRUE)),
             c(rep(TRUE, length(pos)), rep(FALSE, length(neg))))
  }, numeric(1))
  ci <- stats::quantile(boot, c(0.025, 0.975), names = FALSE)
  c(auc = est, lower = ci[1L], upper = ci[2L])
}

# Sanity checks: perfect separation -> 1, label flip -> 1 - AUC, a tie case.
stopifnot(
  all.equal(demo_auc(c(3, 4, 5, 0, 1, 2), c(TRUE, TRUE, TRUE, FALSE, FALSE, FALSE)), 1),
  all.equal(demo_auc(c(1, 2, 1, 3), c(TRUE, TRUE, FALSE, FALSE)), 0.375)
)
```

## A fixed contrast set

We use a fixed, committed set of opposite-sentiment sentences (no cherry-picking).
A handful, for illustration:

```{r}
#| label: prompts
pos <- c(
  "I absolutely loved this film; it was a masterpiece.",
  "The service was outstanding and the staff were welcoming.",
  "What a wonderful day -- everything went perfectly."
)
neg <- c(
  "I hated this film; it was a boring waste of time.",
  "The service was appalling and the staff were rude.",
  "What a miserable day -- everything went wrong."
)
train <- c(pos, neg)
y <- c(rep(TRUE, length(pos)), rep(FALSE, length(neg)))
```

The full runnable demo (a larger contrast set, held-out prompts, and the steering
check) is `tests/demos/demo-A-anatomy-lab.R` in the package repository.

## Trace, probe, and the money plot

With a model loaded, one `llm_trace()` call captures the last-token residual at
every layer; then for each layer we fit a **ridge** logistic probe with
`glmnet::cv.glmnet()` (ridge because the hidden width far exceeds the number of
prompts) and read its out-of-fold AUC.

This is an exploratory analysis: the same cross-validation results select the
regularization parameter, and the best layer is selected afterward. The bootstrap
intervals condition on the selected predictions; they do not account for tuning
or layer selection. Confirmatory performance estimates require nested
cross-validation or an independent test set, with related prompts kept in the
same fold where appropriate.

```{r}
#| label: probe-real
if (run_real) {
  library(relm)
  m <- llm(model_path)
  tr <- llm_trace(m, train, positions = "last", components = "residual")

  layer_matrix <- function(tr, L) {
    x <- as.matrix(tr, layer = L, component = "residual")
    x[order(as.integer(sub("[.].*$", "", rownames(x)))), , drop = FALSE]
  }
  yf <- factor(ifelse(y, "pos", "neg"), levels = c("neg", "pos"))
  set.seed(1)
  foldid <- sample(rep(seq_len(5), length.out = length(y)))

  band <- sort(unique(tr$layer))
  auc <- lower <- upper <- numeric(length(band))
  for (i in seq_along(band)) {
    cvfit <- glmnet::cv.glmnet(layer_matrix(tr, band[i]), yf, family = "binomial",
                               alpha = 0, foldid = foldid, type.measure = "auc",
                               keep = TRUE)
    oof <- cvfit$fit.preval[, match(cvfit$lambda.min, cvfit$lambda)]
    ci <- demo_auc_ci(oof, y)
    auc[i] <- ci[["auc"]]; lower[i] <- ci[["lower"]]; upper[i] <- ci[["upper"]]
  }
  res <- list(layers = band, auc = auc, lower = lower, upper = upper)
  close(m)
}
```

```{r}
#| label: probe-illustrative
if (isFALSE(run_real)) {
  # Illustrative curve (no model available at render time): sentiment becomes
  # linearly readable in the middle layers. Run the chunk above on your own model
  # to reproduce this figure from real activations.
  L <- 1:24
  set.seed(1)
  auc <- pmin(pmax(0.5 + 0.47 / (1 + exp(-(L - 7) / 1.6)) + rnorm(24, 0, 0.01), 0.5), 0.99)
  w <- (0.5 + 0.5 * (1 - auc)) * 0.09
  res <- list(layers = L, auc = auc, lower = pmax(auc - w, 0), upper = pmin(auc + w, 1))
}
```

```{r}
#| label: money-plot
#| fig-width: 8
#| fig-height: 5
#| fig-alt: "Probe AUC by layer, rising from chance to near one in the middle layers."
L <- res$layers
op <- par(mar = c(4.4, 4.4, 3, 1), mgp = c(2.6, 0.7, 0), las = 1)
plot(L, res$auc, type = "n", ylim = c(0.45, 1),
     xlab = "Transformer layer", ylab = "Probe AUC (out-of-fold)",
     main = sprintf("Where sentiment becomes linearly readable (%s)",
                    if (run_real) "real" else "illustrative"))
abline(h = 0.5, lty = 2, col = "grey45")
arrows(L, res$lower, L, res$upper, angle = 90, code = 3, length = 0.03, col = "grey30")
lines(L, res$auc, col = "#3B6EA5", lwd = 1.5)
points(L, res$auc, pch = 21, bg = "#3B6EA5", col = "white", cex = 1.2)
b <- L[which.max(res$auc)]
abline(v = b, lty = 3, col = adjustcolor("#D1495B", 0.4))
points(b, max(res$auc), pch = 21, bg = "#D1495B", col = "white", cex = 1.6)
text(max(L), 0.455, sprintf("best: layer %d (AUC %.2f)", b, max(res$auc)),
     adj = c(1, 0), font = 2, col = "#D1495B")
par(op)
```

The AUC rises from chance in the early layers to a plateau in the middle of the
network: that is *where sentiment lives* for this model.

## Steering along the direction

`llm_steer()` returns a **new** model handle that adds `coef * direction` to the
residual stream at one layer; the original handle is untouched. A concept
direction comes straight from `prcomp()` of the traced activations. Reading the
effect at a **later** layer on held-out prompts shows it propagates forward ---
a genuine downstream shift, not the injected vector read straight back.

```{r}
#| label: steer
#| eval: false
pc <- prcomp(layer_matrix(tr, b), center = TRUE, scale. = FALSE)
direction <- pc$rotation[, 1]
if (demo_auc(pc$x[, 1], y) < 0.5) direction <- -direction

m_pos <- llm_steer(m, layer = b, direction = direction, coef = 6)
llm_generate(m,     "In one sentence, describe your morning.", seed = 1)
llm_generate(m_pos, "In one sentence, describe your morning.", seed = 1)
close(m_pos)
```

The full script adds the quantitative check: a one-sided paired Wilcoxon test on
the held-out prompts' downstream sentiment read-out, positive vs. negative
steering. Everything here is base-R idiom over plain data frames and matrices ---
the model is just another object in your session.

## Reading is not causing

The money plot found where sentiment is *readable*. It is tempting to conclude that
those readable units are the ones the model *uses* for sentiment. They are not --- and
seeing the gap is worth a figure, because it is the representation-space cousin of
"correlation is not causation".

`llm_ablate()` returns a handle with a set of residual units pinned to zero;
`llm_logits()` reads the next-token distribution before and after. We rank the units
of the best layer three ways and measure the mean next-token KL from the un-ablated
model, each against a size-matched **random** control averaged over several draws:

- **load-bearing units** --- the top-*k* by residual RMS: the outlier / "massive
  activation" dimensions that carry most of the residual norm and are load-bearing for
  next-token prediction on *any* prompt;
- **concept-readout units** --- the top-*k* by |probe coefficient|: exactly the units
  the linear probe reads sentiment from;
- **matched random-*k*** --- the control.

```{r}
#| label: decodability-vs-causality
#| fig-width: 8
#| fig-height: 5
#| fig-alt: "Mean next-token KL vs neurons ablated, log scale: load-bearing units far above a matched-random control, while concept-readout units track random."
# Illustrative of the measured RELATIVE pattern (the absolute KL scale varies with
# model and layer -- e.g. impact KL ~ O(1) on Qwen2.5-0.5B, ~ O(0.03) on Gemma 4 E4B
# -- but the orders-of-magnitude separation does not). Reproduce it from your own
# model with tests/demos/demo-A-anatomy-lab.R: run_demo_A(extended = TRUE) writes the
# real three-series figure (demoA-A4-decodability.png).
ks <- c(1, 2, 4, 8, 16, 32, 64)
load_bearing    <- c(0.05, 0.10, 1.20, 2.00, 3.00, 3.50, 4.00) # top-k by residual RMS
concept_readout <- c(1e-4, 2e-4, 5e-4, 1.6e-3, 3e-3, 5e-3, 1.5e-2) # top-k by |probe coef|
matched_random  <- c(4e-4, 7e-4, 2e-3, 6e-3, 1.2e-2, 3e-2, 5e-2) # size-matched control

xx <- log2(ks)
cols <- c(load = "#D1495B", random = "grey45", concept = "#3B6EA5")
op <- par(mar = c(4.4, 4.6, 3, 1), mgp = c(2.7, 0.7, 0), las = 1)
plot(xx, load_bearing, type = "n", log = "y", xaxt = "n",
     ylim = c(1e-4, 6), xlab = "neurons ablated (k)",
     ylab = "Mean next-token KL (log)",
     main = "Reading is not causing: readable units are not the causal locus")
axis(1, at = xx, labels = ks)
abline(h = 10^(-4:0), col = "grey93")
for (s in list(list(matched_random, "random"), list(concept_readout, "concept"),
               list(load_bearing, "load"))) {
  lines(xx, s[[1]], col = cols[[s[[2]]]], lwd = 1.6)
  points(xx, s[[1]], pch = 21, bg = cols[[s[[2]]]], col = "white", cex = 1.3)
}
legend("topleft", bty = "n", cex = 0.9, pch = 21, lwd = 1.6, pt.bg = cols, col = cols,
       legend = c("load-bearing units (top-k residual RMS)",
                  "matched random-k (control)",
                  "concept-readout units (top-k |probe coef|)"))
par(op)
```

Ablating the load-bearing units dwarfs random by one to three orders of magnitude:
intervention is extremely **unit-specific** --- *which* units you remove is everything.
But the concept-readout units --- where sentiment is most *readable* --- are **no more
disruptive than matched random**. Decodability is not the causal locus.

### Sufficiency is not necessity

Does this contradict the steering result just above? No, and the reconciliation is the
whole point. Steering **adds** `coef * direction` and demonstrably moves sentiment: the
readable *direction* is causally **sufficient**. Ablation **pins to zero** the
direction's top basis *units* and changes almost nothing: those units are not
**necessary**. Sufficiency is not necessity --- different operation (add vs. pin-to-zero),
different object (a distributed direction vs. a handful of coordinate units), different
outcome metric (a concept shift vs. a concept-blind KL). The concept is *distributed*,
so no small set of readout coordinates is individually load-bearing. A linear probe
tells you where information is **readable**, not which units the computation **needs**.

### A magnitude caveat

Be careful about the "concept-readout below random" reading in the figure. The KL here
is **concept-blind** --- it measures total next-token disruption, not sentiment --- and
zero-ablating a unit perturbs the residual by roughly its own activation magnitude. The
concept-readout units happen to be low-magnitude, while the random draws are only
*count*-matched (not *magnitude*-matched) and so carry a larger typical magnitude. The
honest claim is therefore "**no more disruptive than matched random**", never "below
random"; a magnitude-matched control would be needed to say more.

### Where unit-level causal importance lives

Ranking units by *measured* causal effect --- ablate a unit (or group), read the KL,
repeat --- is exactly what activation patching and a probe-with-interventions do. That
is Phase-4 `llm_probe()` territory; the pattern already exists in the package's own WP5
honesty fixture (`calibrate_kl` in `intervene_kl.rs`), which selected its targeted
neurons by a per-neuron ablation-KL scan rather than by any linear readout.