diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6e3f8e6..d6ea123 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -81,11 +81,11 @@ Initial release.
 
 ### Added
 
-- Quick API on `MuranoModel`: `find_direction()`, `generate(intervention=...)`, direct activation recording.
+- Quick API on `MuranoModel`: `find_direction()`, `generate(ablate=...)` / `generate(steer=...)`, direct activation recording.
 - Pipeline API: composable `Step` + `Pipeline` with pre-flight validation of `reads`/`writes` contracts.
 - Steps: `Load`, `Record`, `SteeringVector`, `Intervene`, `Probe`, `GenerationMetric`.
 - Logit lens (`LogitLens` step).
-- Datasets: `MuranoDataset` (contrastive) and `LabeledDataset`, with `from_hub()` and `from_template()` factories.
+- Datasets: `MuranoDataset` (contrastive) and `LabeledDataset`, with `contrastive()` / `from_hub()` factories.
 - Direction-based interventions: `ablate_direction`, `steer_direction`.
 - I/O: `save_results()` with structured output layout, `load_steering()`, `save_ablated_model()`.
 - Top-level `__version__` via `importlib.metadata`.
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index d1940be..8ade9eb 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -103,8 +103,8 @@ No bespoke try/except guards: every optional import goes through
 
 ## Code of Conduct
 
-Please note that this project is released with a [Code of Conduct](CODE_OF_CONDUCT.md). By participating in this project you agree to abide by its terms.
+ToDo: adopt and link a Code of Conduct. We ask contributors to be respectful and constructive in the meantime.
 
 ## License
 
-By contributing, you agree that your contributions will be licensed under the MIT License.
+By contributing, you agree that your contributions will be licensed under the Apache License 2.0, the license this project is released under.
diff --git a/README.md b/README.md
index 4ba0aba..ed74d5e 100644
--- a/README.md
+++ b/README.md
@@ -28,9 +28,10 @@ libraries ship as extras, so you install per use case:
 | (base) | recording, steering, intervention, logits, ablation, metrics, paired datasets | nnsight, nnterp, torch, transformers |
 | `probe` | linear probing | scikit-learn |
 | `data` | loading datasets by name from the Hub | datasets |
-| `plot` | figures and visualizations | matplotlib, seaborn, plotly |
+| `plot` | figures and visualizations | plotly, kaleido |
 | `sae` | sparse autoencoder features | sae-lens |
-| `all` | everything above | all of the above |
+| `notebook` | running the example notebooks | jupyter, ipykernel, nltk |
+| `all` | everything above | probe, data, plot, sae, notebook |
 
 ```bash
 pip install "murano-interp[probe,plot]"   # combine as needed
@@ -133,6 +134,13 @@ caught up-front.
 | `WeightAblation`   | `prompts`, `steering` | `intervene`, `weight_ablation` | Project a direction out of model weights, then generate. |
 | `Logits` ‡         | `prompts`             | `final_logits`, `attention_mask`, `target_ids` | Run a forward pass and expose output logits plus next-token targets. With `fn=`, apply an intervention during that pass. |
 | `Ablate` ‡         | `prompts`             | `ablated_logits`, `attention_mask` | Zero, mean, or resample a component and return the logits. |
+| `Patch` ‡          | `corrupt_prompts`, `prompts` | `patched_logits`, `patched_mask` | Interchange (activation) patching: inject source-run activations into the base run. |
+| `PathPatch` ‡      | `prompts`, `corrupt_prompts` | `path_patched_logits`, `path_patched_mask` | Isolate a sender-to-receiver path's effect. |
+| `RecordAttention` ‡ | `prompts`            | `attention_pattern`            | Capture per-head attention weights (needs `enable_attention_probs`). |
+| `AblateAttention` ‡ | `prompts`            | `attn_ablated_logits`, `attn_ablated_mask` | Overwrite per-head attention weights and return the logits. |
+| `LogitAttribution` ‡ | `prompts`           | `logit_attribution`            | Decompose a logit (difference) into per-component contributions (DLA). |
+| `SelectComponents` ‡ | `logit_attribution` (or `sweep`) | `selection`          | Rank per-component scores and keep the strongest as a target set. |
+| `LogitLens`        | `prompts`             | `logit_lens`                   | Project each layer's residual onto the vocabulary. |
 | `Sweep` ‡          | (the swept chain's)   | `sweep`                        | Run a step chain once per item and harvest a metric into a `SweepResult`. |
 | `Probe` §          | `record`              | `probe`                        | Train a linear probe per layer via cross-validation.     |
 | `GenerationMetric` | `intervene`           | `metric`                       | Score baseline vs modified outputs with a user metric.   |
@@ -172,15 +180,18 @@ Pipeline API like the other steps.
 
 ```text
 src/murano/
-  model.py
-  pipeline.py
-  results.py
-  artifacts.py
-  dataset.py
-  io.py
-  evaluation.py
-  steps/
-  plotting/
+  model.py       # MuranoModel: the nnterp-backed model wrapper
+  backend.py     # ModelBackend: the surface steps depend on
+  pipeline.py    # Pipeline / Step orchestration
+  results.py     # Results: the shared state passed between steps
+  keys.py        # canonical Results keys
+  nodes.py       # Node / Edge / NodeSet: component addressing
+  artifacts.py   # PromptBatch, SteeringResult, MetricScore, ...
+  dataset.py     # MuranoDataset, LabeledDataset, CleanCorruptDataset
+  tasks.py       # small canonical tasks (ioi, sentiment)
+  io.py          # save/load of results
+  steps/         # pipeline steps
+  plotting/      # Plotly visualizations
 ```
 
 ## Notebooks
@@ -195,8 +206,11 @@ then pick the application you need:
 - [`notebooks/reproductions/`](notebooks/reproductions/) — published results
   reproduced with Murano.
 
-Every notebook is executed before it is committed, and all of them render on the
+Notebooks are executed before they are committed so their stored outputs are
+current, and all of them render on the
 [documentation site](https://ukplab.github.io/murano/docs/notebooks/getting_started/).
+(This is a maintainer pre-commit step; CI checks that the notebooks import and
+parse, but does not re-execute them, since several need a GPU.)
 
 ## Development
 
diff --git a/docs/src/content/docs/docs/getting-started/installation.md b/docs/src/content/docs/docs/getting-started/installation.md
index 8ba1350..5652915 100644
--- a/docs/src/content/docs/docs/getting-started/installation.md
+++ b/docs/src/content/docs/docs/getting-started/installation.md
@@ -24,9 +24,10 @@ libraries ship as extras, so you install for what you actually do:
 | (base) | recording, steering, intervention, logits, ablation, metrics, paired datasets | nnsight, nnterp, torch, transformers |
 | `probe` | linear probing | scikit-learn |
 | `data` | loading datasets by name from the Hub | datasets |
-| `plot` | figures and visualizations | matplotlib, seaborn, plotly |
+| `plot` | figures and visualizations | plotly, kaleido |
 | `sae` | sparse autoencoder features | sae-lens |
-| `all` | everything above | all of the above |
+| `notebook` | running the example notebooks | jupyter, ipykernel, nltk |
+| `all` | everything above | probe, data, plot, sae, notebook |
 
 Combine extras as needed:
 
diff --git a/docs/src/content/docs/docs/getting-started/quickstart.mdx b/docs/src/content/docs/docs/getting-started/quickstart.mdx
index ed5df1c..e43e3a8 100644
--- a/docs/src/content/docs/docs/getting-started/quickstart.mdx
+++ b/docs/src/content/docs/docs/getting-started/quickstart.mdx
@@ -40,8 +40,8 @@ direction = model.find_direction(
     negative=["What a miserable, dreadful day", "This is awful and depressing"],
     layers=[10, 11, 12, 13],
 )
-print(direction.best_layer)         # e.g. 12
-print(direction.separation_scores)  # {10: 1.3, 11: 2.1, 12: 3.4, 13: 2.8}
+print(direction.best_layer)         # a Node, e.g. L12.resid_post
+print(direction.separation_scores)  # {L10.resid_post: 1.3, L11.resid_post: 2.1, ...}
 ```
 
 ## Steer generation
diff --git a/docs/src/content/docs/docs/guides/patch.md b/docs/src/content/docs/docs/guides/patch.md
index a0b2bb9..b37baa5 100644
--- a/docs/src/content/docs/docs/guides/patch.md
+++ b/docs/src/content/docs/docs/guides/patch.md
@@ -67,13 +67,13 @@ results = Pipeline([
 
 ### Custom intervention functions
 
-`Intervene` accepts any callable with signature `(activation: Tensor, layer_idx: int) -> Tensor`:
+`Intervene` accepts any callable with signature `(activation: Tensor, node: Node) -> Tensor`. The second argument is a `murano.Node` naming the hooked site (`node.layer`, `node.module`, ...), not a bare layer index, so a function that targets specific sites branches on it:
 
 ```python
 import torch
 
-def zero_out(activation: torch.Tensor, layer: int) -> torch.Tensor:
-    """Zero all activations — a sanity check."""
+def zero_out(activation: torch.Tensor, node) -> torch.Tensor:
+    """Zero all activations — a sanity check. `node` names the hooked site."""
     return torch.zeros_like(activation)
 
 results = Pipeline([
@@ -84,7 +84,7 @@ results = Pipeline([
 
 ## Weight-level ablation
 
-The `WeightAblation` step applies an orthogonal projection P = I - dd^T to the model's weight matrices directly. This removes the direction from **all** computations, not just the residual stream at hook points.
+The `WeightAblation` step applies an orthogonal projection P = I - dd^T to the model's weight matrices directly (write matrices as `P @ W`, read matrices as `W @ P`), so the ablation is baked into the weights rather than applied at hook points during the forward pass. It is a stronger, more global edit than activation-level ablation; note that the read-matrix projection interacts with the intervening layer norm, so it does not remove the direction perfectly.
 
 ```python
 from murano.steps import WeightAblation
diff --git a/docs/src/pages/reproductions/ioi.mdx b/docs/src/pages/reproductions/ioi.mdx
index 0fa5aaf..8b62c3e 100644
--- a/docs/src/pages/reproductions/ioi.mdx
+++ b/docs/src/pages/reproductions/ioi.mdx
@@ -63,7 +63,11 @@ from murano.steps.metrics import LogitDiffStep
 from murano.steps.paired import LoadPaired
 from murano.results import Results
 
-model = MuranoModel("gpt2", enable_attention_probs=True)
+import torch
+
+# gpt2 is small, so float32 is cheap and keeps the logit differences clean
+# (bfloat16 measurably reorders close logits at this scale).
+model = MuranoModel("gpt2", dtype=torch.float32, enable_attention_probs=True)
 ds = CleanCorruptDataset(clean=clean, corrupt=corrupt, correct=io_ids, incorrect=s_ids)
 
 base = LoadPaired(ds)(Results())
diff --git a/docs/src/pages/reproductions/tmlr2026.mdx b/docs/src/pages/reproductions/tmlr2026.mdx
index 908d62e..3089cec 100644
--- a/docs/src/pages/reproductions/tmlr2026.mdx
+++ b/docs/src/pages/reproductions/tmlr2026.mdx
@@ -161,6 +161,14 @@ plt.tight_layout(); plt.show()
 
 ## Key results
 
+<p class="aside">
+  Preview: this reproduction is a work in progress. It depends on an external
+  <code>SupervisedMDS</code> package and a placeholder dataset (see above), so the
+  numbers below are the paper's reported figures, not yet independently
+  reproduced with Murano. The other gallery entries (IOI, Geometry of Truth,
+  Function Vectors) are full reproductions.
+</p>
+
 | Concept type      | Model        | Best-fit manifold | Probe accuracy |
 | ----------------- | ------------ | ----------------- | -------------- |
 | Month of year     | Llama-3.2-1B | Circle            | 0.91           |
diff --git a/src/murano/dataset.py b/src/murano/dataset.py
index 997fe9f..92a01ab 100644
--- a/src/murano/dataset.py
+++ b/src/murano/dataset.py
@@ -162,6 +162,16 @@ def from_hub(
         pos_texts = _load_hub_column(positive, n=n_total)
         neg_texts = _load_hub_column(negative, n=n_total)
 
+        # A source with fewer than n_train + n_eval rows would otherwise yield a
+        # short (or empty) eval split silently; fail loudly instead.
+        for label, texts in (("positive", pos_texts), ("negative", neg_texts)):
+            if len(texts) < n_total:
+                raise ValueError(
+                    f"{label} source yielded {len(texts)} examples, but n_train + "
+                    f"n_eval = {n_total} are needed; lower n_train/n_eval or use a "
+                    f"larger split."
+                )
+
         raw_pos = list(pos_texts)
         raw_neg = list(neg_texts)
 
@@ -298,7 +308,14 @@ def from_hub(
             )
         """
         if isinstance(source, tuple):
-            name, config = source if len(source) == 2 else (source[0], None)
+            if len(source) == 2:
+                name, config = source
+            elif len(source) == 1:
+                name, config = source[0], None
+            else:
+                raise ValueError(
+                    f"source tuple must be (name,) or (name, config), got {source!r}"
+                )
         else:
             name, config = source, None
 
diff --git a/src/murano/io.py b/src/murano/io.py
index 9cc9edc..60d18ef 100644
--- a/src/murano/io.py
+++ b/src/murano/io.py
@@ -1,4 +1,12 @@
-"""I/O utilities: saving and loading Murano results."""
+"""I/O utilities: saving and loading Murano results.
+
+Security note: the ``.pt`` loaders (``load_steering``, ``load_logit_lens``,
+``load_attention``, ``load_activation_store``, and the labeled variant) call
+``torch.load(..., weights_only=False)`` because the payloads hold custom objects
+(``Node`` keys, dataclasses) that the safe loader cannot reconstruct. That path
+executes arbitrary pickle, so only load artifacts you produced yourself or
+otherwise trust; do not load ``.pt`` files from an untrusted source.
+"""
 
 from __future__ import annotations
 
@@ -408,14 +416,28 @@ def _restore(value: Any) -> Any:
         # null is how a non-finite scalar was stored; restore it as nan.
         return float("nan") if value is None else value
 
+    # The per-component maps can carry non-finite values too (a head whose frozen
+    # projection produced nan), stored as null; restore them the same way as the
+    # scalar fields so the round-trip is symmetric.
+    contributions = {node: _restore(v) for node, v in data["contributions"].items()}
+    raw_per_example = data.get("per_example")
+    per_example = (
+        None
+        if raw_per_example is None
+        else {
+            node: [_restore(x) for x in values]
+            for node, values in raw_per_example.items()
+        }
+    )
+
     return LogitAttributionResult(
-        contributions=data["contributions"],
+        contributions=contributions,
         embed_contribution=_restore(data["embed_contribution"]),
         other_contribution=_restore(data["other_contribution"]),
         target=data["target"],
         total=_restore(data["total"]),
         completeness_error=_restore(data["completeness_error"]),
-        per_example=data.get("per_example"),
+        per_example=per_example,
         metadata=data.get("metadata", {}),
     )
 
@@ -951,10 +973,12 @@ def serialize_logit_lens(
     ) -> None:
         filename = "logit_lens.pt" if key == keys.LOGIT_LENS else f"{key}.pt"
         save_logit_lens(logit_lens, out / "logit_lens" / filename)
+        # Read shapes off max_probs, not all_probs: the full-vocab tensor is
+        # None unless the step ran with store_full_probs=True.
         metadata[key] = {
             "addresses": [str(a) for a in logit_lens.addresses],
-            "n_layers": logit_lens.all_probs.shape[0],
-            "n_inputs": logit_lens.all_probs.shape[1],
+            "n_layers": logit_lens.max_probs.shape[0],
+            "n_inputs": logit_lens.max_probs.shape[1],
         }
 
     def serialize_attention(
diff --git a/src/murano/model.py b/src/murano/model.py
index 732680b..43e46c2 100644
--- a/src/murano/model.py
+++ b/src/murano/model.py
@@ -6,7 +6,7 @@
 from typing import TYPE_CHECKING, Any, Callable, Sequence, cast
 
 import torch
-from torch import Tensor, bfloat16  # pyright: ignore[reportPrivateImportUsage]
+from torch import Tensor, bfloat16, float32  # pyright: ignore[reportPrivateImportUsage]
 from torch import dtype as TorchDtype  # pyright: ignore[reportPrivateImportUsage]
 from nnterp import StandardizedTransformer
 
@@ -20,6 +20,7 @@
     Node,
     NodeDict,
     canonical_module,
+    unreachable_addresses,
 )
 
 if TYPE_CHECKING:
@@ -32,6 +33,79 @@
 # future work; until then unknown architectures raise from ``attn_out_proj``.
 _ATTN_OUT_PROJ_NAMES = ("o_proj", "out_proj", "c_proj", "dense", "wo")
 
+# float32 is worth the memory up to this many parameters; above it the default
+# resolves to bfloat16 so the model still fits. bf16's 8-bit mantissa measurably
+# reorders close logits at small scale, so the auto default keeps small-model
+# (GPT-2 / <=~3B) interpretability in the precision its logit differences need.
+_FP32_PARAM_LIMIT = 3_000_000_000
+
+
+def _estimate_num_params(config: Any) -> int | None:
+    """Rough parameter count from a model config, for the auto-dtype heuristic.
+
+    Uses only the shape fields every decoder config carries, so it never loads
+    weights. Accounts for grouped-query attention (fewer key/value params) and
+    tied input/output embeddings (counted once), which a naive estimate gets
+    badly wrong on large-vocab models like Gemma. Returns None when the config
+    omits a field the estimate needs.
+    """
+    hidden = getattr(config, "hidden_size", None) or getattr(config, "n_embd", None)
+    layers = getattr(config, "num_hidden_layers", None) or getattr(
+        config, "n_layer", None
+    )
+    vocab = getattr(config, "vocab_size", None)
+    if not hidden or not layers or not vocab:
+        return None
+    inter = getattr(config, "intermediate_size", None) or 4 * hidden
+    n_heads = (
+        getattr(config, "num_attention_heads", None)
+        or getattr(config, "n_head", None)
+        or 1
+    )
+    head_dim = getattr(config, "head_dim", None) or max(hidden // n_heads, 1)
+    n_kv = getattr(config, "num_key_value_heads", None) or n_heads
+    q_out = n_heads * head_dim
+    kv_out = n_kv * head_dim
+    # Attention: q (hidden->q_out) + k + v (hidden->kv_out each) + o (q_out->hidden).
+    attn = hidden * q_out + 2 * hidden * kv_out + q_out * hidden
+    # MLP: 3*hidden*inter upper-bounds a gated MLP (gate/up/down); an over-estimate
+    # for a plain 2-matrix MLP, which is fine for a threshold.
+    mlp = 3 * hidden * inter
+    per_layer = attn + mlp
+    # Embeddings: input + output, counted once when tied (the common case at the
+    # small/mid scale this threshold cares about, e.g. GPT-2, Gemma, Llama-3.2).
+    tied = bool(getattr(config, "tie_word_embeddings", False))
+    embed = vocab * hidden * (1 if tied else 2)
+    return layers * per_layer + embed
+
+
+def _resolve_auto_dtype(load_path: str) -> TorchDtype:
+    """Pick float32 for small models and bfloat16 for large ones, from the config.
+
+    Reads only the config (no weights), estimates the parameter count, and
+    returns float32 at or below :data:`_FP32_PARAM_LIMIT` else bfloat16. Falls
+    back to bfloat16 (the safe large-model default) when the config cannot be
+    read or is missing size fields.
+    """
+    try:
+        from transformers import AutoConfig
+
+        config = AutoConfig.from_pretrained(load_path)
+    except Exception as exc:  # pragma: no cover - exotic/unreadable config
+        logger.debug("auto dtype: could not read config (%s); using bfloat16", exc)
+        return bfloat16
+    n_params = _estimate_num_params(config)
+    if n_params is None:
+        logger.debug("auto dtype: config missing size fields; using bfloat16")
+        return bfloat16
+    chosen = float32 if n_params <= _FP32_PARAM_LIMIT else bfloat16
+    logger.info(
+        "auto dtype: ~%.2fB params -> %s (pass dtype= to override)",
+        n_params / 1e9,
+        chosen,
+    )
+    return chosen
+
 
 def _ensure_downloaded(model_id: str) -> str:
     """Ensure the model is available locally and return its snapshot path.
@@ -64,7 +138,15 @@ class MuranoModel:
     Args:
         model_id: HuggingFace model identifier.
         device_map: Device placement strategy.
-        dtype: Model weight dtype.
+        dtype: Model weight dtype, or ``"auto"`` (the default). ``"auto"`` loads
+            small models (at or below ~3B parameters, e.g. GPT-2) in float32 and
+            larger ones in bfloat16, decided from the config before any weight is
+            read. This matters for correctness: bfloat16's 8-bit mantissa
+            measurably reorders close logits at small scale, so argmax,
+            logit-difference, KL, and answer-rank degrade versus float32 on the
+            small models interpretability work most often targets. Pass an
+            explicit ``torch.dtype`` to override (e.g. ``torch.float16`` for a
+            large model on a memory-tight GPU).
         enable_attention_probs: If True, load with eager attention so nnterp can
             expose the per-head softmax attention weights via
             :attr:`attention_probabilities`. Off by default because eager
@@ -85,7 +167,7 @@ def __init__(
         self,
         model_id: str,
         device_map: str = "auto",
-        dtype: TorchDtype = bfloat16,
+        dtype: TorchDtype | str = "auto",
         enable_attention_probs: bool = False,
         **loader_kwargs: Any,
     ):
@@ -96,6 +178,16 @@ def __init__(
         # making repeated HF API calls that trigger rate limits.
         load_path = _ensure_downloaded(model_id)
 
+        # Resolve the auto default from the config before loading: small models
+        # get float32 (its precision matters for their logit differences), large
+        # ones bfloat16 (so they fit). An explicit torch.dtype is used as given.
+        if isinstance(dtype, str):
+            if dtype != "auto":
+                raise ValueError(
+                    f"dtype string must be 'auto' or a torch.dtype, got {dtype!r}"
+                )
+            dtype = _resolve_auto_dtype(load_path)
+
         # device_map="auto" with nnsight can produce zero/NaN activations
         # for layers beyond the first. Use a single GPU when available.
         if device_map == "auto" and torch.cuda.is_available():
@@ -481,21 +573,23 @@ def _check_directions_reachable(
         """
         if not directions:
             return
-        module_list = [modules] if isinstance(modules, str) else list(modules)
-        hooked_modules = {Node(0, mod).module for mod in module_list}
-        if isinstance(layers, str):
-            # "all" hooks every layer, so a matching module is sufficient.
-            reachable = any(node.module in hooked_modules for node in directions)
-        else:
-            hooked = {Node(layer, mod) for layer in layers for mod in module_list}
-            reachable = bool(set(directions) & hooked)
-        if not reachable:
+        missing = unreachable_addresses(directions, layers, modules)
+        if len(missing) == len(directions):
             raise ValueError(
                 f"No intervention address matches a hooked site "
                 f"(layers={layers!r}, modules={modules!r}); the intervention "
                 f"would do nothing. Addresses: {sorted(directions)}. A bare-int "
                 f"direction key targets 'resid_post'."
             )
+        if missing:
+            logger.warning(
+                "%d intervention address(es) match no hooked site "
+                "(layers=%r, modules=%r) and will be skipped: %s",
+                len(missing),
+                layers,
+                modules,
+                sorted(missing),
+            )
 
     def _layer_indices(self, layers: list[int] | str) -> list[int]:
         if isinstance(layers, str):
diff --git a/src/murano/nodes.py b/src/murano/nodes.py
index 68d42f8..760f1e7 100644
--- a/src/murano/nodes.py
+++ b/src/murano/nodes.py
@@ -348,6 +348,54 @@ def parse(cls, text: str) -> Node:
 AddressLike: TypeAlias = int | tuple[int, str] | str | Node
 
 
+def unreachable_addresses(
+    addresses: "Iterable[Node]",
+    layers: "list[int] | str",
+    modules: "str | list[str]",
+) -> "list[Node]":
+    """Return the addresses that fall on no hooked ``(layer, module)`` site.
+
+    An intervention direction keyed at an address that is never hooked silently
+    does nothing (e.g. a ``resid_post`` key while ``modules="mlp"``). Callers use
+    this to fail loudly when *every* address is unreachable, and to warn when only
+    some are. Module names are canonicalized, so ``"residual"`` matches a
+    ``resid_post`` address.
+
+    Args:
+        addresses: The intervention addresses (:class:`Node` objects).
+        layers: Hooked layer indices, or ``"all"``.
+        modules: Hooked module name(s).
+
+    Returns:
+        The subset of ``addresses`` that no hooked site would touch (empty when
+        all are reachable).
+    """
+    addrs = list(addresses)
+    module_list = [modules] if isinstance(modules, str) else list(modules)
+    hooked_modules = {Node(0, mod).module for mod in module_list}
+    if isinstance(layers, str):
+        # "all" hooks every layer, so a matching module is sufficient.
+        return [node for node in addrs if node.module not in hooked_modules]
+    hooked = {Node(layer, mod) for layer in layers for mod in module_list}
+    return [node for node in addrs if node not in hooked]
+
+
+def addresses_reachable(
+    addresses: "Iterable[Node]",
+    layers: "list[int] | str",
+    modules: "str | list[str]",
+) -> bool:
+    """Whether at least one address falls on a hooked ``(layer, module)`` site.
+
+    An empty ``addresses`` is a deliberate no-op and returns ``True``. See
+    :func:`unreachable_addresses` for the per-address detail callers warn on.
+    """
+    addrs = list(addresses)
+    if not addrs:
+        return True
+    return len(unreachable_addresses(addrs, layers, modules)) < len(addrs)
+
+
 def _parse_int(value: str, name: str, raw: str) -> int:
     """Parse ``value`` as an int, naming ``name`` and ``raw`` on failure.
 
@@ -613,5 +661,7 @@ def __ror__(self, other) -> NodeDict:
     "NodeDict",
     "NodeSet",
     "Side",
+    "addresses_reachable",
     "canonical_module",
+    "unreachable_addresses",
 ]
diff --git a/src/murano/plotting/attention.py b/src/murano/plotting/attention.py
index e8e0619..79512d0 100644
--- a/src/murano/plotting/attention.py
+++ b/src/murano/plotting/attention.py
@@ -88,7 +88,11 @@ def plot_head_matrix(
         matrix: A ``[n_layers, n_heads]`` tensor or nested list, e.g. the output
             of ``AttentionResult.entropy()``.
         title: Plot title.
-        layers: Layer indices for the y-axis labels; defaults to ``0..n-1``.
+        layers: Layer indices for the y-axis labels; defaults to ``0..n-1``. When
+            the matrix came from an ``AttentionResult`` captured on a subset of
+            layers (e.g. ``RecordAttention(layers=[3, 7, 9])``), pass
+            ``layers=result.layers`` so the row labels name the real layers rather
+            than a contiguous ``0..n-1``.
         value_label: Label for the colorbar (the statistic being shown).
         color_scale: Plotly colorscale name.
         zmid: Value anchored to the middle of the colorscale. Pass ``0`` for a
diff --git a/src/murano/plotting/sae.py b/src/murano/plotting/sae.py
index 2dbb8f2..70d62e2 100644
--- a/src/murano/plotting/sae.py
+++ b/src/murano/plotting/sae.py
@@ -17,6 +17,8 @@
     topk,  # pyright: ignore[reportPrivateImportUsage]
 )
 
+from murano._optional import require_optional
+
 if TYPE_CHECKING:
     import plotly.graph_objects as go
 
@@ -263,6 +265,7 @@ def plot_sae_feature_logit_effects(
         A Plotly figure containing a positive/negative token table and a
         histogram over every vocabulary logit effect.
     """
+    require_optional("plot", "plotly")
     import plotly.graph_objects as go
     from plotly.subplots import make_subplots
 
@@ -383,6 +386,7 @@ def plot_sae_token_activations(
     Returns:
         A Plotly figure whose token backgrounds encode activation strength.
     """
+    require_optional("plot", "plotly")
     import plotly.graph_objects as go
 
     # Validate the public row contract once, then keep rendering simple.
diff --git a/src/murano/steps/ablate.py b/src/murano/steps/ablate.py
index 97655e6..1ef21ec 100644
--- a/src/murano/steps/ablate.py
+++ b/src/murano/steps/ablate.py
@@ -480,11 +480,6 @@ def __call__(self, results: Results) -> Results:
         attention_mask: torch.Tensor = tokens["attention_mask"]
         batch, seq = tokens["input_ids"].shape
 
-        # zero needs no capture; mean from a precomputed table needs none either.
-        captured: dict[Site, torch.Tensor] = {}
-        if self.method == "resample" or (self.method == "mean" and self.means is None):
-            captured = self._capture(tokens)
-
         # resample draws its replacement from a second batch: raw source= prompts
         # or a source_key= prompt batch already in results (the cross-run patch).
         source_prompts: Sequence[str] | None = None
@@ -494,6 +489,18 @@ def __call__(self, results: Results) -> Results:
             elif self.source_key is not None:
                 source_prompts = results[self.source_key].prompts
 
+        # Capture the base activations only when they are actually used: for the
+        # batch mean (mean without a precomputed table) or the within-batch
+        # resample permutation (resample with no source). A cross-run resample
+        # replaces from source_captured, so the base capture would be a wasted
+        # forward pass. zero needs no capture at all.
+        need_base_capture = (self.method == "mean" and self.means is None) or (
+            self.method == "resample" and source_prompts is None
+        )
+        captured: dict[Site, torch.Tensor] = (
+            self._capture(tokens) if need_base_capture else {}
+        )
+
         source_captured: dict[Site, torch.Tensor] | None = None
         if source_prompts is not None:
             source_captured = self._capture(
diff --git a/src/murano/steps/attention.py b/src/murano/steps/attention.py
index e7ba7ad..1cca675 100644
--- a/src/murano/steps/attention.py
+++ b/src/murano/steps/attention.py
@@ -146,10 +146,11 @@ def _pattern_entropy(pattern: Tensor, mask: Tensor) -> Tensor:
 def _pattern_sink(pattern: Tensor, mask: Tensor, index: int) -> Tensor:
     """Return the per-head mean attention mass on the ``index``-th real key.
 
-    The key is taken relative to each example's first real token, not absolute
-    column ``index``: Murano's tokenizers left-pad by default, so column 0 is a
-    padding token for the shorter sequences in a batch. ``index=0`` is the usual
-    attention sink (the first real token).
+    The key is taken relative to each example's first real token (found from the
+    mask), not absolute column ``index``, so it is correct whichever side the
+    batch is padded on: under left padding column 0 is a padding token for the
+    shorter sequences. ``index=0`` is the usual attention sink (the first real
+    token).
     """
     batch, _, _, k_len = pattern.shape
     first_real = mask.long().argmax(dim=1)  # first key with mask == 1, per example
@@ -260,8 +261,9 @@ def attention_to(self, query=None, key=None) -> Tensor:
         :func:`~murano.steps.metrics._answer_positions` form (an int, a per-example
         sequence or tensor, negatives from the end); ``None`` uses each example's
         last real token (the natural query for next-token prediction). Explicit
-        positive positions are absolute columns, so under Murano's default left
-        padding use ``-1`` for the last real token and prefer negative indices.
+        positive positions are absolute columns, so on a left-padded batch they
+        point into the padding for the shorter sequences; prefer negative indices
+        (``-1`` is the last real token) unless you know the batch is right-padded.
         """
         q = _positions(query, self.attention_mask)
         k = _positions(key, self.attention_mask)
@@ -574,8 +576,8 @@ class AblateAttention(Step):
             :func:`~murano.steps.metrics._answer_positions` form (an int or
             per-example sequence, negatives allowed); ``None`` overwrites every
             query row. These are absolute columns (negatives count from the end),
-            so under Murano's default left padding use ``-1`` for the last real
-            token.
+            so on a left-padded batch a positive index points into the padding for
+            the shorter sequences; use ``-1`` for the last real token.
         prompts_key: Results key to read the batch to run from.
         logits_key: Results key to write the intervened logits under.
         mask_key: Results key to write the attention mask under.
diff --git a/src/murano/steps/intervene.py b/src/murano/steps/intervene.py
index 69bf9f6..1461d64 100644
--- a/src/murano/steps/intervene.py
+++ b/src/murano/steps/intervene.py
@@ -10,7 +10,7 @@
 from murano import keys
 from murano.artifacts import GenerationComparison, PromptBatch
 from murano.logging import logger
-from murano.nodes import Node, NodeDict
+from murano.nodes import Node, NodeDict, unreachable_addresses
 from murano.results import Results
 from murano.steps.base import Step
 
@@ -60,7 +60,9 @@ def _normalize_directions(directions: dict[Node, Tensor]) -> NodeDict:
 def ablate_direction(directions: dict[Node, Tensor]) -> Callable:
     """Return an intervention function that projects out a direction.
 
-    Removes the component along the direction from the residual stream.
+    Removes the component along the direction from the residual stream. Each
+    direction is unit-normalized first, so its magnitude does not matter for
+    ablation (the projection removes the full component regardless of scale).
 
     Args:
         directions: ``{address: tensor [d_model]}`` directions to ablate. Keys
@@ -85,12 +87,18 @@ def fn(activation: Tensor, key: Node) -> Tensor:
 def steer_direction(directions: dict[Node, Tensor], alpha: float) -> Callable:
     """Return an intervention function that adds a scaled direction.
 
-    Adds alpha * direction to the residual stream at each layer/module.
+    Each direction is unit-normalized first, then ``alpha * unit_direction`` is
+    added to the residual stream at each layer/module. So ``alpha`` is the
+    absolute magnitude added (in residual-norm units), not a multiplier on the
+    stored vector's length: the direction's original magnitude is discarded, and
+    an upstream ``SteeringVector(normalize=False)`` therefore does not change the
+    steering strength (scale it with ``alpha`` instead).
 
     Args:
         directions: ``{address: tensor [d_model]}`` directions to add. Keys are
             coerced to canonical :class:`Node` addresses.
-        alpha: Scaling factor. Positive = strengthen, negative = suppress.
+        alpha: Absolute magnitude added along the unit direction. Positive =
+            strengthen, negative = suppress.
 
     Returns:
         Callable(activation, node) -> modified activation.
@@ -274,10 +282,40 @@ def _resolve_fn(self, results: Results) -> Callable:
         # ablate_direction normalize the direction keys to canonical Nodes.
         steering = results[self.direction_key]
         directions = self._select_directions(steering)
+        self._check_reachable(directions)
         if self.mode == "steer":
             return steer_direction(directions, self.alpha)
         return ablate_direction(directions)
 
+    def _check_reachable(self, directions: dict[Node, Tensor]) -> None:
+        """Fail loudly (or warn) when directions land on no hooked site.
+
+        The quick-API ``model.generate`` guards this; the step must too, or a
+        ``modules=`` that does not match where the direction was recorded (e.g.
+        a ``resid_post`` direction with ``modules="mlp"``) silently returns the
+        clean generation as the "intervened" one and reads as a null result.
+        """
+        if not directions:
+            return
+        missing = unreachable_addresses(directions, self.layers, self.modules)
+        if len(missing) == len(directions):
+            raise ValueError(
+                f"No steering direction matches a hooked site "
+                f"(layers={self.layers!r}, modules={self.modules!r}); the "
+                f"intervention would do nothing. Directions: {sorted(directions)}. "
+                f"Align modules= with where the direction was recorded (a "
+                f"SteeringVector direction defaults to 'resid_post')."
+            )
+        if missing:
+            logger.warning(
+                "%d steering direction(s) match no hooked site (layers=%r, "
+                "modules=%r) and will be skipped: %s",
+                len(missing),
+                self.layers,
+                self.modules,
+                sorted(missing),
+            )
+
     def _select_directions(self, steering: object) -> dict[Node, Tensor]:
         """Return the subset of recorded directions ``direction_layers`` selects.
 
diff --git a/src/murano/steps/logit_attribution.py b/src/murano/steps/logit_attribution.py
index d987df7..11d2420 100644
--- a/src/murano/steps/logit_attribution.py
+++ b/src/murano/steps/logit_attribution.py
@@ -331,6 +331,26 @@ def _true_target(
             target = target - _gather_answer(answer_logits, incorrect_ids)
         return target
 
+    def _uses_unit_offset_norm(self, ln: Any) -> bool:
+        """Whether the final norm scales by ``(1 + weight)`` instead of ``weight``.
+
+        Gemma-family RMSNorm (``Gemma2RMSNorm`` etc.) applies ``x_normed * (1 +
+        weight)``; the frozen-norm decomposition must add 1 to the stored weight
+        or the reconstruction misses the identity term and completeness breaks.
+        Detected by the norm class name or the model type, both robust to the
+        nnsight Envoy wrapper.
+        """
+        try:
+            if "Gemma" in type(getattr(ln, "_module", ln)).__name__:
+                return True
+        except Exception:  # pragma: no cover - defensive
+            pass
+        try:
+            config = getattr(self.model.hf_model, "config", None)
+            return str(getattr(config, "model_type", "")).startswith("gemma")
+        except Exception:  # pragma: no cover - defensive
+            return False
+
     def _frozen_norm_params(
         self, r_pos: Tensor
     ) -> tuple[Tensor, Tensor, Tensor | None]:
@@ -339,9 +359,13 @@ def _frozen_norm_params(
         LayerNorm centers and divides by the residual's standard deviation;
         RMSNorm divides by its root-mean-square. ``beta`` is returned only for
         LayerNorm, where it is a constant added once to the normalized vector.
+        Gemma-family RMSNorm scales by ``(1 + weight)``, so its effective gamma
+        adds one to the stored weight.
         """
         ln = self.model.final_norm
         gamma = ln.weight.to(device=r_pos.device, dtype=float32)
+        if self._uses_unit_offset_norm(ln):
+            gamma = gamma + 1.0
         eps = getattr(ln, "eps", None)
         if eps is None:
             eps = getattr(ln, "variance_epsilon", 1e-5)
diff --git a/src/murano/steps/logit_lens.py b/src/murano/steps/logit_lens.py
index 0ad8d14..3c2757e 100644
--- a/src/murano/steps/logit_lens.py
+++ b/src/murano/steps/logit_lens.py
@@ -30,7 +30,11 @@ class LogitLensResult:
     """Per-layer next-token probability distributions.
 
     Attributes:
-        all_probs: Tensor [n_layers, n_inputs, seq, vocab] of softmax outputs.
+        all_probs: Tensor [n_layers, n_inputs, seq, vocab] of softmax outputs, or
+            ``None``. Kept only when the step ran with ``store_full_probs=True``;
+            it is ``n_layers * n_inputs * seq * vocab`` floats, which is tens to
+            hundreds of gigabytes on a real model, so it is off by default. The
+            reduced fields below (and the standard heatmap) do not need it.
         max_probs: Tensor [n_layers, n_inputs, seq] of argmax probabilities.
         predicted_tokens: Tensor [n_layers, n_inputs, seq] of argmax token IDs.
         predicted_words: Decoded tokens, shape [n_layers][n_inputs][seq].
@@ -42,7 +46,7 @@ class LogitLensResult:
             output, so these are ``resid_post`` nodes.
     """
 
-    all_probs: Tensor
+    all_probs: Tensor | None
     max_probs: Tensor
     predicted_tokens: Tensor
     predicted_words: list[list[list[str]]]
@@ -70,6 +74,13 @@ class LogitLens(Step):
     Args:
         model: MuranoModel to record from.
         layers: Layer indices to record, or ``"all"`` for every layer.
+        store_full_probs: Keep the full ``[n_layers, n_inputs, seq, vocab]``
+            softmax on the result. Off by default: that tensor is tens to
+            hundreds of gigabytes on a real model (every layer, position, and
+            vocabulary entry), and the reduced ``max_probs`` / ``predicted_tokens``
+            / ``predicted_words`` fields (and the standard heatmap) do not need
+            it. Turn it on only for a targeted, small run that inspects the full
+            distribution.
 
     Raises:
         ValueError: If ``layers`` is a string other than ``"all"``.
@@ -84,8 +95,10 @@ def __init__(
         self,
         model: ModelBackend,
         layers: list[int] | str = "all",
+        store_full_probs: bool = False,
     ):
         self.model = model
+        self.store_full_probs = store_full_probs
         if isinstance(layers, str):
             if layers != "all":
                 raise ValueError(f"layers as string must be 'all', got {layers!r}")
@@ -113,6 +126,13 @@ def __call__(self, results: Results) -> Results:
             for layer in self.layers:
                 saved[layer] = self.model.layer(layer).output.save()
 
+        # Reduce each layer's [n_inputs, seq, vocab] softmax to the argmax
+        # probability and token as we go, so the full-vocab tensor for a layer is
+        # transient. Keeping every layer's full distribution (the old behavior) is
+        # n_layers * n_inputs * seq * vocab floats: tens to hundreds of GB on a
+        # real model. Retain it only when the caller opts in via store_full_probs.
+        layer_max: list[Tensor] = []
+        layer_pred: list[Tensor] = []
         layer_probs: list[Tensor] = []
         with no_grad():
             for layer in self.layers:
@@ -124,12 +144,19 @@ def __call__(self, results: Results) -> Results:
                 if isinstance(output, tuple):
                     output = output[0]
                 # output: [n_inputs, seq, d_model]; project to vocab.
-                logits = self.model.project_on_vocab(output)
-                layer_probs.append(softmax(logits, dim=-1).detach().cpu())
-
-        # all_probs: [n_layers, n_inputs, seq, vocab]
-        all_probs = stack(layer_probs, dim=0)
-        max_probs, predicted_tokens = all_probs.max(dim=-1)
+                probs = (
+                    softmax(self.model.project_on_vocab(output), dim=-1).detach().cpu()
+                )
+                mp, pt = probs.max(dim=-1)
+                layer_max.append(mp)
+                layer_pred.append(pt)
+                if self.store_full_probs:
+                    layer_probs.append(probs)
+
+        # max_probs / predicted_tokens: [n_layers, n_inputs, seq]
+        max_probs = stack(layer_max, dim=0)
+        predicted_tokens = stack(layer_pred, dim=0)
+        all_probs = stack(layer_probs, dim=0) if self.store_full_probs else None
 
         input_words = [
             [self.model.tokenizer.decode([int(t)]) for t in seq.tolist()]
diff --git a/src/murano/steps/metrics.py b/src/murano/steps/metrics.py
index 24a9aef..4d81525 100644
--- a/src/murano/steps/metrics.py
+++ b/src/murano/steps/metrics.py
@@ -188,8 +188,10 @@ def _answer_positions(
     Resolution order: an explicit ``positions`` wins (negative indices count
     from the end); otherwise the last real token is read from
     ``attention_mask`` (correct for either padding side); otherwise the final
-    column is used, which is valid for the left padding Murano's tokenizers
-    default to.
+    column is used, which is the last real token only when the batch is
+    unpadded (a single prompt, or equal-length prompts). Murano does not set a
+    tokenizer padding side, so pass ``attention_mask`` for any padded batch
+    rather than relying on the final-column fallback.
 
     Args:
         logits: Output logits ``[B, S, V]``.
diff --git a/src/murano/steps/path_patch.py b/src/murano/steps/path_patch.py
index 7e89d1c..7f347c9 100644
--- a/src/murano/steps/path_patch.py
+++ b/src/murano/steps/path_patch.py
@@ -2,12 +2,15 @@
 
 Activation patching (:class:`~murano.steps.patch.Patch`) replaces one site and lets
 the change propagate through every downstream route, so it measures a component's
-*total* effect. Path patching measures the *direct* effect: it injects the sender's
-activation from a second run but freezes every other component at its base value, so
-the perturbation can only reach the receiver along the direct residual path. This is
-the primitive behind circuit localization (name movers, s-inhibition, backup name
-movers in Wang et al.'s IOI work); the algorithm is the one from Goldowsky-Dill et
-al., "Localizing Model Behavior with Path Patching".
+*total* effect. Path patching isolates the sender path: it injects the sender's
+activation from a second run while freezing every other attention head at its base
+value (and, when ``freeze_mlps=True``, the MLPs too). With the MLPs frozen the
+perturbation can only reach the receiver along the direct residual path; with the
+default ``freeze_mlps=False`` the MLPs recompute and mediate it, so the measured
+quantity is the attention-mediated effect. This is the primitive behind circuit
+localization (name movers, s-inhibition, backup name movers in Wang et al.'s IOI
+work); the algorithm is the one from Goldowsky-Dill et al., "Localizing Model
+Behavior with Path Patching".
 
 For senders ``S`` and receiver ``R``, with ``base`` = the run to measure and
 ``source`` = the run supplying the sender activations:
@@ -177,10 +180,22 @@ class PathPatch(Step):
     """Path-patch senders into a receiver and write the resulting logits.
 
     Runs the base prompts while injecting the source prompts' sender activations
-    along only the direct path to the receiver (every other component frozen), then
-    stores the logits so a metric step can score the direct effect. With the
-    defaults (base = clean ``prompts``, source = corrupt ``corrupt_prompts``) this
-    reads how much of the behaviour the direct sender path carries.
+    and freezing every attention head at its base value, then stores the logits so
+    a metric step can score the sender path. With the defaults (base = clean
+    ``prompts``, source = corrupt ``corrupt_prompts``) this reads how much of the
+    behaviour the sender path carries.
+
+    By default (``freeze_mlps=False``) the MLPs recompute, so they mediate the
+    sender's perturbation and the measured quantity is the attention-mediated
+    effect, not the pure direct residual path. Set ``freeze_mlps=True`` to freeze
+    the MLPs too, leaving only the direct sender-to-receiver residual path.
+
+    Note the default direction is the reverse of :class:`~murano.steps.patch.Patch`:
+    ``Patch`` defaults to denoising (base = corrupt, source = clean) while
+    ``PathPatch`` defaults to noising (base = clean ``prompts``, source = corrupt
+    ``corrupt_prompts``). So with :class:`~murano.steps.metrics.RecoveredMetricStep`
+    the recovered fraction reads oppositely: swap ``base_key`` / ``source_key`` (or
+    read the raw metric) if you want the denoising convention.
 
     Reads from results:
         results[base_key]: PromptBatch (default ``prompts``), the run to measure.
diff --git a/src/murano/steps/select.py b/src/murano/steps/select.py
index 4af5723..d803c0a 100644
--- a/src/murano/steps/select.py
+++ b/src/murano/steps/select.py
@@ -126,7 +126,10 @@ def __init__(
             self.modules = {canonical_module(name) for name in names}
         self.output_key = output_key
         self.reads = [source_key]
-        self.read_types = {source_key: (LogitAttributionResult, SweepResult)}
+        # _extract_scores also accepts a bare {Node: float} mapping, so include
+        # dict here or the pipeline's pre-flight type check would reject the very
+        # input the step supports.
+        self.read_types = {source_key: (LogitAttributionResult, SweepResult, dict)}
         self.writes = [output_key]
         self.write_types = {output_key: ComponentSelection}
 
diff --git a/src/murano/steps/sweep.py b/src/murano/steps/sweep.py
index 4019464..d180e9e 100644
--- a/src/murano/steps/sweep.py
+++ b/src/murano/steps/sweep.py
@@ -11,7 +11,10 @@
 pipeline built before it), then forks the ``Results`` per item so each item sees
 the same starting state, and collects the harvested keys into a
 :class:`~murano.artifacts.SweepResult`. Because the forks are shallow copies, the
-swept steps' writes are scratch: they never reach the pipeline that follows.
+swept steps' writes are scratch: they never reach the pipeline that follows. The
+isolation is key-level only, though: a swept step that mutates a shared value in
+place (rather than writing a new key) would leak that change across items, so
+swept steps must treat the incoming ``Results`` values as read-only.
 
 The swept steps are built by a callback, so what varies is not restricted to a
 component. Sweeping over :class:`~murano.nodes.Node` addresses yields a component
diff --git a/src/murano/steps/train.py b/src/murano/steps/train.py
index b847972..d070f8b 100644
--- a/src/murano/steps/train.py
+++ b/src/murano/steps/train.py
@@ -48,7 +48,12 @@ class SteeringVector(Step):
     Args:
         method: Estimation method. Currently only ``"contrastive_mean_diff"``
             is supported.
-        normalize: If True, normalize each per-layer direction to unit norm.
+        normalize: If True, normalize each stored per-layer direction to unit
+            norm. Note the intervention functions (``steer_direction`` /
+            ``ablate_direction``) unit-normalize again at apply time, so
+            ``normalize=False`` only affects the magnitude of the stored vector
+            (e.g. if you read ``direction_per_layer`` yourself); it does not
+            change steering or ablation strength.
 
     Raises:
         ValueError: If ``method`` is not a supported value.
diff --git a/src/murano/tasks.py b/src/murano/tasks.py
index 5d211f0..b171105 100644
--- a/src/murano/tasks.py
+++ b/src/murano/tasks.py
@@ -34,8 +34,10 @@
 
 IOI_TEMPLATE: Final = "When {a} and {b} went to the store, {giver} gave a drink to"
 
-# Each pair is (indirect object, subject). Both names are single tokens with a
-# leading space in GPT-2's vocabulary, so the answer is one token either way.
+# Each pair is (indirect object, subject), chosen so both names tokenize to a
+# single token with a leading space in GPT-2's vocabulary (the metric scores the
+# first token of a multi-token answer, so single-token names keep the logit
+# difference exact).
 _IOI_NAMES: Final = [
     ("Mary", "John"),
     ("Alice", "Bob"),
diff --git a/tests/test_auto_dtype.py b/tests/test_auto_dtype.py
new file mode 100644
index 0000000..9c91cf6
--- /dev/null
+++ b/tests/test_auto_dtype.py
@@ -0,0 +1,78 @@
+"""Auto-dtype heuristic: small models -> float32, large -> bfloat16.
+
+Regression guard for the parameter estimate, which must account for tied
+embeddings and grouped-query attention or it misclassifies large-vocab models
+like Gemma-2-2b (a 2.6B model that should load in float32).
+"""
+
+from __future__ import annotations
+
+from murano.model import _FP32_PARAM_LIMIT, _estimate_num_params
+
+
+class _Cfg:
+    def __init__(self, **kw):
+        self.__dict__.update(kw)
+
+
+def test_gpt2_is_small():
+    cfg = _Cfg(
+        n_embd=768, n_layer=12, vocab_size=50257, n_head=12, tie_word_embeddings=True
+    )
+    n = _estimate_num_params(cfg)
+    assert n is not None and n <= _FP32_PARAM_LIMIT, f"gpt2 estimate {n}"
+
+
+def test_gemma2_2b_is_small_despite_large_vocab():
+    # Real Gemma-2-2b: ~2.6B params. Large vocab (256k), GQA (8 heads / 4 kv,
+    # head_dim 256 independent of hidden), tied embeddings.
+    cfg = _Cfg(
+        hidden_size=2304,
+        num_hidden_layers=26,
+        vocab_size=256000,
+        intermediate_size=9216,
+        num_attention_heads=8,
+        head_dim=256,
+        num_key_value_heads=4,
+        tie_word_embeddings=True,
+    )
+    n = _estimate_num_params(cfg)
+    assert n is not None
+    assert n <= _FP32_PARAM_LIMIT, f"gemma-2-2b estimate {n / 1e9:.2f}B should be <= 3B"
+    assert 2.0e9 < n < 3.0e9, f"gemma-2-2b estimate {n / 1e9:.2f}B off from real ~2.6B"
+
+
+def test_llama_3_2_1b_is_small():
+    cfg = _Cfg(
+        hidden_size=2048,
+        num_hidden_layers=16,
+        vocab_size=128256,
+        intermediate_size=8192,
+        num_attention_heads=32,
+        head_dim=64,
+        num_key_value_heads=8,
+        tie_word_embeddings=True,
+    )
+    n = _estimate_num_params(cfg)
+    assert n is not None and n <= _FP32_PARAM_LIMIT, f"llama-3.2-1b estimate {n}"
+
+
+def test_qwen2_5_7b_is_large():
+    cfg = _Cfg(
+        hidden_size=3584,
+        num_hidden_layers=28,
+        vocab_size=152064,
+        intermediate_size=18944,
+        num_attention_heads=28,
+        head_dim=128,
+        num_key_value_heads=4,
+        tie_word_embeddings=False,
+    )
+    n = _estimate_num_params(cfg)
+    assert n is not None and n > _FP32_PARAM_LIMIT, (
+        f"qwen2.5-7b estimate {n / 1e9:.2f}B"
+    )
+
+
+def test_missing_fields_returns_none():
+    assert _estimate_num_params(_Cfg(hidden_size=1024)) is None
diff --git a/tests/test_logit_attribution.py b/tests/test_logit_attribution.py
index e3e9a1e..04fca60 100644
--- a/tests/test_logit_attribution.py
+++ b/tests/test_logit_attribution.py
@@ -253,3 +253,83 @@ def test_plot_returns_figure(self, murano_model):
         fig = plot_logit_attribution(result, top_k=5)
         assert fig.data
         assert len(fig.data[0].y) <= 5
+
+
+# ── Gemma-family RMSNorm (1 + weight) regression ──────────────────────
+
+
+class TestLogitAttributionGemmaNorm:
+    """Gemma RMSNorm scales by (1 + weight); the frozen-norm DLA must match.
+
+    Without the unit-offset correction the reconstruction misses the identity
+    term and completeness_error blows up (seen as ~18 on real gemma-2-2b). This
+    builds a tiny Gemma-2 and asserts completeness holds.
+    """
+
+    @pytest.fixture(scope="class")
+    def gemma_model(self, tmp_path_factory):
+        from pathlib import Path
+
+        from tokenizers import Tokenizer
+        from tokenizers.models import WordLevel
+        from tokenizers.pre_tokenizers import Whitespace
+        from transformers import (
+            Gemma2Config,
+            Gemma2ForCausalLM,
+            PreTrainedTokenizerFast,
+        )
+
+        from murano.model import MuranoModel
+
+        vocab = {
+            "<pad>": 0,
+            "<s>": 1,
+            "</s>": 2,
+            "<unk>": 3,
+            "hello": 4,
+            "world": 5,
+            "good": 6,
+            "bad": 7,
+        }
+        path = Path(tmp_path_factory.mktemp("tiny_gemma2"))
+        tok = Tokenizer(WordLevel(vocab=dict(vocab), unk_token="<unk>"))
+        tok.pre_tokenizer = Whitespace()
+        PreTrainedTokenizerFast(
+            tokenizer_object=tok,
+            unk_token="<unk>",
+            pad_token="<pad>",
+            bos_token="<s>",
+            eos_token="</s>",
+            model_max_length=64,
+        ).save_pretrained(path)
+        config = Gemma2Config(
+            vocab_size=len(vocab),
+            hidden_size=32,
+            intermediate_size=64,
+            num_hidden_layers=2,
+            num_attention_heads=4,
+            num_key_value_heads=2,
+            head_dim=8,
+            max_position_embeddings=64,
+            sliding_window=64,
+            pad_token_id=0,
+            bos_token_id=1,
+            eos_token_id=2,
+        )
+        torch.manual_seed(0)
+        Gemma2ForCausalLM(config).save_pretrained(path)
+        return MuranoModel(str(path), device_map="cpu", dtype=torch.float32)
+
+    def test_completeness_holds_on_gemma(self, gemma_model):
+        result = _run(gemma_model, correct=4, incorrect=5)
+        assert result.metadata["norm"] == "rmsnorm"
+        # With the (1 + weight) correction this is ~machine-epsilon; without it,
+        # it is large (order 1-20).
+        assert result.completeness_error < 1e-2, (
+            f"gemma DLA completeness_error={result.completeness_error} — the "
+            f"frozen norm is not reconstructing Gemma's (1 + weight) RMSNorm"
+        )
+
+    def test_unit_offset_norm_detected(self, gemma_model):
+        step = LogitAttribution(gemma_model, correct=4)
+        assert step._uses_unit_offset_norm(gemma_model.final_norm) is True
diff --git a/tests/test_logit_lens.py b/tests/test_logit_lens.py
index 68a06de..328cb7f 100644
--- a/tests/test_logit_lens.py
+++ b/tests/test_logit_lens.py
@@ -92,15 +92,35 @@ def test_pipeline_produces_logit_lens_result(self, model):
         assert isinstance(result, LogitLensResult)
         n_layers = model.n_layers
         n_inputs = 2
+        # The full-vocab tensor is off by default (it OOMs on real models); the
+        # reduced fields carry the standard logit-lens signal.
+        assert result.all_probs is None
+        assert result.max_probs.ndim == 3
+        assert result.max_probs.shape[0] == n_layers
+        assert result.max_probs.shape[1] == n_inputs
+        assert result.predicted_tokens.shape == result.max_probs.shape
+        assert result.addresses == [Node(i, RESID_POST) for i in range(n_layers)]
+        assert len(result.input_words) == n_inputs
+        assert len(result.predicted_words) == n_layers
+
+    def test_store_full_probs_keeps_full_vocab_tensor(self, model):
+        pipe = Pipeline(
+            [
+                LoadPrompts(["hello world", "good world"]),
+                LogitLens(model, store_full_probs=True),
+            ]
+        )
+        result = pipe.run()["logit_lens"]
+        n_layers, n_inputs = model.n_layers, 2
         assert result.all_probs.ndim == 4
         assert result.all_probs.shape[0] == n_layers
         assert result.all_probs.shape[1] == n_inputs
         assert result.all_probs.shape[3] == model.tokenizer.vocab_size
         assert result.max_probs.shape == result.all_probs.shape[:-1]
-        assert result.predicted_tokens.shape == result.all_probs.shape[:-1]
-        assert result.addresses == [Node(i, RESID_POST) for i in range(n_layers)]
-        assert len(result.input_words) == n_inputs
-        assert len(result.predicted_words) == n_layers
+        # The reduced fields must equal the reduction of the full tensor.
+        mp, pt = result.all_probs.max(dim=-1)
+        assert torch.allclose(result.max_probs, mp, atol=1e-6)
+        assert torch.equal(result.predicted_tokens, pt)
 
     def test_layers_subset(self, model):
         pipe = Pipeline(
@@ -111,7 +131,7 @@ def test_layers_subset(self, model):
         )
         results = pipe.run()
         result = results["logit_lens"]
-        assert result.all_probs.shape[0] == 1
+        assert result.max_probs.shape[0] == 1
         assert result.addresses == [Node(0, RESID_POST)]
 
     def test_invalid_layers_string_raises(self, model):
@@ -122,7 +142,7 @@ def test_probabilities_sum_to_one(self, model):
         pipe = Pipeline(
             [
                 LoadPrompts(["hello world"]),
-                LogitLens(model),
+                LogitLens(model, store_full_probs=True),
             ]
         )
         results = pipe.run()
@@ -130,6 +150,23 @@ def test_probabilities_sum_to_one(self, model):
         sums = all_probs.sum(dim=-1)
         assert torch.allclose(sums, torch.ones_like(sums), atol=1e-4)
 
+    def test_last_layer_lens_matches_model_logits(self, model):
+        """Ground truth: the last-layer lens is the model's real output.
+
+        ``project_on_vocab`` applies the model's own final norm and unembedding,
+        so the logit lens at the final ``resid_post`` must reproduce the model's
+        true next-token distribution. This catches a wrong norm, a wrong layer,
+        or a wrong projection, which a shape/sum-to-one check cannot.
+        """
+        prompts = ["hello world", "good world"]
+        result = Pipeline([LoadPrompts(prompts), LogitLens(model)]).run()["logit_lens"]
+
+        true_logits = model.logits(prompts)  # [B, S, V], the model's real output
+        true_pred = true_logits.argmax(dim=-1)  # [B, S]
+        # predicted_tokens is [n_layers, B, S]; the last layer is the model output.
+        last_layer_pred = result.predicted_tokens[-1]  # [B, S]
+        assert torch.equal(last_layer_pred, true_pred)
+
 
 class TestLogitLensSave:
     """Round-trip persistence via Pipeline + Save."""
@@ -138,7 +175,7 @@ def test_save_writes_logit_lens_artifact(self, model, tmp_path):
         pipe = Pipeline(
             [
                 LoadPrompts(["hello world", "good world"]),
-                LogitLens(model, layers=[0, 1]),
+                LogitLens(model, layers=[0, 1], store_full_probs=True),
                 Save(output_dir=str(tmp_path)),
             ]
         )
@@ -183,7 +220,7 @@ def test_load_logit_lens_roundtrip(self, model, tmp_path):
         results = Pipeline(
             [
                 LoadPrompts(["hello world", "good world"]),
-                LogitLens(model, layers=[0, 1]),
+                LogitLens(model, layers=[0, 1], store_full_probs=True),
                 Save(output_dir=str(tmp_path)),
             ]
         ).run()
diff --git a/tests/test_plotting_probing.py b/tests/test_plotting_probing.py
new file mode 100644
index 0000000..a5a7604
--- /dev/null
+++ b/tests/test_plotting_probing.py
@@ -0,0 +1,64 @@
+"""Smoke tests for the probing plots (previously untested public API)."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import torch
+
+pytest.importorskip("plotly")
+pytest.importorskip("sklearn")
+
+from sklearn.linear_model import LogisticRegression
+
+from murano.nodes import Node, RESID_POST
+from murano.plotting.probing import plot_confusion_matrix
+from murano.steps.probe import ProbeResult
+from murano.steps.record import LabeledActivationStore
+
+
+def _labeled_store(n_per_class: int = 8, d: int = 6):
+    # Two linearly separable clusters so the refit classifier is meaningful.
+    g = torch.Generator().manual_seed(0)
+    pos = torch.randn(n_per_class, d, generator=g) + 3.0
+    neg = torch.randn(n_per_class, d, generator=g) - 3.0
+    acts = torch.cat([pos, neg], dim=0)
+    labels = torch.tensor([1] * n_per_class + [0] * n_per_class)
+    return LabeledActivationStore(
+        activations={Node(0, RESID_POST): acts},
+        labels=labels,
+        position="last",
+        per_head=False,
+    )
+
+
+def _probe_result(store, refit: bool):
+    node = Node(0, RESID_POST)
+    classifiers = {}
+    if refit:
+        clf = LogisticRegression(max_iter=200)
+        clf.fit(store.activations[node].numpy(), store.labels.numpy())
+        classifiers[node] = clf
+    return ProbeResult(
+        accuracy_per_layer={node: 1.0},
+        cv_scores={node: np.array([1.0, 1.0])},
+        best_layer=node,
+        classifiers=classifiers,
+        label_names=["neg", "pos"],
+    )
+
+
+def test_plot_confusion_matrix_returns_figure_when_refit():
+    store = _labeled_store()
+    probe = _probe_result(store, refit=True)
+    fig = plot_confusion_matrix(probe, store)
+    assert fig is not None
+    data = fig.to_dict()
+    assert data["data"][0]["type"] == "heatmap"
+
+
+def test_plot_confusion_matrix_returns_none_without_refit():
+    store = _labeled_store()
+    probe = _probe_result(store, refit=False)
+    # No refitted classifier at the best layer -> None, not a crash.
+    assert plot_confusion_matrix(probe, store) is None
diff --git a/tests/test_sae.py b/tests/test_sae.py
index 7c47432..d029c8e 100644
--- a/tests/test_sae.py
+++ b/tests/test_sae.py
@@ -544,6 +544,103 @@ def test_max_length_truncates_when_set(self, model):
         assert store.tokens.shape[1] == 4
 
 
+class _IdentitySAE(_FakeSAE):
+    """A fake SAE whose encode is the identity, so the captured residual is
+    readable straight off ``store.activations``.
+
+    Lets a test assert the *value* SAEEncode captured for each hook kind, which
+    the zero-returning ``_FakeSAE`` hides. Requires ``d_sae == d_model``.
+    """
+
+    def __init__(self, d_model: int, hook_name: str, hook_layer: int):
+        super().__init__(
+            d_sae=d_model, d_model=d_model, hook_name=hook_name, hook_layer=hook_layer
+        )
+
+    def encode(self, residual: torch.Tensor) -> torch.Tensor:
+        return residual.clone()
+
+
+class TestSAEEncodeCaptureValues:
+    """SAEEncode captures the right residual for each hook kind (value-level)."""
+
+    def _expected_residual(self, model, prompts, layer, kind):
+        """Independently trace the residual SAEEncode should capture for ``kind``."""
+        from murano._proxy import unwrap_traced
+
+        # Single prompt, so no padding: token ids match SAEEncode's tokenization
+        # regardless of padding side.
+        tokens = model.tokenizer(
+            prompts, return_tensors="pt", return_token_type_ids=False
+        )
+        saved = {}
+        with model.trace(tokens):
+            if kind == "resid_post":
+                saved["v"] = model.layer(layer).output.save()
+            elif kind == "resid_pre":
+                saved["v"] = model.layer(layer).input.save()
+            else:  # mlp_out
+                saved["v"] = model.resolve_module(layer, "mlp").output.save()
+        value = unwrap_traced(saved["v"])
+        if isinstance(value, tuple):
+            value = value[0]
+        return value.detach().float().cpu()
+
+    @pytest.mark.parametrize(
+        "kind,hook_name",
+        [
+            ("resid_post", "blocks.1.hook_resid_post"),
+            ("resid_pre", "blocks.1.hook_resid_pre"),
+            ("mlp_out", "blocks.1.hook_mlp_out"),
+        ],
+    )
+    def test_captured_residual_matches_independent_trace(self, model, kind, hook_name):
+        from murano.artifacts import PromptBatch
+        from murano.results import Results
+
+        prompts = ["hello world good"]
+        d_model = model._lm.lm_head.in_features
+
+        results = Results()
+        results["prompts"] = PromptBatch(prompts=prompts)
+        step = SAEEncode(model, release="test/repo", sae_id="test/id")
+        step.sae_model._sae = _IdentitySAE(
+            d_model=d_model, hook_name=hook_name, hook_layer=1
+        )
+        store = step(results)["sae_record"]
+
+        # Under the identity encode, activations == the residual SAEEncode fed in.
+        expected = self._expected_residual(model, prompts, layer=1, kind=kind)
+        assert store.activations.shape == expected.shape
+        assert torch.allclose(store.activations, expected, atol=1e-5), (
+            f"{kind}: SAEEncode captured a different residual than a direct trace "
+            f"of {hook_name}"
+        )
+
+    def test_resid_pre_differs_from_resid_post(self, model):
+        """A guard that the hook-kind mapping is not collapsing to one site."""
+        from murano.artifacts import PromptBatch
+        from murano.results import Results
+
+        d_model = model._lm.lm_head.in_features
+
+        def encode_kind(hook_name):
+            step = SAEEncode(model, release="r", sae_id="i")
+            step.sae_model._sae = _IdentitySAE(
+                d_model=d_model, hook_name=hook_name, hook_layer=1
+            )
+            res = Results()
+            res["prompts"] = PromptBatch(prompts=["hello world good"])
+            return step(res)["sae_record"].activations
+
+        pre = encode_kind("blocks.1.hook_resid_pre")
+        post = encode_kind("blocks.1.hook_resid_post")
+        assert not torch.allclose(pre, post), (
+            "resid_pre and resid_post captured the same tensor; the hook-kind "
+            "selection is not distinguishing the sites."
+        )
+
+
 class TestSAETopActivations:
     """Synthetic activations, deterministic top-K behavior."""
 
diff --git a/tests/test_weight_ablation.py b/tests/test_weight_ablation.py
new file mode 100644
index 0000000..75a7e61
--- /dev/null
+++ b/tests/test_weight_ablation.py
@@ -0,0 +1,174 @@
+"""Value-level tests for the WeightAblation projection.
+
+The step-level suite (test_steps.py) monkeypatches ``ablate_model_weights`` out,
+so the real read/write projection (``W @ P`` for read matrices, ``P @ W`` for
+write matrices), the embedding special-case, and the architecture guard are never
+exercised. These tests run the real projection on a fresh tiny model and assert
+the direction is annihilated on the correct side of each matrix.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+import torch
+
+pytest.importorskip("nnsight")
+pytest.importorskip("transformers")
+
+from tokenizers import Tokenizer
+from tokenizers.models import WordLevel
+from tokenizers.pre_tokenizers import Whitespace
+from transformers import (
+    GPT2Config,
+    GPT2LMHeadModel,
+    LlamaConfig,
+    LlamaForCausalLM,
+    PreTrainedTokenizerFast,
+)
+
+from murano.model import MuranoModel
+from murano.steps.weight_ablation import (
+    ProjectionOperator,
+    ablate_model_weights,
+)
+
+_VOCAB = {
+    "<pad>": 0,
+    "<s>": 1,
+    "</s>": 2,
+    "<unk>": 3,
+    "hello": 4,
+    "world": 5,
+    "good": 6,
+    "bad": 7,
+}
+
+
+def _save_tokenizer(path: Path) -> None:
+    tok = Tokenizer(WordLevel(vocab=dict(_VOCAB), unk_token="<unk>"))
+    tok.pre_tokenizer = Whitespace()
+    PreTrainedTokenizerFast(
+        tokenizer_object=tok,
+        unk_token="<unk>",
+        pad_token="<pad>",
+        bos_token="<s>",
+        eos_token="</s>",
+        model_max_length=64,
+    ).save_pretrained(path)
+
+
+def _build_tiny_llama(path: Path) -> None:
+    _save_tokenizer(path)
+    config = LlamaConfig(
+        vocab_size=len(_VOCAB),
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        max_position_embeddings=64,
+        pad_token_id=_VOCAB["<pad>"],
+        bos_token_id=_VOCAB["<s>"],
+        eos_token_id=_VOCAB["</s>"],
+    )
+    torch.manual_seed(0)
+    LlamaForCausalLM(config).save_pretrained(path)
+
+
+def _build_tiny_gpt2(path: Path) -> None:
+    _save_tokenizer(path)
+    config = GPT2Config(
+        vocab_size=len(_VOCAB),
+        n_embd=32,
+        n_layer=2,
+        n_head=4,
+        n_positions=64,
+        pad_token_id=_VOCAB["<pad>"],
+        bos_token_id=_VOCAB["<s>"],
+        eos_token_id=_VOCAB["</s>"],
+    )
+    torch.manual_seed(0)
+    GPT2LMHeadModel(config).save_pretrained(path)
+
+
+@pytest.fixture
+def llama(tmp_path):
+    # Fresh model per test: ablate_model_weights mutates weights in place.
+    path = tmp_path / "llama"
+    _build_tiny_llama(path)
+    return MuranoModel(str(path), device_map="cpu", dtype=torch.float32)
+
+
+@pytest.fixture
+def gpt2(tmp_path):
+    path = tmp_path / "gpt2"
+    _build_tiny_gpt2(path)
+    return MuranoModel(str(path), device_map="cpu", dtype=torch.float32)
+
+
+def _unit_direction(d_model: int, seed: int) -> torch.Tensor:
+    g = torch.Generator().manual_seed(seed)
+    v = torch.randn(d_model, generator=g)
+    return v / v.norm()
+
+
+class TestWeightAblationProjection:
+    def test_n_modified_counts_every_matrix(self, llama):
+        direction = _unit_direction(llama.d_model, seed=1)
+        n = ablate_model_weights(llama, ProjectionOperator(direction))
+        # embedding (1) + per layer q,k,v,o,gate,up,down (7).
+        assert n == 1 + llama.n_layers * 7
+
+    def test_projection_annihilates_direction_on_correct_side(self, llama):
+        direction = _unit_direction(llama.d_model, seed=2)
+        ablate_model_weights(llama, ProjectionOperator(direction))
+
+        d = direction.to(torch.float32)
+        hf = llama.hf_model
+        for layer_idx in range(llama.n_layers):
+            layer = hf.layers[layer_idx]
+            attn, mlp = layer.self_attn, layer.mlp
+
+            # Read matrices (W @ P): remove the direction from the *input* side,
+            # so W @ d == 0.
+            for read_w in (
+                attn.q_proj.weight,
+                attn.k_proj.weight,
+                attn.v_proj.weight,
+                mlp.gate_proj.weight,
+                mlp.up_proj.weight,
+            ):
+                got = read_w.detach().float() @ d
+                assert torch.allclose(got, torch.zeros_like(got), atol=1e-5)
+
+            # Write matrices (P @ W): remove the direction from the *output* side,
+            # so d @ W == 0.
+            for write_w in (attn.o_proj.weight, mlp.down_proj.weight):
+                got = d @ write_w.detach().float()
+                assert torch.allclose(got, torch.zeros_like(got), atol=1e-5)
+
+        # The embedding writes the residual (W @ P form): each row orthogonal to d.
+        got = hf.embed_tokens.weight.detach().float() @ d
+        assert torch.allclose(got, torch.zeros_like(got), atol=1e-5)
+
+    def test_other_direction_is_not_annihilated(self, llama):
+        """Sanity: the projection is targeted, not a global zeroing."""
+        direction = _unit_direction(llama.d_model, seed=3)
+        other = _unit_direction(llama.d_model, seed=99)
+        ablate_model_weights(llama, ProjectionOperator(direction))
+
+        o_proj = llama.hf_model.layers[0].self_attn.o_proj.weight.detach().float()
+        got = other.float() @ o_proj
+        assert not torch.allclose(got, torch.zeros_like(got), atol=1e-3)
+
+    def test_unsupported_architecture_raises_before_mutation(self, gpt2):
+        """GPT-2 (fused c_attn / Conv1D) is not a Llama layout; must raise."""
+        direction = _unit_direction(gpt2.d_model, seed=4)
+        # Snapshot a weight to prove nothing was mutated on the failure path.
+        before = next(gpt2.hf_model.parameters()).detach().clone()
+        with pytest.raises(NotImplementedError, match="Llama-family"):
+            ablate_model_weights(gpt2, ProjectionOperator(direction))
+        after = next(gpt2.hf_model.parameters()).detach()
+        assert torch.equal(before, after)