From fd205cfd1568e40c4444bad27d92379ca3daf2fe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Emirhan=20B=C3=B6ge?= Date: Fri, 31 Jul 2026 12:35:07 +0200 Subject: [PATCH 1/6] feat(model): add grad_forward, a graph-preserving forward pass --- src/murano/backend.py | 13 +++++++++- src/murano/model.py | 56 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 68 insertions(+), 1 deletion(-) diff --git a/src/murano/backend.py b/src/murano/backend.py index 685524a..09209e4 100644 --- a/src/murano/backend.py +++ b/src/murano/backend.py @@ -25,7 +25,7 @@ from torch import Tensor if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Callable, Sequence from murano.nodes import Node @@ -116,6 +116,17 @@ def forward_logits( """ ... + def grad_forward( + self, input_ids: Tensor, layers: "Sequence[int]" + ) -> dict[int, Tensor]: + """Run a native forward pass and return live residual streams per layer. + + Unlike ``forward_logits``, nothing is detached: the returned + ``{layer: [batch, seq, d_model]}`` block outputs stay connected to the + autograd graph so gradient-based steps can differentiate through them. + """ + ... + def generate_with_hooks( self, text: str, diff --git a/src/murano/model.py b/src/murano/model.py index 43e46c2..a608d6e 100644 --- a/src/murano/model.py +++ b/src/murano/model.py @@ -515,6 +515,62 @@ def forward_logits( value = unwrap_traced(saved) return value.detach().float().cpu() + def grad_forward( + self, input_ids: Tensor, layers: Sequence[int] + ) -> dict[int, Tensor]: + """Run a native forward pass and return live residual streams per layer. + + The graph-preserving counterpart of :meth:`forward_logits`: every + proxy-based capture path detaches its tensors, so gradient-based steps + (:class:`~murano.steps.gradients.RecordGradients`, + :class:`~murano.steps.gfc.GFCOperator`) need a forward whose + captured activations stay connected to the autograd graph. This runs + the underlying base module natively, no nnsight trace, with plain + forward hooks on the requested decoder layers, the same raw-module + seam the generation interventions use, and returns the hooked block + outputs without detaching. Callers take logits from the last layer's + residual via :meth:`project_on_vocab`, which is likewise + graph-preserving. + + Args: + input_ids: ``[batch, seq]`` token ids. Pass unpadded sequences + (typically one rollout at a time): the native pass applies no + padding mask, so a padded batch would attend to pad tokens. + layers: Decoder layer indices whose block outputs to return. + + Returns: + ``{layer: hidden}`` with each hidden ``[batch, seq, d_model]`` + still attached to the autograd graph. + + Raises: + RuntimeError: If a requested layer's forward never ran (its hook + captured nothing). + """ + module = self.hf_model + base = getattr(module, "_module", module) + device = next(base.parameters()).device + stash: dict[int, Tensor] = {} + handles = [] + for idx in layers: + + def hook(module, inputs, output, idx=idx): + stash[idx] = output[0] if isinstance(output, tuple) else output + + handles.append(self.raw_layer(idx).register_forward_hook(hook)) + try: + with torch.enable_grad(): + base(input_ids=input_ids.to(device), use_cache=False) + finally: + for handle in handles: + handle.remove() + missing = [idx for idx in layers if idx not in stash] + if missing: + raise RuntimeError( + f"grad_forward captured no output for layer(s) {missing}; " + f"the forward pass did not reach them." + ) + return stash + @property def hf_model(self): """Underlying HuggingFace module. From e9988417accf1793d984ec64983d3c960d804180 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Emirhan=20B=C3=B6ge?= Date: Fri, 31 Jul 2026 12:35:21 +0200 Subject: [PATCH 2/6] feat(steps): record residual-stream gradients over saved rollouts --- src/murano/keys.py | 4 + src/murano/steps/gradients.py | 331 ++++++++++++++++++++++++++++++++++ tests/conftest.py | 13 ++ tests/test_gradients.py | 196 ++++++++++++++++++++ 4 files changed, 544 insertions(+) create mode 100644 src/murano/steps/gradients.py create mode 100644 tests/test_gradients.py diff --git a/src/murano/keys.py b/src/murano/keys.py index 1a3e3cf..2c2cae9 100644 --- a/src/murano/keys.py +++ b/src/murano/keys.py @@ -24,6 +24,10 @@ FEATURE_EXAMPLES: Final = "feature_examples" SELECTION: Final = "selection" SWEEP: Final = "sweep" +ROLLOUTS: Final = "rollouts" +GRADIENT_RECORD: Final = "gradient_record" +GFC_OPERATOR: Final = "gfc_operator" +GSAE: Final = "gsae" METRIC: Final = "metric" WEIGHT_ABLATION: Final = "weight_ablation" OUTPUT_DIR: Final = "output_dir" diff --git a/src/murano/steps/gradients.py b/src/murano/steps/gradients.py new file mode 100644 index 0000000..6a61ecc --- /dev/null +++ b/src/murano/steps/gradients.py @@ -0,0 +1,331 @@ +"""Gradient recording: teacher-forced backward passes over saved rollouts. + +Forward capture (:class:`~murano.steps.record.Record`) answers "what did the +model represent"; this module answers "what would training move". It records, +for a frozen checkpoint, the gradient of the completion log-likelihood with +respect to the residual stream, one vector per completion token. That gradient +is the direction factor of the first-epoch policy-gradient update (the +advantage enters only as a per-rollout scalar), so recording it on saved +rollouts reads the update a trainer *would* apply without ever taking an +optimizer step. + +Two design decisions differ from the rest of the capture code, on purpose: + +* Every other capture path detaches at the earliest opportunity; a gradient + pass cannot, so this module owns its forward end to end through + :meth:`~murano.model.MuranoModel.grad_forward` (the graph-preserving + counterpart of ``forward_logits``) and detaches only at the store boundary. +* Rollouts are processed one at a time rather than in padded batches. The + rollouts of one corpus differ in length, and padding would put the completion + window at a different sequence offset per row; per-rollout passes keep the + prompt/completion boundary exact and keep memory proportional to one + sequence. + +The loss is the plain summed log-likelihood of the realized completion tokens +over a fixed window: no advantage weighting, no normalization. Rewards and +GRPO group structure travel with the rollouts (:class:`RolloutBatch`) so that +downstream analyses can weight per-rollout results by the standardized +advantage, but the recorded gradient itself is reward-free. + +Typical flow:: + + rollouts = RolloutBatch(input_ids=ids, prompt_lengths=lengths) + pipeline = Pipeline([ + LoadRollouts(rollouts), + RecordGradients(model, layer=30, window=512), + ]) + results = pipeline.run() + store = results["gradient_record"] + store.gradients[0] # [T, d_model] gradient at layer 30, first rollout +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +from torch import Tensor, as_tensor, empty, float32, isfinite, log_softmax # pyright: ignore[reportPrivateImportUsage] +from torch import long, tensor, zeros_like # pyright: ignore[reportPrivateImportUsage] + +from murano import keys +from murano.logging import logger +from murano.results import Results +from murano.steps.base import Step + +if TYPE_CHECKING: + from murano.backend import ModelBackend + + +@dataclass +class RolloutBatch: + """A corpus of saved rollouts: full token ids plus the completion boundary. + + A rollout is a prompt and the completion some model sampled for it. For + gradient recording only the fixed token ids matter: every checkpoint is + teacher-forced through the same ids, so results differ only in weights, + never in text. Rewards and group ids are optional passengers for + advantage-weighted analyses. + + Attributes: + input_ids: One 1-D long tensor of token ids per rollout, prompt and + completion concatenated. + prompt_lengths: Number of leading prompt tokens per rollout. Loss terms + and gradient reads start at this offset. Must be at least 1 (the + first completion token is predicted from the last prompt position) + and leave at least one completion token. + rewards: Optional per-rollout scalar verifier rewards. + groups: Optional per-rollout group ids (e.g. the GRPO task group), used + by :meth:`advantages` to standardize rewards within each group. + """ + + input_ids: list[Tensor] + prompt_lengths: list[int] + rewards: list[float] | None = None + groups: list[int] | None = None + + def __post_init__(self) -> None: + if len(self.input_ids) != len(self.prompt_lengths): + raise ValueError( + f"input_ids and prompt_lengths must align: got " + f"{len(self.input_ids)} rollouts but " + f"{len(self.prompt_lengths)} prompt lengths." + ) + coerced: list[Tensor] = [] + for n, ids in enumerate(self.input_ids): + ids = as_tensor(ids, dtype=long) + if ids.dim() != 1: + raise ValueError( + f"Rollout {n}: input_ids must be 1-D token ids, got shape " + f"{tuple(ids.shape)}." + ) + prompt_len = self.prompt_lengths[n] + if not 1 <= prompt_len < ids.shape[0]: + raise ValueError( + f"Rollout {n}: prompt_length must be in [1, {ids.shape[0]}) " + f"so at least one completion token remains, got {prompt_len}." + ) + coerced.append(ids) + self.input_ids = coerced + for name, extra in ( + ("rewards", self.rewards), + ("groups", self.groups), + ): + if extra is not None and len(extra) != len(self.input_ids): + raise ValueError( + f"{name} must have one entry per rollout: got {len(extra)} " + f"for {len(self.input_ids)} rollouts." + ) + + def __len__(self) -> int: + return len(self.input_ids) + + def advantages(self, eps: float = 1e-4) -> Tensor: + """Standardized within-group advantages ``(r - mean) / (std + eps)``. + + Groups with a single outcome (all rollouts passed, or all failed) carry + no within-group signal and get advantage 0 rather than a division + artifact. + + Args: + eps: Stabilizer added to the within-group standard deviation. + + Returns: + Float tensor ``[n_rollouts]`` of advantages. + + Raises: + ValueError: If ``rewards`` or ``groups`` was not provided. + """ + if self.rewards is None or self.groups is None: + raise ValueError( + "advantages() needs both rewards and groups on the RolloutBatch." + ) + rewards = tensor(self.rewards, dtype=float32) + out = zeros_like(rewards) + for group in set(self.groups): + mask = tensor([g == group for g in self.groups]) + r = rewards[mask] + std = r.std(unbiased=False) + if std.item() == 0.0: + continue + out[mask] = (r - r.mean()) / (std + eps) + return out + + +class LoadRollouts(Step): + """Load a rollout corpus into the pipeline results. + + Writes to results: + results['rollouts']: RolloutBatch + + Args: + rollouts: The rollout corpus to make available to downstream steps. + """ + + reads = [] + writes = [keys.ROLLOUTS] + write_types = {keys.ROLLOUTS: RolloutBatch} + + def __init__(self, rollouts: RolloutBatch): + if not isinstance(rollouts, RolloutBatch): + raise TypeError( + f"LoadRollouts expects a RolloutBatch, got {type(rollouts).__name__}." + ) + self.rollouts = rollouts + + def __call__(self, results: Results) -> Results: + results[keys.ROLLOUTS] = self.rollouts + return results + + +@dataclass +class GradientStore: + """Per-rollout residual-stream gradients from a teacher-forced pass. + + Attributes: + gradients: One ``[T, d_model]`` float32 CPU tensor per rollout: the + gradient of the completion log-likelihood with respect to the + residual stream after block ``layer``, at each completion position + inside the window. ``T`` is the per-rollout window length. + nll: ``[n_rollouts]`` mean negative log-likelihood per completion token + over the same window, the functional-change scale used to match + finetuning runs. + finite: ``[n_rollouts]`` bool; False where the backward produced + non-finite values (near-random checkpoints can), in which case that + rollout's gradient tensor is empty and its nll is NaN. + layer: Residual-stream layer the gradients were read at. + window: Maximum completion tokens read per rollout. + """ + + gradients: list[Tensor] + nll: Tensor + finite: Tensor + layer: int + window: int + + def __post_init__(self) -> None: + n = len(self.gradients) + if self.nll.shape[0] != n or self.finite.shape[0] != n: + raise ValueError( + f"nll and finite must have one entry per rollout: got " + f"{self.nll.shape[0]} and {self.finite.shape[0]} for {n} " + f"gradient tensors." + ) + + def __len__(self) -> int: + return len(self.gradients) + + +def _completion_window(ids: Tensor, prompt_length: int, window: int) -> tuple[int, int]: + """Return the ``[start, end)`` sequence positions of the read window.""" + return prompt_length, min(int(ids.shape[0]), prompt_length + window) + + +def _completion_loglik( + model: ModelBackend, hidden_last: Tensor, ids: Tensor, start: int, end: int +) -> Tensor: + """Summed log-likelihood of completion tokens ``ids[start:end]``. + + The token at position ``t`` is predicted from the logits at ``t - 1``, so + the loss reads logits at ``[start-1, end-1)``. Computed in float32 for a + stable softmax regardless of the model dtype. + """ + logits = model.project_on_vocab(hidden_last[:, start - 1 : end - 1]) + logprobs = log_softmax(logits.float(), dim=-1) + targets = ids[start:end].to(logprobs.device) + return logprobs[0].gather(1, targets.unsqueeze(1)).sum() + + +class RecordGradients(Step): + """Record residual-stream gradients of the completion log-likelihood. + + For each rollout, run the frozen model teacher-forced over the saved + tokens, form the summed log-likelihood of the completion tokens inside the + window, and record its gradient with respect to the residual stream after + ``layer``, at every completion position in the window. One forward and one + backward per rollout; the weights are never touched. + + Reads from results: + results['rollouts']: RolloutBatch + + Writes to results: + results['gradient_record']: GradientStore + + Args: + model: Wrapped model to read gradients from. + layer: Residual-stream layer (block output, 0-based) to read at. + window: Maximum completion tokens contributing loss terms and gradient + reads per rollout. The window bounds memory; shorter rollouts use + all their completion tokens. + + Raises: + ValueError: If ``layer`` is out of range or ``window`` is not positive. + + Note: + Gradients are stored per rollout as ``[T, d_model]`` float32 on CPU: + memory is ``n_rollouts * window * d_model`` floats, the gradient-side + analogue of full-position recording. Load large models in float32 for + this step where feasible; bfloat16 gradients are noticeably noisier. + """ + + reads = [keys.ROLLOUTS] + writes = [keys.GRADIENT_RECORD] + read_types = {keys.ROLLOUTS: RolloutBatch} + write_types = {keys.GRADIENT_RECORD: GradientStore} + + def __init__(self, model: ModelBackend, layer: int, window: int = 512): + if not 0 <= layer < model.n_layers: + raise ValueError(f"layer must be in [0, {model.n_layers}), got {layer}.") + if window < 1: + raise ValueError(f"window must be >= 1, got {window}.") + self.model = model + self.layer = layer + self.window = window + + def __call__(self, results: Results) -> Results: + rollouts: RolloutBatch = results[keys.ROLLOUTS] + last = self.model.n_layers - 1 + layers = sorted({self.layer, last}) + + gradients: list[Tensor] = [] + nll: list[float] = [] + finite: list[bool] = [] + logger.info( + "Recording gradients: %d rollouts, layer %d, window %d", + len(rollouts), + self.layer, + self.window, + ) + for n, ids in enumerate(rollouts.input_ids): + start, end = _completion_window( + ids, rollouts.prompt_lengths[n], self.window + ) + hidden = self.model.grad_forward(ids.unsqueeze(0), layers) + loglik = _completion_loglik(self.model, hidden[last], ids, start, end) + (grad,) = torch.autograd.grad(loglik, hidden[self.layer]) + grad_window = grad[0, start:end].detach().float().cpu() + if bool(isfinite(grad_window).all()): + gradients.append(grad_window) + nll.append(-loglik.item() / (end - start)) + finite.append(True) + else: + logger.warning( + "Rollout %d produced non-finite gradients; storing empty.", n + ) + gradients.append(empty(0, self.model.d_model)) + nll.append(float("nan")) + finite.append(False) + + results[keys.GRADIENT_RECORD] = GradientStore( + gradients=gradients, + nll=tensor(nll, dtype=float32), + finite=tensor(finite), + layer=self.layer, + window=self.window, + ) + logger.info( + "Gradient recording done: %d/%d rollouts finite.", + sum(finite), + len(finite), + ) + return results diff --git a/tests/conftest.py b/tests/conftest.py index d5e72c2..29c8907 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,6 +4,7 @@ import sys from pathlib import Path +from typing import TYPE_CHECKING import pytest import torch @@ -18,6 +19,18 @@ PreTrainedTokenizerFast, ) +if TYPE_CHECKING: + from murano.steps.gradients import RolloutBatch + + +def toy_rollouts(n: int, length: int, prompt_len: int, seed: int) -> RolloutBatch: + """Deterministic fake rollouts over the tiny fixture vocab (ids 4..9).""" + from murano.steps.gradients import RolloutBatch + + generator = torch.Generator().manual_seed(seed) + ids = [torch.randint(4, 10, (length,), generator=generator) for _ in range(n)] + return RolloutBatch(input_ids=ids, prompt_lengths=[prompt_len] * n) + _VOCAB = { "": 0, diff --git a/tests/test_gradients.py b/tests/test_gradients.py new file mode 100644 index 0000000..3ae1fc2 --- /dev/null +++ b/tests/test_gradients.py @@ -0,0 +1,196 @@ +"""Tests for gradient recording (RolloutBatch, LoadRollouts, RecordGradients).""" + +from __future__ import annotations + +import pytest +import torch + +from murano import keys +from murano.pipeline import Pipeline +from murano.results import Results +from murano.steps.gradients import ( + GradientStore, + LoadRollouts, + RecordGradients, + RolloutBatch, +) +from tests.conftest import toy_rollouts + + +def _toy_rollouts(n: int = 3, length: int = 14, prompt_len: int = 3) -> RolloutBatch: + return toy_rollouts(n, length, prompt_len, seed=7) + + +def test_rollout_batch_validates_alignment(): + with pytest.raises(ValueError, match="align"): + RolloutBatch(input_ids=[torch.tensor([4, 5, 6])], prompt_lengths=[1, 2]) + + +def test_rollout_batch_rejects_bad_prompt_length(): + with pytest.raises(ValueError, match="prompt_length"): + RolloutBatch(input_ids=[torch.tensor([4, 5, 6])], prompt_lengths=[0]) + with pytest.raises(ValueError, match="prompt_length"): + RolloutBatch(input_ids=[torch.tensor([4, 5, 6])], prompt_lengths=[3]) + + +def test_rollout_batch_rejects_misaligned_extras(): + with pytest.raises(ValueError, match="rewards"): + RolloutBatch( + input_ids=[torch.tensor([4, 5, 6])], + prompt_lengths=[1], + rewards=[1.0, 0.0], + ) + + +def test_rollout_batch_advantages_standardize_within_group(): + batch = RolloutBatch( + input_ids=[torch.tensor([4, 5, 6])] * 5, + prompt_lengths=[1] * 5, + rewards=[1.0, 0.0, 0.0, 1.0, 1.0], + groups=[0, 0, 0, 1, 1], + ) + adv = batch.advantages() + # Group 1 is all-pass: no within-group signal, advantage 0. + assert torch.allclose(adv[3:], torch.zeros(2)) + # Group 0 is centered and the passing rollout sits above the failing ones. + assert abs(adv[:3].mean().item()) < 1e-6 + assert adv[0] > 0 > adv[1] + + +def test_rollout_batch_advantages_require_rewards_and_groups(): + batch = _toy_rollouts() + with pytest.raises(ValueError, match="rewards and groups"): + batch.advantages() + + +def test_load_rollouts_writes_key(): + batch = _toy_rollouts() + results = LoadRollouts(batch)(Results()) + assert results[keys.ROLLOUTS] is batch + + +def test_load_rollouts_rejects_non_batch(): + with pytest.raises(TypeError, match="RolloutBatch"): + LoadRollouts([torch.tensor([4, 5])]) # pyright: ignore[reportArgumentType] + + +def test_record_gradients_validates_arguments(murano_model): + with pytest.raises(ValueError, match="layer"): + RecordGradients(murano_model, layer=5) + with pytest.raises(ValueError, match="window"): + RecordGradients(murano_model, layer=1, window=0) + + +def test_record_gradients_shapes_and_finiteness(murano_model): + batch = _toy_rollouts(n=3, length=14, prompt_len=3) + pipeline = Pipeline( + [LoadRollouts(batch), RecordGradients(murano_model, layer=1, window=8)] + ) + store = pipeline.run()[keys.GRADIENT_RECORD] + assert isinstance(store, GradientStore) + assert len(store) == 3 + for grad in store.gradients: + # Window of 8 completion tokens, residual width of the tiny model. + assert grad.shape == (8, murano_model.d_model) + assert torch.isfinite(grad).all() + assert grad.abs().sum() > 0 + assert store.finite.all() + assert torch.isfinite(store.nll).all() + + +def test_record_gradients_window_clips_to_completion(murano_model): + batch = _toy_rollouts(n=1, length=10, prompt_len=6) + store = Pipeline( + [LoadRollouts(batch), RecordGradients(murano_model, layer=0, window=512)] + ).run()[keys.GRADIENT_RECORD] + # Only 4 completion tokens exist, so the window clips to them. + assert store.gradients[0].shape == (4, murano_model.d_model) + + +def test_record_gradients_is_deterministic(murano_model): + batch = _toy_rollouts() + step = RecordGradients(murano_model, layer=1, window=8) + first = Pipeline([LoadRollouts(batch), step]).run()[keys.GRADIENT_RECORD] + second = Pipeline([LoadRollouts(batch), step]).run()[keys.GRADIENT_RECORD] + for a, b in zip(first.gradients, second.gradients): + assert torch.equal(a, b) + + +def test_record_gradients_causality(murano_model): + """With one loss term, the gradient at the read position is exactly zero. + + The token at the window's only position is predicted from the logits one + position earlier, and the model is causal, so the residual stream *at* + that position cannot influence its own prediction. + """ + batch = _toy_rollouts(n=1, length=12, prompt_len=4) + store = Pipeline( + [LoadRollouts(batch), RecordGradients(murano_model, layer=1, window=1)] + ).run()[keys.GRADIENT_RECORD] + assert torch.allclose( + store.gradients[0], torch.zeros_like(store.gradients[0]), atol=1e-6 + ) + + +def test_record_gradients_nll_matches_nnsight_forward(murano_model): + """The native graph-preserving forward agrees with the nnsight path. + + The teacher-forced NLL computed inside the gradient pass must match the + NLL recomputed from ``forward_logits`` (a completely separate capture + path), pinning the prompt/completion alignment and the logits themselves. + """ + batch = _toy_rollouts(n=2, length=12, prompt_len=3) + store = Pipeline( + [LoadRollouts(batch), RecordGradients(murano_model, layer=1, window=6)] + ).run()[keys.GRADIENT_RECORD] + + for n, ids in enumerate(batch.input_ids): + tokens = { + "input_ids": ids.unsqueeze(0), + "attention_mask": torch.ones(1, ids.shape[0], dtype=torch.long), + } + logits = murano_model.forward_logits(tokens) + logprobs = torch.log_softmax(logits.float(), dim=-1) + start, end = 3, 3 + 6 + picked = logprobs[0, start - 1 : end - 1].gather(1, ids[start:end].unsqueeze(1)) + expected = -picked.mean().item() + assert store.nll[n].item() == pytest.approx(expected, abs=1e-4) + + +def test_record_gradients_on_gpt2(gpt2_model): + """The native forward path also works on the GPT-2 architecture.""" + batch = _toy_rollouts(n=1, length=10, prompt_len=2) + store = Pipeline( + [LoadRollouts(batch), RecordGradients(gpt2_model, layer=1, window=4)] + ).run()[keys.GRADIENT_RECORD] + assert store.gradients[0].shape == (4, gpt2_model.d_model) + assert torch.isfinite(store.gradients[0]).all() + assert store.gradients[0].abs().sum() > 0 + + +def test_gradient_store_save_load_roundtrip(murano_model, tmp_path): + from murano.io import load_gradient_store, save_gradient_store + + batch = _toy_rollouts(n=2) + store = Pipeline( + [LoadRollouts(batch), RecordGradients(murano_model, layer=1, window=5)] + ).run()[keys.GRADIENT_RECORD] + path = tmp_path / "gradient_record.pt" + save_gradient_store(store, path) + loaded = load_gradient_store(path) + assert loaded.layer == store.layer + assert loaded.window == store.window + assert torch.equal(loaded.nll, store.nll) + for a, b in zip(loaded.gradients, store.gradients): + assert torch.equal(a, b) + + +def test_save_results_dispatches_gradient_store(murano_model, tmp_path): + from murano.io import save_results + + batch = _toy_rollouts(n=1) + results = Pipeline( + [LoadRollouts(batch), RecordGradients(murano_model, layer=1, window=4)] + ).run() + out = save_results(results, output_dir=str(tmp_path), model_id="tiny") + assert (out / "gradients" / "gradient_record.pt").exists() From 457ed83f6bae1ace006bea4e67fc6ae442d68349 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Emirhan=20B=C3=B6ge?= Date: Fri, 31 Jul 2026 12:35:32 +0200 Subject: [PATCH 3/6] feat(steps): add the GFC routing operator and its comparisons --- src/murano/steps/gfc.py | 586 ++++++++++++++++++++++++++++++++++++++++ tests/test_gfc.py | 383 ++++++++++++++++++++++++++ 2 files changed, 969 insertions(+) create mode 100644 src/murano/steps/gfc.py create mode 100644 tests/test_gfc.py diff --git a/src/murano/steps/gfc.py b/src/murano/steps/gfc.py new file mode 100644 index 0000000..be2716f --- /dev/null +++ b/src/murano/steps/gfc.py @@ -0,0 +1,586 @@ +"""Gradient Feature Circuits (GFC): where the gradient travels between layers. + +The method asks not what a checkpoint computes but how it *routes credit*: +which directions of the residual stream at a late source layer connect, +through the frozen network's own backward transport, to which coordinates at +an earlier target layer. Comparing that routing map across checkpoints (early +pretraining, base, RL steps, and finetuned or perturbed variants) separates +"training rewired the paths" from "training re-weighted reads along fixed +paths". + +The construction, per rollout: + +1. Read the completion log-likelihood gradient ``g_t`` at the source layer + (the same quantity :class:`~murano.steps.gradients.RecordGradients` stores). +2. Project it onto ``k`` fixed random orthonormal directions ``d_i`` drawn + once from a seed: read strengths ``a_it = `` and seeds + ``s_i(t) = a_it * d_i``. The basis is fixed before any checkpoint is seen, + so any cross-checkpoint agreement belongs to the models, not the readout. +3. Transport every direction's seed field to the target layer with one + batched vector-Jacobian product against the frozen forward map, evaluated + at the same teacher-forced activations (chunkable via ``direction_chunk``). + Signs are kept through the cross-position sum; the absolute value is taken + per coordinate only afterwards. +4. Row ``i`` of the operator is the transported magnitude summed over target + positions: ``R_i = sum_tau |(J^T s_i)(tau)|``, giving ``R`` of shape + ``[k, d_model]``. Rollouts accumulate by summation. + +Two switches carry the method's controls. ``gradient_off=True`` sets every +read strength to 1, so the operator reads the frozen transport alone: any +surviving overlap is carried by the network's paths, not by the gradient read +along them. ``per_position`` keeps the target position open instead of +summing, giving one operator per completion position for opening-versus-rest +comparisons. + +Typical flow:: + + rollouts = RolloutBatch(input_ids=ids, prompt_lengths=lengths) + step = GFCOperator(model, source_layer=30, target_layer=25, k=128) + results = Pipeline([LoadRollouts(rollouts), step]).run() + base_map = results["gfc_operator"].operator # [k, d_model] + ... # rerun with another checkpoint's model + score = pairing_overlap(base_map, other_map) # 1 = same routing +""" + +from __future__ import annotations + +import contextlib +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +from torch import Generator, Tensor, full, isfinite, ones, outer # pyright: ignore[reportPrivateImportUsage] +from torch import float32, long, randperm, tensor, zeros, zeros_like # pyright: ignore[reportPrivateImportUsage] + +from murano import keys +from murano.logging import logger +from murano.results import Results +from murano.steps.base import Step +from murano.steps.gradients import ( + RolloutBatch, + _completion_loglik, + _completion_window, +) + +if TYPE_CHECKING: + from murano.backend import ModelBackend + + +def read_directions(d_model: int, k: int, seed: int = 0) -> Tensor: + """Draw ``k`` fixed random orthonormal read directions in ``d_model`` dims. + + A ``[d_model, k]`` matrix of iid standard Gaussians is drawn from ``seed`` + and orthonormalized by QR; the directions are the first ``k`` columns of + the orthonormal factor. The same seed gives the same basis on every call, + which is the point: the basis is fixed once, before any checkpoint is + read. + + Args: + d_model: Ambient dimension (the residual-stream width). + k: Number of directions; must not exceed ``d_model``. + seed: Seed for the Gaussian draw. + + Returns: + Float32 tensor ``[k, d_model]`` with orthonormal rows. + + Raises: + ValueError: If ``k`` exceeds ``d_model`` or either is not positive. + """ + if not 0 < k <= d_model: + raise ValueError( + f"k must be in [1, d_model={d_model}], got {k}: more orthonormal " + f"directions than dimensions do not exist." + ) + generator = Generator().manual_seed(seed) + gaussian = torch.randn(d_model, k, generator=generator, dtype=float32) + q, _ = torch.linalg.qr(gaussian) + return q[:, :k].T.contiguous() + + +def marginal(operator: Tensor) -> Tensor: + """The rank-one background fixed by an operator's row and column totals. + + ``marg(R) = rowsum(R) * colsum(R)^T / total(R)`` is what the totals would + produce if sources and targets were unlinked: it encodes how loud each + source direction and each target coordinate is, but nothing about which + source wires to which target. Two operators with identical marginals can + route completely differently, so every routing comparison subtracts it. + + Args: + operator: ``[k, d_model]`` routing operator. + + Returns: + Rank-one tensor of the same shape. Zero when the operator sums to 0. + """ + total = operator.sum() + if total.abs().item() == 0.0: + return zeros_like(operator) + return outer(operator.sum(dim=1), operator.sum(dim=0)) / total + + +def background_subtract(operator: Tensor) -> Tensor: + """Remove the rank-one marginal background: ``R - marg(R)``. + + The result has (numerically) zero row sums and zero column sums; what + remains is the pairing structure, who wires to whom beyond loudness. + """ + return operator - marginal(operator) + + +def pairing_operator(operator: Tensor, rank: int = 8) -> Tensor: + """Keep the top singular pairs at unit strength: ``sum_i u_i v_i^T``. + + The SVD of a (background-subtracted) operator splits it into paired source + and target directions with strengths. Replacing the top ``rank`` singular + values by 1 keeps *which* sources pair with *which* targets and discards + how hard each route is driven, so the comparison reads routing, not gain. + The result always has Frobenius norm ``sqrt(rank)``. + + Args: + operator: ``[k, d_model]`` operator, normally background-subtracted. + rank: Number of singular pairs to keep. + + Returns: + ``[k, d_model]`` unit-strength pairing operator. + + Raises: + ValueError: If ``rank`` is not in ``[1, min(operator.shape)]``. + """ + if not 0 < rank <= min(operator.shape): + raise ValueError( + f"rank must be in [1, {min(operator.shape)}] for an operator of " + f"shape {tuple(operator.shape)}, got {rank}." + ) + u, _, vh = torch.linalg.svd(operator.float(), full_matrices=False) + return u[:, :rank] @ vh[:rank, :] + + +def pairing_overlap( + a: Tensor, b: Tensor, rank: int = 8, subtract: bool = True +) -> float: + """Pairing-only overlap between two routing operators. + + Background-subtracts both operators (unless ``subtract=False``), reduces + each to its unit-strength top-``rank`` pairing, and returns their Frobenius + cosine. 1 means the two checkpoints pair the same source directions with + the same target coordinates (identical routing, up to consistent + relabeling of the pairs); values near the :func:`permutation_floor` mean + chance agreement. + + Args: + a: ``[k, d_model]`` routing operator. + b: ``[k, d_model]`` routing operator on the same basis and layer pair. + rank: Singular pairs kept per operator. + subtract: Background-subtract before pairing. Raw operators share most + of their energy through the marginal alone, so leaving this on is + what makes the floor sit near zero. + + Returns: + The overlap as a float. + + Raises: + ValueError: If the operators' shapes differ. + """ + if a.shape != b.shape: + raise ValueError( + f"Operators must share a shape, got {tuple(a.shape)} vs " + f"{tuple(b.shape)}; they are only comparable on the same basis " + f"and layer pair." + ) + if subtract: + a = background_subtract(a) + b = background_subtract(b) + pa = pairing_operator(a, rank=rank) + pb = pairing_operator(b, rank=rank) + denom = pa.norm() * pb.norm() + return ((pa * pb).sum() / denom).item() + + +def permutation_floor( + a: Tensor, + b: Tensor, + rank: int = 8, + subtract: bool = True, + draws: int = 20, + seed: int = 0, +) -> tuple[float, float]: + """Chance level of :func:`pairing_overlap`, by destroying the pairing. + + Randomly permutes the target columns of ``b`` (which preserves both + operators' row and column totals but breaks who-wires-to-whom) and + recomputes the overlap. On background-subtracted operators the floor sits + near zero; on raw operators it is high, which is why subtraction is the + default everywhere. + + Args: + a: ``[k, d_model]`` routing operator. + b: ``[k, d_model]`` routing operator to permute. + rank: Singular pairs kept per operator. + subtract: Background-subtract before pairing. + draws: Number of random permutations. + seed: Seed for the permutations. + + Returns: + ``(mean, sd)`` of the overlap across draws. + """ + generator = Generator().manual_seed(seed) + scores = [] + for _ in range(draws): + perm = randperm(b.shape[1], generator=generator) + scores.append(pairing_overlap(a, b[:, perm], rank=rank, subtract=subtract)) + stacked = tensor(scores) + return stacked.mean().item(), stacked.std().item() + + +def per_position_overlap( + a: Tensor, b: Tensor, rank: int = 8, subtract: bool = True +) -> Tensor: + """Position-resolved routing overlap ``o(tau)`` between two runs. + + Applies :func:`pairing_overlap` independently at every completion position + of two per-position operator stacks (the ``per_position`` output of + :class:`GFCOperator`). A dip over the first few positions with + recovery afterwards localizes a routing difference to how rollouts open. + + Args: + a: ``[positions, k, d_model]`` per-position operator stack. + b: Stack of the same shape from another checkpoint. + rank: Singular pairs kept per operator. + subtract: Background-subtract each position's operator. + + Returns: + Float tensor ``[positions]`` of overlaps. + + Raises: + ValueError: If the stacks' shapes differ. + """ + if a.shape != b.shape: + raise ValueError( + f"Per-position stacks must share a shape, got {tuple(a.shape)} vs " + f"{tuple(b.shape)}." + ) + return tensor( + [ + pairing_overlap(a[p], b[p], rank=rank, subtract=subtract) + for p in range(a.shape[0]) + ] + ) + + +@dataclass +class GFCOperatorResult: + """A routing operator and its capture settings. + + Attributes: + operator: ``[k, d_model]`` float32 operator, summed over target + positions and rollouts. + source_layer: Layer the gradient is read at. + target_layer: Layer the seeds are transported to. + k: Number of read directions. + per_position: ``[positions, k, d_model]`` stack with the target + position kept open, when requested; ``None`` otherwise. Position 0 + is the first completion token. + position_counts: ``[positions]`` rollouts contributing at each + position (rollouts shorter than the horizon stop contributing); + ``None`` when ``per_position`` was not requested. + nll: ``[n_rollouts]`` mean teacher-forced negative log-likelihood per + completion token on the same window, the functional-change scale + for matching runs; NaN for dropped rollouts. + kept: ``[n_rollouts]`` bool; False where a rollout was dropped for + non-finite gradients or transport. + basis_seed: Seed of the direction draw, or ``None`` when an explicit + ``directions`` basis was passed instead of a seeded draw. + gradient_off: True when every read strength was held at 1 (the + gradient-off control); False for the gradient-weighted operator. + window: Maximum completion tokens read per rollout. + """ + + operator: Tensor + source_layer: int + target_layer: int + k: int + per_position: Tensor | None = None + position_counts: Tensor | None = None + nll: Tensor | None = None + kept: Tensor | None = None + basis_seed: int | None = 0 + gradient_off: bool = False + window: int = 512 + + +class GFCOperator(Step): + """Build the GFC routing operator for one checkpoint over a rollout corpus. + + For each rollout: one teacher-forced forward preserving the graph, one + backward for the source-layer gradient (skipped in gradient-off mode), + then one batched vector-Jacobian product transporting every read + direction's seed field from the source layer to the target layer at the + frozen activations (chunkable via ``direction_chunk``). Transported + magnitudes accumulate into ``[k, d_model]`` rows; rollouts accumulate by + summation. + + Reads from results: + results['rollouts']: RolloutBatch + + Writes to results: + results['gfc_operator']: GFCOperatorResult + + Args: + model: Wrapped model to read. + source_layer: Residual-stream layer (block output) the gradient is + read and seeds are placed at. + target_layer: Earlier layer the seeds are transported to; must be + strictly below ``source_layer``. + k: Number of fixed random orthonormal read directions; at most + ``d_model``. + basis_seed: Seed of the direction draw. Use the same seed for every + checkpoint being compared; redraw with several seeds to check + basis robustness. Ignored when ``directions`` is given. + window: Maximum completion tokens read per rollout (loss terms, seed + positions, and target sum alike). + gradient_off: If True, hold every read strength at 1 instead of + reading it from the gradient, so the operator reads the frozen + transport alone (the paper's gradient-off control). + per_position: If given, additionally keep the target position open for + the first ``per_position`` completion positions, producing the + per-position operator stack for opening-versus-rest analyses. + directions: A pre-drawn ``[k, d_model]`` basis with unit-norm rows, + for reproducing an analysis whose read directions were drawn + elsewhere. When omitted, the basis is drawn from ``basis_seed`` + via :func:`read_directions`. + direction_chunk: Directions transported per vector-Jacobian product. + ``None`` (the default) batches all ``k`` in one VJP via + ``is_grads_batched``, the fast path; an integer transports that + many at a time, trading speed for peak memory; ``1`` uses plain + unbatched VJPs, the escape hatch for an architecture whose + backward does not compose with ``torch.vmap``. Every setting + produces the identical operator. + + Raises: + ValueError: If the layer pair, ``k``, ``window``, ``per_position``, + ``directions``, or ``direction_chunk`` is invalid. + + Note: + Cost scales with ``k * n_rollouts`` VJP work; batching the directions + (the default) amortizes the backward traversal so the wall clock is + far below ``k`` sequential backwards. Peak memory adds one + ``[chunk, seq, d_model]`` seed-field stack in the model dtype on top + of one sequence's activations. Rollouts whose gradients are + non-finite (near-random checkpoints can produce them) are dropped and + counted in ``kept`` (non-finite transport likewise), matching + :class:`~murano.steps.gradients.RecordGradients`. + """ + + reads = [keys.ROLLOUTS] + writes = [keys.GFC_OPERATOR] + read_types = {keys.ROLLOUTS: RolloutBatch} + write_types = {keys.GFC_OPERATOR: GFCOperatorResult} + + def __init__( + self, + model: ModelBackend, + source_layer: int, + target_layer: int, + k: int = 128, + basis_seed: int = 0, + window: int = 512, + gradient_off: bool = False, + per_position: int | None = None, + directions: Tensor | None = None, + direction_chunk: int | None = None, + ): + if not 0 <= target_layer < source_layer < model.n_layers: + raise ValueError( + f"Need 0 <= target_layer < source_layer < n_layers=" + f"{model.n_layers}; got source {source_layer}, target " + f"{target_layer}. The operator transports backwards, from a " + f"later layer to an earlier one." + ) + if not 0 < k <= model.d_model: + raise ValueError(f"k must be in [1, d_model={model.d_model}], got {k}.") + if window < 1: + raise ValueError(f"window must be >= 1, got {window}.") + if per_position is not None and per_position < 1: + raise ValueError( + f"per_position must be a positive horizon or None, got {per_position}." + ) + if directions is not None: + if directions.shape != (k, model.d_model): + raise ValueError( + f"directions must be [k={k}, d_model={model.d_model}], got " + f"{tuple(directions.shape)}." + ) + norms = directions.float().norm(dim=1) + if not bool((norms - 1.0).abs().lt(1e-3).all()): + raise ValueError( + "directions must have unit-norm rows; draw them with " + "read_directions or normalize before passing." + ) + if direction_chunk is not None and direction_chunk < 1: + raise ValueError( + f"direction_chunk must be a positive count or None, got " + f"{direction_chunk}." + ) + self.model = model + self.source_layer = source_layer + self.target_layer = target_layer + self.k = k + self.basis_seed = basis_seed + self.window = window + self.gradient_off = gradient_off + self.per_position = per_position + self.directions = directions + self.direction_chunk = direction_chunk + + def __call__(self, results: Results) -> Results: + rollouts: RolloutBatch = results[keys.ROLLOUTS] + model = self.model + last = model.n_layers - 1 + layers = sorted({self.target_layer, self.source_layer, last}) + directions = ( + self.directions.float().cpu() + if self.directions is not None + else read_directions(model.d_model, self.k, self.basis_seed) + ) + + operator = zeros(self.k, model.d_model) + horizon = self.per_position + per_position = ( + zeros(horizon, self.k, model.d_model) if horizon is not None else None + ) + position_counts = zeros(horizon, dtype=long) if horizon is not None else None + nll = full((len(rollouts),), float("nan")) + kept = full((len(rollouts),), False) + + logger.info( + "GFC operator: %d rollouts, layers %d->%d, k=%d, gradient_off=%s", + len(rollouts), + self.source_layer, + self.target_layer, + self.k, + self.gradient_off, + ) + for n, ids in enumerate(rollouts.input_ids): + start, end = _completion_window( + ids, rollouts.prompt_lengths[n], self.window + ) + hidden = model.grad_forward(ids.unsqueeze(0), layers) + source = hidden[self.source_layer] + target = hidden[self.target_layer] + + strengths = self._read_strengths( + model, hidden[last], source, ids, start, end, directions, nll, n + ) + if strengths is None: + continue + + contribution = zeros_like(operator) + pos_contribution = ( + zeros(horizon, self.k, model.d_model) if horizon is not None else None + ) + chunk = self.direction_chunk or self.k + for begin in range(0, self.k, chunk): + stop = min(begin + chunk, self.k) + last_chunk = stop == self.k + # Seed field per direction in the chunk: a_it * d_i at every + # window position, zero elsewhere. [c, T, d_model] built in one + # broadcast, then placed into the full-sequence field stack. + seeds = strengths[:, begin:stop].T.unsqueeze(-1) * directions[ + begin:stop + ].unsqueeze(1) + fields = zeros( + stop - begin, + *source.shape, + dtype=source.dtype, + device=source.device, + ) + fields[:, 0, start:end] = seeds.to(source.device, source.dtype) + if fields.shape[0] == 1: + # Plain unbatched VJP: the escape hatch for architectures + # whose backward does not compose with torch.vmap. + (transported,) = torch.autograd.grad( + source, + target, + grad_outputs=fields[0], + retain_graph=not last_chunk, + ) + transported = transported.unsqueeze(0) + else: + (transported,) = torch.autograd.grad( + source, + target, + grad_outputs=fields, + is_grads_batched=True, + retain_graph=not last_chunk, + ) + magnitudes = transported[:, 0, start:end].detach().float().cpu().abs() + contribution[begin:stop] = magnitudes.sum(dim=1) + if pos_contribution is not None: + p = min(magnitudes.shape[1], pos_contribution.shape[0]) + pos_contribution[:p, begin:stop] = magnitudes[:, :p].transpose(0, 1) + + if not bool(isfinite(contribution).all()): + logger.warning("Rollout %d produced non-finite transport; dropped.", n) + nll[n] = float("nan") + continue + operator += contribution + if horizon is not None: + assert per_position is not None + assert pos_contribution is not None + assert position_counts is not None + per_position += pos_contribution + position_counts[: min(end - start, per_position.shape[0])] += 1 + kept[n] = True + + results[keys.GFC_OPERATOR] = GFCOperatorResult( + operator=operator, + per_position=per_position, + position_counts=position_counts, + nll=nll, + kept=kept, + source_layer=self.source_layer, + target_layer=self.target_layer, + k=self.k, + basis_seed=None if self.directions is not None else self.basis_seed, + gradient_off=self.gradient_off, + window=self.window, + ) + logger.info( + "GFC operator done: kept %d/%d rollouts.", + int(kept.sum()), + len(rollouts), + ) + return results + + def _read_strengths( + self, + model: ModelBackend, + hidden_last: Tensor, + source: Tensor, + ids: Tensor, + start: int, + end: int, + directions: Tensor, + nll: Tensor, + n: int, + ) -> Tensor | None: + """Per-position read strengths ``[T, k]``, or None to drop the rollout. + + Normally the strengths are the gradient's components along each + direction and the graph is retained for the transport VJPs; with + ``gradient_off`` every strength is 1 and only the (grad-free) NLL is + computed. Fills ``nll[n]`` as a side effect. + """ + context = torch.no_grad() if self.gradient_off else contextlib.nullcontext() + with context: + loglik = _completion_loglik(model, hidden_last, ids, start, end) + nll[n] = -loglik.item() / (end - start) + if self.gradient_off: + return ones(end - start, self.k) + (grad,) = torch.autograd.grad(loglik, source, retain_graph=True) + grad_window = grad[0, start:end].detach().float().cpu() + if not bool(isfinite(grad_window).all()): + logger.warning("Rollout %d produced non-finite gradients; dropped.", n) + nll[n] = float("nan") + return None + return grad_window @ directions.T diff --git a/tests/test_gfc.py b/tests/test_gfc.py new file mode 100644 index 0000000..e1debbb --- /dev/null +++ b/tests/test_gfc.py @@ -0,0 +1,383 @@ +"""Tests for the GFC routing operator and its comparison functions.""" + +from __future__ import annotations + +import pytest +import torch + +from murano import keys +from murano.pipeline import Pipeline +from murano.steps.gradients import ( + LoadRollouts, + RolloutBatch, + _completion_window, +) +from murano.steps.gfc import ( + GFCOperator, + GFCOperatorResult, + background_subtract, + marginal, + read_directions, + pairing_operator, + per_position_overlap, + permutation_floor, + pairing_overlap, +) +from tests.conftest import toy_rollouts + + +def _toy_rollouts(n: int = 2, length: int = 14, prompt_len: int = 3) -> RolloutBatch: + return toy_rollouts(n, length, prompt_len, seed=11) + + +def _random_operator(seed: int, k: int = 16, d: int = 32) -> torch.Tensor: + generator = torch.Generator().manual_seed(seed) + return torch.rand(k, d, generator=generator) + + +# ---------------------------------------------------------------- pure helpers + + +def test_read_directions_are_orthonormal(): + directions = read_directions(32, 16, seed=0) + assert directions.shape == (16, 32) + gram = directions @ directions.T + assert torch.allclose(gram, torch.eye(16), atol=1e-5) + + +def test_read_directions_seeded(): + same = read_directions(32, 8, seed=3) + again = read_directions(32, 8, seed=3) + other = read_directions(32, 8, seed=4) + assert torch.equal(same, again) + assert not torch.allclose(same, other) + + +def test_read_directions_rejects_k_above_d(): + with pytest.raises(ValueError, match="d_model"): + read_directions(16, 32) + + +def test_background_subtract_zeroes_margins(): + operator = _random_operator(0) + subtracted = background_subtract(operator) + assert torch.allclose(subtracted.sum(dim=0), torch.zeros(32), atol=1e-5) + assert torch.allclose(subtracted.sum(dim=1), torch.zeros(16), atol=1e-5) + # marg(R) itself is rank one. + assert torch.linalg.matrix_rank(marginal(operator)).item() == 1 + + +def test_pairing_operator_has_unit_strength_norm(): + operator = background_subtract(_random_operator(1)) + for rank in (1, 4, 8): + paired = pairing_operator(operator, rank=rank) + assert paired.norm().item() == pytest.approx(rank**0.5, abs=1e-4) + + +def test_pairing_operator_rejects_bad_rank(): + with pytest.raises(ValueError, match="rank"): + pairing_operator(_random_operator(2), rank=17) + + +def test_self_overlap_is_one(): + operator = _random_operator(3) + assert pairing_overlap(operator, operator) == pytest.approx(1.0, abs=1e-5) + + +def test_overlap_rejects_shape_mismatch(): + with pytest.raises(ValueError, match="shape"): + pairing_overlap(_random_operator(0), _random_operator(0, k=8)) + + +def test_permutation_floor_near_zero_after_subtraction(): + """Column-shuffled operators agree only by chance once the background + (which the shuffle preserves) is removed.""" + a = _random_operator(4) + b = _random_operator(5) + mean, sd = permutation_floor(a, b, draws=10) + assert abs(mean) < 0.15 + assert sd < 0.15 + # The floor sits far below a genuine self-match. + assert pairing_overlap(a, a) - mean > 0.8 + + +def test_independent_operators_overlap_near_floor(): + a = _random_operator(6) + b = _random_operator(7) + assert abs(pairing_overlap(a, b)) < 0.3 + + +def test_per_position_overlap_shape(): + a = torch.stack([_random_operator(i) for i in range(4)]) + curve = per_position_overlap(a, a) + assert curve.shape == (4,) + assert torch.allclose(curve, torch.ones(4), atol=1e-5) + with pytest.raises(ValueError, match="shape"): + per_position_overlap(a, a[:2]) + + +# ------------------------------------------------------------------- the step + + +def test_gfc_operator_validates_arguments(murano_model): + with pytest.raises(ValueError, match="target_layer < source_layer"): + GFCOperator(murano_model, source_layer=0, target_layer=1) + with pytest.raises(ValueError, match="k must be"): + GFCOperator(murano_model, source_layer=1, target_layer=0, k=64) + with pytest.raises(ValueError, match="per_position"): + GFCOperator(murano_model, source_layer=1, target_layer=0, k=4, per_position=0) + with pytest.raises(ValueError, match="window"): + GFCOperator(murano_model, source_layer=1, target_layer=0, k=4, window=0) + + +def test_gfc_operator_shape_and_positivity(murano_model): + batch = _toy_rollouts(n=2) + result = Pipeline( + [ + LoadRollouts(batch), + GFCOperator(murano_model, source_layer=1, target_layer=0, k=8, window=6), + ] + ).run()[keys.GFC_OPERATOR] + assert isinstance(result, GFCOperatorResult) + # The operator is [k, d_model] with nonnegative entries (summed magnitudes). + assert result.operator.shape == (8, murano_model.d_model) + assert (result.operator >= 0).all() + assert result.operator.sum() > 0 + assert result.kept is not None and result.kept.all() + assert result.nll is not None and torch.isfinite(result.nll).all() + + +def test_per_position_stack_sums_to_pooled_operator(murano_model): + """Opening the target position is a refinement, not a different object: + summing the per-position stack over positions recovers the pooled operator + exactly (same VJPs, same magnitudes).""" + batch = _toy_rollouts(n=2, length=12, prompt_len=3) + result = Pipeline( + [ + LoadRollouts(batch), + GFCOperator( + murano_model, + source_layer=1, + target_layer=0, + k=6, + window=6, + per_position=6, + ), + ] + ).run()[keys.GFC_OPERATOR] + assert result.per_position is not None + assert result.per_position.shape == (6, 6, murano_model.d_model) + assert torch.allclose(result.per_position.sum(dim=0), result.operator, atol=1e-5) + assert result.position_counts is not None + assert result.position_counts.tolist() == [2] * 6 + + +def test_gradient_off_matches_manual_bare_direction_vjp(murano_model): + """Gradient-off means the loss never enters: the operator must equal the + hand-computed transport of the bare directions at every window position.""" + batch = _toy_rollouts(n=1, length=10, prompt_len=2) + k, window = 4, 5 + result = Pipeline( + [ + LoadRollouts(batch), + GFCOperator( + murano_model, + source_layer=1, + target_layer=0, + k=k, + window=window, + gradient_off=True, + ), + ] + ).run()[keys.GFC_OPERATOR] + + ids = batch.input_ids[0] + start, end = _completion_window(ids, 2, window) + hidden = murano_model.grad_forward(ids.unsqueeze(0), [0, 1]) + directions = read_directions(murano_model.d_model, k) + manual = torch.zeros(k, murano_model.d_model) + for i in range(k): + field = torch.zeros_like(hidden[1]) + field[0, start:end] = directions[i] + (transported,) = torch.autograd.grad( + hidden[1], hidden[0], grad_outputs=field, retain_graph=True + ) + manual[i] = transported[0, start:end].detach().float().abs().sum(dim=0) + assert torch.allclose(result.operator, manual, atol=1e-4) + + +def test_gradient_and_gradient_off_operators_differ(murano_model): + batch = _toy_rollouts(n=1) + results = {} + for gradient_off in (False, True): + results[gradient_off] = Pipeline( + [ + LoadRollouts(batch), + GFCOperator( + murano_model, + source_layer=1, + target_layer=0, + k=6, + window=6, + gradient_off=gradient_off, + ), + ] + ).run()[keys.GFC_OPERATOR] + assert not torch.allclose(results[False].operator, results[True].operator) + + +def test_direction_chunks_produce_identical_operators(murano_model): + """Batched, chunked, and plain per-direction VJPs are the same math. + + ``direction_chunk=None`` runs one vmap-batched VJP, ``3`` chunks it, and + ``1`` falls back to plain unbatched VJPs; the operator must be identical + (up to float accumulation) in every case. + """ + batch = _toy_rollouts(n=2, length=12, prompt_len=3) + operators = [] + for chunk in (None, 3, 1): + result = Pipeline( + [ + LoadRollouts(batch), + GFCOperator( + murano_model, + source_layer=1, + target_layer=0, + k=8, + window=6, + per_position=6, + direction_chunk=chunk, + ), + ] + ).run()[keys.GFC_OPERATOR] + operators.append(result) + for other in operators[1:]: + assert torch.allclose(operators[0].operator, other.operator, atol=1e-5) + assert torch.allclose( + operators[0].per_position, other.per_position, atol=1e-5 + ) + + +def test_explicit_directions_match_seeded_draw(murano_model): + """Passing the basis explicitly reproduces the seeded internal draw.""" + batch = _toy_rollouts(n=1) + basis = read_directions(murano_model.d_model, 6, seed=9) + explicit = Pipeline( + [ + LoadRollouts(batch), + GFCOperator( + murano_model, + source_layer=1, + target_layer=0, + k=6, + window=6, + directions=basis, + ), + ] + ).run()[keys.GFC_OPERATOR] + seeded = Pipeline( + [ + LoadRollouts(batch), + GFCOperator( + murano_model, + source_layer=1, + target_layer=0, + k=6, + window=6, + basis_seed=9, + ), + ] + ).run()[keys.GFC_OPERATOR] + assert torch.allclose(explicit.operator, seeded.operator, atol=1e-6) + + +def test_explicit_directions_validation(murano_model): + with pytest.raises(ValueError, match="directions must be"): + GFCOperator( + murano_model, + source_layer=1, + target_layer=0, + k=4, + directions=torch.eye(3), + ) + with pytest.raises(ValueError, match="unit-norm"): + GFCOperator( + murano_model, + source_layer=1, + target_layer=0, + k=4, + directions=torch.ones(4, murano_model.d_model), + ) + with pytest.raises(ValueError, match="direction_chunk"): + GFCOperator( + murano_model, + source_layer=1, + target_layer=0, + k=4, + direction_chunk=0, + ) + + +def test_gfc_operator_deterministic_and_self_consistent(murano_model): + batch = _toy_rollouts(n=2) + step = GFCOperator(murano_model, source_layer=1, target_layer=0, k=8, window=6) + first = Pipeline([LoadRollouts(batch), step]).run()[keys.GFC_OPERATOR] + second = Pipeline([LoadRollouts(batch), step]).run()[keys.GFC_OPERATOR] + assert torch.equal(first.operator, second.operator) + + +def test_pipeline_validates_wiring(murano_model): + batch = _toy_rollouts(n=1) + pipeline = Pipeline( + [ + LoadRollouts(batch), + GFCOperator(murano_model, source_layer=1, target_layer=0, k=4), + ] + ) + pipeline.validate() + # Without the loader the chain is broken and validation says so. + with pytest.raises(KeyError, match="rollouts"): + Pipeline( + [GFCOperator(murano_model, source_layer=1, target_layer=0, k=4)] + ).validate() + + +def test_gfc_operator_save_load_roundtrip(murano_model, tmp_path): + from murano.io import load_gfc_operator, save_gfc_operator + + batch = _toy_rollouts(n=1) + result = Pipeline( + [ + LoadRollouts(batch), + GFCOperator( + murano_model, + source_layer=1, + target_layer=0, + k=4, + window=5, + per_position=5, + ), + ] + ).run()[keys.GFC_OPERATOR] + path = tmp_path / "gfc_operator.pt" + save_gfc_operator(result, path) + loaded = load_gfc_operator(path) + assert torch.equal(loaded.operator, result.operator) + assert loaded.per_position is not None + assert torch.equal(loaded.per_position, result.per_position) + assert loaded.gradient_off == result.gradient_off + assert (loaded.source_layer, loaded.target_layer) == (1, 0) + + +def test_save_results_dispatches_gfc_operator(murano_model, tmp_path): + from murano.io import save_results + + batch = _toy_rollouts(n=1) + results = Pipeline( + [ + LoadRollouts(batch), + GFCOperator(murano_model, source_layer=1, target_layer=0, k=4), + ] + ).run() + out = save_results(results, output_dir=str(tmp_path), model_id="tiny") + assert (out / "gfc" / "gfc_operator.pt").exists() From a3baded27eed2d945a73732787ea412c37af2e18 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Emirhan=20B=C3=B6ge?= Date: Fri, 31 Jul 2026 12:36:27 +0200 Subject: [PATCH 4/6] feat(steps): add GSAE, a TopK autoencoder over recorded gradients --- src/murano/steps/gsae.py | 407 +++++++++++++++++++++++++++++++++++++++ tests/test_gsae.py | 225 ++++++++++++++++++++++ 2 files changed, 632 insertions(+) create mode 100644 src/murano/steps/gsae.py create mode 100644 tests/test_gsae.py diff --git a/src/murano/steps/gsae.py b/src/murano/steps/gsae.py new file mode 100644 index 0000000..33ec87d --- /dev/null +++ b/src/murano/steps/gsae.py @@ -0,0 +1,407 @@ +"""Gradient Sparse Autoencoders (GSAE): a TopK dictionary over recorded gradients. + +The forward-capture SAE path (:mod:`murano.steps.sae`) loads pre-trained +dictionaries; gradients have none, so the GSAE ships its own minimal trainer. +The object of study is the Gradient Feature Circuits feature census: +train one shared dictionary on the pooled gradient inputs of several +checkpoints (base and RL alike), then compare how often each feature fires per +checkpoint. A feature is *introduced* when it is nearly silent on the base +gradients but common after RL, and *turned up* (the paper also says +"recruited") when RL multiplies an existing firing rate. Because the +dictionary is shared and fit once, a firing-rate shift is a property of the +gradients, not of a per-checkpoint refit. + +The encoder input follows the paper: per token position, +``x = sign(A) * g / RMS(g)``, where ``g`` is the recorded residual-stream +gradient and ``A`` the rollout's standardized advantage, so the code sees the +advantage's sign but not its size, and every input has unit RMS regardless of +how loud its position was. + +The autoencoder is the TopK variant: ``z = TopK_k(W_enc (x - b_pre))`` keeps +the ``k`` largest coordinates and zeroes the rest, ``xhat = W_dec z + b_pre``, +decoder rows are renormalized to unit length after every step, and features +that never fired in an epoch are re-initialized to the residual of a +badly-reconstructed input so the dictionary stays alive. + +Typical flow:: + + inputs = normalized_gradient_inputs(store, advantages) # per checkpoint + gsae = GSAE.train(torch.cat(all_inputs), m=16384, k=64) # fit ONCE, pooled + base_rates = firing_rates(gsae, base_inputs) + rl_rates = firing_rates(gsae, rl_inputs).amax(dim=0) # max over RL ckpts + census = classify_features(base_rates, rl_rates) + census["introduced"] # feature ids nearly silent before RL, common after +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +from torch import ( + Generator, # pyright: ignore[reportPrivateImportUsage] + Tensor, + bool as torch_bool, # pyright: ignore[reportPrivateImportUsage] + cat, # pyright: ignore[reportPrivateImportUsage] + randn, + randperm, # pyright: ignore[reportPrivateImportUsage] + zeros, # pyright: ignore[reportPrivateImportUsage] +) + +from murano.logging import logger +from murano.steps.gradients import GradientStore + +if TYPE_CHECKING: + from murano.backend import ModelBackend + +CENSUS_CLASSES = ("introduced", "turned_up", "suppressed", "stable") + + +def normalized_gradient_inputs( + store: GradientStore, + advantages: Tensor | None = None, + subsample: float | None = None, + seed: int = 0, +) -> Tensor: + """Pool a gradient store into GSAE inputs ``sign(A) * g / RMS(g)``. + + Every kept rollout's per-position gradients are RMS-normalized (each row + ends at unit RMS, so loud and quiet positions weigh alike) and multiplied + by the sign of that rollout's advantage. Rollouts with advantage 0 (an + all-pass or all-fail group carries no within-group signal) are dropped, + matching the paper's capture. + + Args: + store: Recorded gradients, one ``[T, d_model]`` tensor per rollout. + advantages: Per-rollout advantages; ``None`` keeps every finite + rollout with sign +1. + subsample: Keep this fraction of pooled positions, drawn without + replacement (the paper keeps 5% at scale); ``None`` keeps all. + seed: Seed for the subsample draw. + + Returns: + Float32 tensor ``[n_positions, d_model]`` of encoder inputs. + + Raises: + ValueError: If ``advantages`` misaligns with the store or + ``subsample`` is not in (0, 1]. + """ + if advantages is not None and len(advantages) != len(store): + raise ValueError( + f"advantages must have one entry per rollout: got " + f"{len(advantages)} for {len(store)} rollouts." + ) + if subsample is not None and not 0.0 < subsample <= 1.0: + raise ValueError(f"subsample must be in (0, 1], got {subsample}.") + rows = [] + for n, grad in enumerate(store.gradients): + if not bool(store.finite[n]) or grad.numel() == 0: + continue + sign = 1.0 + if advantages is not None: + a = float(advantages[n]) + if a == 0.0: + continue + sign = 1.0 if a > 0 else -1.0 + rms = grad.float().pow(2).mean(dim=1, keepdim=True).sqrt().clamp_min(1e-8) + rows.append(sign * grad.float() / rms) + if not rows: + raise ValueError("No usable rollouts: all dropped or zero-advantage.") + pooled = cat(rows) + if subsample is not None and subsample < 1.0: + generator = Generator().manual_seed(seed) + n_keep = max(1, int(pooled.shape[0] * subsample)) + keep = randperm(pooled.shape[0], generator=generator)[:n_keep] + pooled = pooled[keep] + return pooled + + +@dataclass +class GSAE: + """A trained TopK gradient sparse autoencoder. + + Attributes: + w_enc: Encoder weight ``[m, d_model]``. + w_dec: Decoder weight ``[m, d_model]``, unit-norm rows; row ``j`` is + feature ``j``'s direction in the residual stream. + b_pre: Pre-encoder bias ``[d_model]``, subtracted before encoding and + added back after decoding. + k: Active features per input. + fvu: Fraction of variance unexplained on the training inputs after the + final epoch (0 = perfect reconstruction). + """ + + w_enc: Tensor + w_dec: Tensor + b_pre: Tensor + k: int + fvu: float = float("nan") + + @property + def m(self) -> int: + """Dictionary size (number of features).""" + return int(self.w_enc.shape[0]) + + def encode(self, inputs: Tensor) -> Tensor: + """Sparse codes ``[n, m]``: TopK of the pre-activations, rest zeroed. + + Args: + inputs: ``[n, d_model]`` encoder inputs. + + Returns: + ``[n, m]`` codes with exactly ``min(k, m)`` nonzeros per row. + """ + pre = (inputs.float() - self.b_pre) @ self.w_enc.T + values, indices = pre.topk(min(self.k, self.m), dim=-1) + codes = zeros(pre.shape, dtype=pre.dtype, device=pre.device) + return codes.scatter_(-1, indices, values) + + def decode(self, codes: Tensor) -> Tensor: + """Reconstruction ``[n, d_model]`` from sparse codes.""" + return codes @ self.w_dec + self.b_pre + + def feature_direction(self, feature: int) -> Tensor: + """Feature ``feature``'s unit decoder direction ``[d_model]``.""" + return self.w_dec[feature] + + @classmethod + def train( + cls, + inputs: Tensor, + m: int, + k: int, + lr: float = 5e-4, + batch_size: int = 4096, + epochs: int = 3, + seed: int = 0, + ) -> "GSAE": + """Fit a TopK autoencoder on pooled gradient inputs. + + Plain reconstruction training: Adam on the squared error, decoder rows + renormalized to unit length after every step (so a feature's strength + lives in its code, not its direction), and features that never fired + in an epoch re-initialized between epochs to the residual of the + worst-reconstructed input of a probe batch, so the dictionary stays + alive instead of accumulating dead features. + + Args: + inputs: ``[n, d_model]`` pooled encoder inputs (pool the + checkpoints first: the census needs one shared dictionary). + m: Dictionary size. + k: Active features per input; at most ``m``. + lr: Adam learning rate. + batch_size: Minibatch size. + epochs: Passes over the inputs. + seed: Seed for initialization and batch order. + + Returns: + The trained :class:`GSAE`, with ``fvu`` filled in. + + Raises: + ValueError: If ``inputs`` is not 2-D or ``k``/``m`` are invalid. + """ + if inputs.dim() != 2: + raise ValueError( + f"inputs must be [n, d_model], got shape {tuple(inputs.shape)}." + ) + if not 0 < k <= m: + raise ValueError(f"Need 0 < k <= m, got k={k}, m={m}.") + inputs = inputs.float() + n, d_model = inputs.shape + generator = Generator().manual_seed(seed) + w_dec = randn(m, d_model, generator=generator) + w_dec = w_dec / w_dec.norm(dim=1, keepdim=True) + gsae = cls( + w_enc=w_dec.clone().requires_grad_(True), + w_dec=w_dec.clone().requires_grad_(True), + b_pre=zeros(d_model).requires_grad_(True), + k=k, + ) + optimizer = torch.optim.Adam([gsae.w_enc, gsae.w_dec, gsae.b_pre], lr=lr) + variance = inputs.var(dim=0, unbiased=False).sum().clamp_min(1e-12) + + for epoch in range(epochs): + order = randperm(n, generator=generator) + fired = zeros(m, dtype=torch_bool) + for begin in range(0, n, batch_size): + batch = inputs[order[begin : begin + batch_size]] + codes = gsae.encode(batch) + reconstruction = gsae.decode(codes) + loss = (reconstruction - batch).pow(2).sum(dim=1).mean() + optimizer.zero_grad() + loss.backward() + optimizer.step() + with torch.no_grad(): + gsae.w_dec /= gsae.w_dec.norm(dim=1, keepdim=True).clamp_min(1e-8) + fired |= codes.ne(0).any(dim=0) + logger.info( + "GSAE epoch %d/%d: %d/%d features alive", + epoch + 1, + epochs, + int(fired.sum()), + m, + ) + if epoch < epochs - 1: + _reinit_dead_features(gsae, inputs, fired, generator) + + gsae.w_enc = gsae.w_enc.detach() + gsae.w_dec = gsae.w_dec.detach() + gsae.b_pre = gsae.b_pre.detach() + with torch.no_grad(): + error = 0.0 + for begin in range(0, n, batch_size): + batch = inputs[begin : begin + batch_size] + residual = batch - gsae.decode(gsae.encode(batch)) + error += float(residual.pow(2).sum()) + gsae.fvu = error / n / float(variance) + logger.info("GSAE trained: FVU %.4f on the training inputs", gsae.fvu) + return gsae + + +def _reinit_dead_features( + gsae: GSAE, inputs: Tensor, fired: Tensor, generator: Generator +) -> None: + """Point dead features at what the live ones reconstruct worst. + + For every feature that fired on nothing this epoch, pick a high-error + input from a probe batch and set the feature's decoder row to that input's + residual (renormalized), with the encoder row matched to it, so the next + epoch can recruit the feature for exactly the structure the dictionary is + missing. + + Each revived feature needs its own residual, so the probe grows with the + dead count. When more features are dead than there are inputs to draw from, + the remainder keep their current directions and get another chance next + epoch rather than the pass failing. + """ + dead = (~fired).nonzero(as_tuple=True)[0] + if dead.numel() == 0: + return + with torch.no_grad(): + n_probe = min(max(4096, dead.numel()), inputs.shape[0]) + probe = inputs[randperm(inputs.shape[0], generator=generator)[:n_probe]] + residual = probe - gsae.decode(gsae.encode(probe)) + errors = residual.pow(2).sum(dim=1) + revived = dead[: probe.shape[0]] + worst = errors.argsort(descending=True)[: revived.numel()] + directions = residual[worst] + directions = directions / directions.norm(dim=1, keepdim=True).clamp_min(1e-8) + gsae.w_dec[revived] = directions + gsae.w_enc[revived] = directions + logger.info( + "Re-initialized %d of %d dead features", + int(revived.numel()), + int(dead.numel()), + ) + + +def firing_rates(gsae: GSAE, inputs: Tensor, batch_size: int = 4096) -> Tensor: + """Fraction of inputs on which each feature is among the selected k. + + Args: + gsae: The trained autoencoder. + inputs: ``[n, d_model]`` encoder inputs of one checkpoint. + batch_size: Encode batch size. + + Returns: + Float tensor ``[m]`` of per-feature firing rates in [0, 1]. + """ + counts = zeros(gsae.m) + with torch.no_grad(): + for begin in range(0, inputs.shape[0], batch_size): + codes = gsae.encode(inputs[begin : begin + batch_size]) + counts += codes.ne(0).sum(dim=0).float() + return counts / inputs.shape[0] + + +def classify_features( + base_rates: Tensor, + rl_rates: Tensor, + introduced: tuple[float, float] = (0.05, 0.15), + turned_up: tuple[float, float] = (3.0, 0.005), +) -> dict[str, list[int]]: + """The paper's firing-rate census, first matching class per feature. + + With base rate ``f0`` and RL rate ``fR`` (at scale, the max over the RL + checkpoints): + + * introduced: ``f0 < introduced[0]`` and ``fR > introduced[1]``. The + feature barely existed in the base gradients and is common after RL. + * turned_up: else ``fR >= turned_up[0] * f0`` and + ``fR - f0 > turned_up[1]``. An existing feature RL fires much more + often (the paper also calls this "recruited"). + * suppressed: else ``f0 >= turned_up[0] * fR``. + * stable: everything else. + + Args: + base_rates: ``[m]`` firing rates on the base checkpoint's inputs. + rl_rates: ``[m]`` firing rates after RL (max over RL checkpoints). + introduced: ``(base_below, rl_above)`` thresholds. + turned_up: ``(ratio, min_gain)`` thresholds, shared with suppressed. + + Returns: + ``{class_name: sorted feature ids}`` covering every feature exactly + once, keys ordered as :data:`CENSUS_CLASSES`. + + Raises: + ValueError: If the rate vectors' shapes differ. + """ + if base_rates.shape != rl_rates.shape: + raise ValueError( + f"Rate vectors must share a shape, got {tuple(base_rates.shape)} " + f"vs {tuple(rl_rates.shape)}." + ) + census: dict[str, list[int]] = {name: [] for name in CENSUS_CLASSES} + # A feature silent everywhere (f0 = fR = 0) lands in "suppressed" because + # 0 >= ratio * 0 holds; the criteria are kept exactly as the paper states + # them (no extra guard), and the trained dictionaries keep every feature + # alive via the dead-feature re-initialization, so the edge is theoretical. + for j in range(base_rates.shape[0]): + f0, fr = float(base_rates[j]), float(rl_rates[j]) + if f0 < introduced[0] and fr > introduced[1]: + census["introduced"].append(j) + elif fr >= turned_up[0] * f0 and (fr - f0) > turned_up[1]: + census["turned_up"].append(j) + elif f0 >= turned_up[0] * fr: + census["suppressed"].append(j) + else: + census["stable"].append(j) + return census + + +def promoted_tokens( + gsae: GSAE, model: ModelBackend, feature: int, top_k: int = 5 +) -> list[str]: + """The tokens a feature's direction promotes through the unembedding. + + Projects the decoder direction through the final norm's gain and the + unembedding, ``W_U (gamma * w_j)``, and returns the ``top_k`` most + promoted token strings; the cheap readout the paper uses to say what a + recruited gradient feature would push the model toward. + + Args: + gsae: The trained autoencoder. + model: Backend supplying ``final_norm``, ``unembed_weight``, and the + tokenizer. Use one fixed checkpoint for every feature. + feature: Feature id. + top_k: Tokens to return. + + Returns: + Token strings, most promoted first. + + Raises: + ValueError: If ``feature`` is out of range or ``top_k`` is not + positive. + """ + if not 0 <= feature < gsae.m: + raise ValueError(f"feature must be in [0, {gsae.m}), got {feature}.") + if top_k < 1: + raise ValueError(f"top_k must be >= 1, got {top_k}.") + direction = gsae.feature_direction(feature) + gain = model.final_norm.weight + direction = direction.to(gain.device, gain.dtype) * gain + unembed = model.unembed_weight + logits = unembed.float() @ direction.to(unembed.device).float() + ids = logits.topk(top_k).indices.tolist() + return model.tokenizer.convert_ids_to_tokens(ids) diff --git a/tests/test_gsae.py b/tests/test_gsae.py new file mode 100644 index 0000000..fcc2cfc --- /dev/null +++ b/tests/test_gsae.py @@ -0,0 +1,225 @@ +"""Tests for the GSAE: trainer, sparsity, census, and readout.""" + +from __future__ import annotations + +import pytest +import torch + +from murano import keys +from murano.steps.gradients import GradientStore +from murano.steps.gsae import ( + GSAE, + _reinit_dead_features, + classify_features, + firing_rates, + normalized_gradient_inputs, + promoted_tokens, +) + + +def _synthetic_inputs(n: int = 512, d: int = 16, rank: int = 4) -> torch.Tensor: + """Low-rank data a small dictionary can actually reconstruct.""" + generator = torch.Generator().manual_seed(0) + basis = torch.randn(rank, d, generator=generator) + codes = torch.randn(n, rank, generator=generator) + return codes @ basis + + +def _synthetic_store(n: int = 4, t: int = 6, d: int = 16) -> GradientStore: + generator = torch.Generator().manual_seed(1) + gradients = [torch.randn(t, d, generator=generator) for _ in range(n)] + return GradientStore( + gradients=gradients, + nll=torch.zeros(n), + finite=torch.ones(n).bool(), + layer=0, + window=t, + ) + + +# ------------------------------------------------------------------- inputs + + +def test_normalized_inputs_have_unit_rms_and_signed_rows(): + store = _synthetic_store(n=3) + advantages = torch.tensor([1.5, -0.5, 0.0]) + inputs = normalized_gradient_inputs(store, advantages) + # The zero-advantage rollout is dropped: 2 rollouts x 6 positions remain. + assert inputs.shape == (12, 16) + rms = inputs.pow(2).mean(dim=1).sqrt() + assert torch.allclose(rms, torch.ones_like(rms), atol=1e-5) + # The negative-advantage rollout's rows are sign-flipped copies. + grad = store.gradients[1].float() + expected = -grad / grad.pow(2).mean(dim=1, keepdim=True).sqrt() + assert torch.allclose(inputs[6:], expected, atol=1e-5) + + +def test_normalized_inputs_validation(): + store = _synthetic_store(n=2) + with pytest.raises(ValueError, match="advantages"): + normalized_gradient_inputs(store, torch.ones(3)) + with pytest.raises(ValueError, match="subsample"): + normalized_gradient_inputs(store, subsample=1.5) + + +def test_normalized_inputs_subsample(): + store = _synthetic_store(n=4, t=8) + inputs = normalized_gradient_inputs(store, subsample=0.25, seed=0) + assert inputs.shape == (8, 16) + + +# ------------------------------------------------------------------ trainer + + +def test_encode_is_exactly_k_sparse(): + gsae = GSAE.train(_synthetic_inputs(), m=32, k=3, epochs=1, batch_size=128) + codes = gsae.encode(_synthetic_inputs(n=64)) + assert codes.shape == (64, 32) + assert (codes.ne(0).sum(dim=1) == 3).all() + + +def test_train_shapes_and_unit_decoder_rows(): + gsae = GSAE.train(_synthetic_inputs(), m=32, k=4, epochs=2, batch_size=128) + assert gsae.w_enc.shape == (32, 16) + assert gsae.w_dec.shape == (32, 16) + assert gsae.b_pre.shape == (16,) + norms = gsae.w_dec.norm(dim=1) + assert torch.allclose(norms, torch.ones(32), atol=1e-4) + assert not gsae.w_enc.requires_grad and not gsae.w_dec.requires_grad + + +def test_training_reduces_reconstruction_error(): + inputs = _synthetic_inputs() + quick = GSAE.train(inputs, m=32, k=4, epochs=1, batch_size=128) + longer = GSAE.train(inputs, m=32, k=4, epochs=3, batch_size=128) + assert longer.fvu < quick.fvu + assert longer.fvu < 1.0 + + +def test_train_validation(): + with pytest.raises(ValueError, match="inputs"): + GSAE.train(torch.zeros(4), m=8, k=2) + with pytest.raises(ValueError, match="k <= m"): + GSAE.train(_synthetic_inputs(), m=8, k=9) + + +def test_dead_feature_reinit_replaces_and_normalizes(): + gsae = GSAE.train(_synthetic_inputs(), m=16, k=2, epochs=1, batch_size=128) + fired = torch.ones(16).bool() + fired[3] = False + fired[7] = False + before = gsae.w_dec.clone() + _reinit_dead_features( + gsae, _synthetic_inputs(), fired, torch.Generator().manual_seed(0) + ) + changed = ~torch.isclose(gsae.w_dec, before, atol=1e-8).all(dim=1) + assert set(changed.nonzero(as_tuple=True)[0].tolist()) == {3, 7} + assert torch.allclose(gsae.w_dec[3].norm(), torch.tensor(1.0), atol=1e-5) + + +def test_dead_feature_reinit_when_dead_outnumber_inputs(): + """More dead features than inputs to draw replacements from. + + Each revived feature takes one input's residual, so a dictionary far larger + than the pooled corpus leaves some dead features with nothing to take. They + keep their directions until the next epoch instead of failing the pass. + """ + inputs = _synthetic_inputs(n=12) + gsae = GSAE.train(inputs, m=64, k=2, epochs=1, batch_size=16) + fired = torch.zeros(64).bool() + fired[:2] = True + before = gsae.w_dec.clone() + _reinit_dead_features(gsae, inputs, fired, torch.Generator().manual_seed(0)) + changed = ~torch.isclose(gsae.w_dec, before, atol=1e-8).all(dim=1) + assert int(changed.sum()) == inputs.shape[0] + assert torch.allclose( + gsae.w_dec[changed].norm(dim=1), torch.ones(int(changed.sum())), atol=1e-5 + ) + + +def test_train_with_dictionary_larger_than_corpus(): + """The same shape through the public API: m far above the number of inputs.""" + gsae = GSAE.train(_synthetic_inputs(n=32), m=256, k=2, epochs=2, batch_size=16) + assert gsae.m == 256 + assert torch.isfinite(torch.tensor(gsae.fvu)) + + +# ------------------------------------------------------------------- census + + +def test_firing_rates_match_manual_count(): + gsae = GSAE.train(_synthetic_inputs(), m=16, k=4, epochs=1, batch_size=128) + inputs = _synthetic_inputs(n=50) + rates = firing_rates(gsae, inputs, batch_size=16) + manual = gsae.encode(inputs).ne(0).float().mean(dim=0) + assert torch.allclose(rates, manual, atol=1e-6) + # Every input selects k features, so the rates sum to k. + assert rates.sum().item() == pytest.approx(4.0, abs=1e-4) + + +def test_classify_features_thresholds(): + # introduced turned_up suppressed stable boundary + base_rates = torch.tensor([0.01, 0.10, 0.30, 0.20, 0.05]) + rl_rates = torch.tensor([0.20, 0.40, 0.05, 0.22, 0.90]) + census = classify_features(base_rates, rl_rates) + assert census["introduced"] == [0] + assert census["turned_up"] == [1, 4] # f0=0.05 is NOT < 0.05: not introduced + assert census["suppressed"] == [2] + assert census["stable"] == [3] + # Every feature lands in exactly one class. + assert sum(len(ids) for ids in census.values()) == 5 + + +def test_classify_features_shape_mismatch(): + with pytest.raises(ValueError, match="shape"): + classify_features(torch.zeros(3), torch.zeros(4)) + + +def test_classify_features_min_gain_guard(): + # A 3x ratio on a tiny base rate still needs an absolute gain > 0.005. + base_rates = torch.tensor([0.001]) + rl_rates = torch.tensor([0.004]) + census = classify_features(base_rates, rl_rates) + assert census["turned_up"] == [] + assert census["stable"] == [0] + + +# ------------------------------------------------------- readout and io + + +def test_promoted_tokens_returns_token_strings(murano_model): + gsae = GSAE.train( + _synthetic_inputs(d=murano_model.d_model), + m=8, + k=2, + epochs=1, + batch_size=128, + ) + tokens = promoted_tokens(gsae, murano_model, feature=0, top_k=3) + assert len(tokens) == 3 + assert all(isinstance(token, str) for token in tokens) + + +def test_gsae_save_load_roundtrip(tmp_path): + from murano.io import load_gsae, save_gsae + + gsae = GSAE.train(_synthetic_inputs(), m=16, k=4, epochs=1, batch_size=128) + path = tmp_path / "gsae.pt" + save_gsae(gsae, path) + loaded = load_gsae(path) + assert torch.equal(loaded.w_enc, gsae.w_enc) + assert torch.equal(loaded.w_dec, gsae.w_dec) + assert loaded.k == gsae.k + assert loaded.fvu == pytest.approx(gsae.fvu) + + +def test_save_results_dispatches_gsae(tmp_path): + from murano.io import save_results + from murano.results import Results + + results = Results() + results[keys.GSAE] = GSAE.train( + _synthetic_inputs(), m=8, k=2, epochs=1, batch_size=128 + ) + out = save_results(results, output_dir=str(tmp_path), model_id="tiny") + assert (out / "gsae" / "gsae.pt").exists() From aaeae5700fb83cf359a0ac3308d65b95793ff935 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Emirhan=20B=C3=B6ge?= Date: Fri, 31 Jul 2026 12:36:39 +0200 Subject: [PATCH 5/6] feat(io): persist gradient stores, GFC operators and GSAEs --- src/murano/__init__.py | 31 ++++++ src/murano/io.py | 199 ++++++++++++++++++++++++++++++++++- src/murano/steps/__init__.py | 44 ++++++++ 3 files changed, 273 insertions(+), 1 deletion(-) diff --git a/src/murano/__init__.py b/src/murano/__init__.py index eed629e..1f4e4c1 100644 --- a/src/murano/__init__.py +++ b/src/murano/__init__.py @@ -26,9 +26,12 @@ from murano.io import ( load_activation_store, load_attention, + load_gradient_store, load_metric_score, load_labeled_activation_store, load_logit_lens, + load_gfc_operator, + load_gsae, load_sae_activations, load_sae_examples, load_sae_labels, @@ -56,6 +59,14 @@ ) from murano.steps.logit_lens import LogitLens, LogitLensResult from murano.steps.logits import Logits + from murano.steps.gradients import GradientStore, RolloutBatch + from murano.steps.gfc import ( + GFCOperator, + GFCOperatorResult, + permutation_floor, + pairing_overlap, + ) + from murano.steps.gsae import GSAE from murano.steps.select import SelectComponents from murano.steps.sae import ( SAEActivationStore, @@ -87,6 +98,13 @@ "ComponentSelection": ("murano.artifacts", "ComponentSelection"), "SweepResult": ("murano.artifacts", "SweepResult"), "SelectComponents": ("murano.steps.select", "SelectComponents"), + "RolloutBatch": ("murano.steps.gradients", "RolloutBatch"), + "GradientStore": ("murano.steps.gradients", "GradientStore"), + "GFCOperator": ("murano.steps.gfc", "GFCOperator"), + "GFCOperatorResult": ("murano.steps.gfc", "GFCOperatorResult"), + "pairing_overlap": ("murano.steps.gfc", "pairing_overlap"), + "permutation_floor": ("murano.steps.gfc", "permutation_floor"), + "GSAE": ("murano.steps.gsae", "GSAE"), "MuranoDataset": ("murano.dataset", "MuranoDataset"), "LabeledDataset": ("murano.dataset", "LabeledDataset"), "CleanCorruptDataset": ("murano.dataset", "CleanCorruptDataset"), @@ -122,6 +140,9 @@ "load_sae_examples": ("murano.io", "load_sae_examples"), "load_sae_labels": ("murano.io", "load_sae_labels"), "load_metric_score": ("murano.io", "load_metric_score"), + "load_gradient_store": ("murano.io", "load_gradient_store"), + "load_gfc_operator": ("murano.io", "load_gfc_operator"), + "load_gsae": ("murano.io", "load_gsae"), "save_ablated_model": ("murano.io", "save_ablated_model"), } @@ -142,6 +163,13 @@ "ComponentSelection", "SweepResult", "SelectComponents", + "RolloutBatch", + "GradientStore", + "GFCOperator", + "GFCOperatorResult", + "pairing_overlap", + "permutation_floor", + "GSAE", "MuranoDataset", "LabeledDataset", "CleanCorruptDataset", @@ -176,6 +204,9 @@ "load_sae_examples", "load_sae_labels", "load_metric_score", + "load_gradient_store", + "load_gfc_operator", + "load_gsae", "save_ablated_model", ] diff --git a/src/murano/io.py b/src/murano/io.py index 60d18ef..2496906 100644 --- a/src/murano/io.py +++ b/src/murano/io.py @@ -1,7 +1,8 @@ """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 +``load_attention``, ``load_activation_store`` and the labeled variant, +``load_gradient_store``, ``load_gfc_operator``, and ``load_gsae``) 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 @@ -773,6 +774,144 @@ def load_sae_labels(path: str | Path) -> Any: ) +def save_gradient_store(store: Any, path: Path) -> None: + """Save a GradientStore to a .pt file. + + Args: + store: GradientStore to serialize. + path: Output path. Parent directory is created if missing. + """ + path.parent.mkdir(parents=True, exist_ok=True) + torch.save( + { + "gradients": store.gradients, + "nll": store.nll, + "finite": store.finite, + "layer": store.layer, + "window": store.window, + }, + path, + ) + logger.info("Saved gradient store to %s", path) + + +def load_gradient_store(path: str | Path) -> Any: + """Load a GradientStore from a .pt file. + + Args: + path: Path to the gradient_record.pt file. + + Returns: + GradientStore ready for downstream analyses. + """ + from murano.steps.gradients import GradientStore + + data = torch.load(path, weights_only=False) + return GradientStore( + gradients=data["gradients"], + nll=data["nll"], + finite=data["finite"], + layer=data["layer"], + window=data["window"], + ) + + +def save_gfc_operator(result: Any, path: Path) -> None: + """Save a GFCOperatorResult to a .pt file. + + Args: + result: GFCOperatorResult to serialize. + path: Output path. Parent directory is created if missing. + """ + path.parent.mkdir(parents=True, exist_ok=True) + torch.save( + { + "operator": result.operator, + "per_position": result.per_position, + "position_counts": result.position_counts, + "nll": result.nll, + "kept": result.kept, + "source_layer": result.source_layer, + "target_layer": result.target_layer, + "k": result.k, + "basis_seed": result.basis_seed, + "gradient_off": result.gradient_off, + "window": result.window, + }, + path, + ) + logger.info("Saved GFC operator to %s", path) + + +def load_gfc_operator(path: str | Path) -> Any: + """Load a GFCOperatorResult from a .pt file. + + Args: + path: Path to the gfc_operator.pt file. + + Returns: + GFCOperatorResult ready for overlap comparisons. + """ + from murano.steps.gfc import GFCOperatorResult + + data = torch.load(path, weights_only=False) + return GFCOperatorResult( + operator=data["operator"], + per_position=data["per_position"], + position_counts=data["position_counts"], + nll=data["nll"], + kept=data["kept"], + source_layer=data["source_layer"], + target_layer=data["target_layer"], + k=data["k"], + basis_seed=data["basis_seed"], + gradient_off=data["gradient_off"], + window=data["window"], + ) + + +def save_gsae(gsae: Any, path: Path) -> None: + """Save a trained GSAE to a .pt file. + + Args: + gsae: GSAE to serialize. + path: Output path. Parent directory is created if missing. + """ + path.parent.mkdir(parents=True, exist_ok=True) + torch.save( + { + "w_enc": gsae.w_enc, + "w_dec": gsae.w_dec, + "b_pre": gsae.b_pre, + "k": gsae.k, + "fvu": gsae.fvu, + }, + path, + ) + logger.info("Saved GSAE to %s", path) + + +def load_gsae(path: str | Path) -> Any: + """Load a trained GSAE from a .pt file. + + Args: + path: Path to the gsae.pt file. + + Returns: + GSAE ready for encoding and the firing-rate census. + """ + from murano.steps.gsae import GSAE + + data = torch.load(path, weights_only=False) + return GSAE( + w_enc=data["w_enc"], + w_dec=data["w_dec"], + b_pre=data["b_pre"], + k=data["k"], + fvu=data["fvu"], + ) + + def save_prompts(prompt_batch: PromptBatch, path: Path) -> None: """Save a PromptBatch to JSON. @@ -853,6 +992,9 @@ def register_artifact_serializer( def _serializer_registry() -> list[tuple[type, ArtifactSerializer]]: from murano.artifacts import ComponentSelection, SweepResult from murano.steps.attention import AttentionResult + from murano.steps.gradients import GradientStore + from murano.steps.gfc import GFCOperatorResult + from murano.steps.gsae import GSAE from murano.steps.logit_attribution import LogitAttributionResult from murano.steps.logit_lens import LogitLensResult from murano.steps.probe import ProbeResult @@ -1134,6 +1276,52 @@ def serialize_sae_labels( "n_labeled": len(labels.feat_ids), } + def serialize_gradient_store( + key: str, + store: Any, + out: Path, + _results: Any, + metadata: dict[str, Any], + ) -> None: + filename = "gradient_record.pt" if key == keys.GRADIENT_RECORD else f"{key}.pt" + save_gradient_store(store, out / "gradients" / filename) + metadata[key] = { + "n_rollouts": len(store.gradients), + "n_finite": int(store.finite.sum()), + "layer": store.layer, + "window": store.window, + } + + def serialize_gfc_operator( + key: str, + result: Any, + out: Path, + _results: Any, + metadata: dict[str, Any], + ) -> None: + filename = "gfc_operator.pt" if key == keys.GFC_OPERATOR else f"{key}.pt" + save_gfc_operator(result, out / "gfc" / filename) + metadata[key] = { + "source_layer": result.source_layer, + "target_layer": result.target_layer, + "k": result.k, + "basis_seed": result.basis_seed, + "gradient_off": result.gradient_off, + "window": result.window, + "n_kept": int(result.kept.sum()) if result.kept is not None else None, + } + + def serialize_gsae( + key: str, + gsae: Any, + out: Path, + _results: Any, + metadata: dict[str, Any], + ) -> None: + filename = "gsae.pt" if key == keys.GSAE else f"{key}.pt" + save_gsae(gsae, out / "gsae" / filename) + metadata[key] = {"m": gsae.m, "k": gsae.k, "fvu": gsae.fvu} + register_artifact_serializer(registry, PromptBatch, serialize_prompts) register_artifact_serializer(registry, SteeringResult, serialize_steering) register_artifact_serializer(registry, GenerationComparison, serialize_generations) @@ -1160,6 +1348,9 @@ def serialize_sae_labels( ) register_artifact_serializer(registry, SAEFeatureExamples, serialize_sae_examples) register_artifact_serializer(registry, SAEFeatureLabels, serialize_sae_labels) + register_artifact_serializer(registry, GradientStore, serialize_gradient_store) + register_artifact_serializer(registry, GFCOperatorResult, serialize_gfc_operator) + register_artifact_serializer(registry, GSAE, serialize_gsae) return registry @@ -1196,6 +1387,12 @@ def save_results( ├── sae/ # SAE activations + per-feature examples │ ├── sae_record.pt │ └── feature_examples.json + ├── gradients/ # recorded residual-stream gradients + │ └── gradient_record.pt + ├── gfc/ # GFC routing operators + │ └── gfc_operator.pt + ├── gsae/ # trained gradient sparse autoencoders + │ └── gsae.pt └── metadata.json Args: diff --git a/src/murano/steps/__init__.py b/src/murano/steps/__init__.py index 364c83a..92bf4c2 100644 --- a/src/murano/steps/__init__.py +++ b/src/murano/steps/__init__.py @@ -5,6 +5,31 @@ from murano.steps.prompts import LoadPrompts from murano.steps.paired import LoadPaired from murano.steps.record import Record +from murano.steps.gradients import ( + GradientStore, + LoadRollouts, + RecordGradients, + RolloutBatch, +) +from murano.steps.gfc import ( + GFCOperator, + GFCOperatorResult, + background_subtract, + marginal, + read_directions, + pairing_operator, + per_position_overlap, + permutation_floor, + pairing_overlap, +) +from murano.steps.gsae import ( + CENSUS_CLASSES, + GSAE, + classify_features, + firing_rates, + normalized_gradient_inputs, + promoted_tokens, +) from murano.steps.save import Save from murano.steps.train import SteeringVector from murano.steps.intervene import Intervene @@ -55,6 +80,25 @@ "LoadPrompts", "LoadPaired", "Record", + "RolloutBatch", + "LoadRollouts", + "GradientStore", + "RecordGradients", + "GFCOperator", + "GFCOperatorResult", + "read_directions", + "marginal", + "background_subtract", + "pairing_operator", + "pairing_overlap", + "permutation_floor", + "per_position_overlap", + "CENSUS_CLASSES", + "GSAE", + "classify_features", + "firing_rates", + "normalized_gradient_inputs", + "promoted_tokens", "Save", "SteeringVector", "Intervene", From 7152f1d2db81edfb400445ad23239835dae562e4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Emirhan=20B=C3=B6ge?= Date: Fri, 31 Jul 2026 12:37:07 +0200 Subject: [PATCH 6/6] docs: add the gradient interpretability notebook and wire the docs --- CHANGELOG.md | 6 +- docs/scripts/doc_check.py | 3 + docs/scripts/gen_api_docs.py | 3 + docs/scripts/gen_notebook_docs.py | 1 + .../gradient_interpretability.ipynb | 719 ++++++++++++++++++ tests/test_notebook_structure.py | 2 + 6 files changed, 733 insertions(+), 1 deletion(-) create mode 100644 notebooks/applications/gradient_interpretability.ipynb diff --git a/CHANGELOG.md b/CHANGELOG.md index d6ea123..c0d605b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- Gradient recording: `RolloutBatch` carries a saved rollout corpus (token ids plus the prompt/completion boundary, with optional rewards and GRPO groups), and the `LoadRollouts` / `RecordGradients` steps at `murano.steps` teacher-force a frozen checkpoint through those fixed tokens and store the completion log-likelihood gradient at a residual layer into a `GradientStore`, one `[T, d_model]` tensor per rollout plus the teacher-forced NLL. This is the direction factor of the first-epoch policy-gradient update read without an optimizer step, owned end to end by the new `MuranoModel.grad_forward`. +- `GFCOperator` step and `GFCOperatorResult` artifact, the gradient feature circuit routing operator: transport `k` seeded orthonormal read directions from a source layer to an earlier target layer with batched vector-Jacobian products against the frozen forward map, weighted by the recorded gradient (`gradient_off=True` is the transport-only control). `pairing_overlap` and `permutation_floor` score two operators against a measured chance level; the remaining comparison kernels (`read_directions`, `marginal`, `background_subtract`, `pairing_operator`, `per_position_overlap`) live at `murano.steps`, and `notebooks/applications/gradient_interpretability.ipynb` walks the whole flow on a tiny model. +- `GSAE`: a TopK sparse autoencoder over recorded gradients, with the minimal in-house trainer the census needs (murano's existing SAE path only loads pre-trained dictionaries, and none exist for gradients). `GSAE.train` fits one shared dictionary with unit-norm decoder rows and dead-feature re-initialization; the census workflow around it (`normalized_gradient_inputs`, `firing_rates`, `classify_features`, `CENSUS_CLASSES`, `promoted_tokens`) lives at `murano.steps`, and `notebooks/applications/gradient_interpretability.ipynb` runs the census next to the routing operator on the same tiny model. +- `save_gradient_store` / `load_gradient_store`, `save_gfc_operator` / `load_gfc_operator`, and `save_gsae` / `load_gsae` in `murano.io`, registered with `save_results` so the artifacts persist from a pipeline run like every other store. - `SelectComponents` step and `ComponentSelection` artifact: rank an attribution result (for example `LogitAttribution`) by magnitude, signed value, or most-negative, keep the top `top_k` or everything past a `threshold`, and write the chosen addresses for a downstream step to read. `Patch` / `PathPatch` / `Ablate` accept a `targets_key` / `senders_key` naming that selection, so attribute-then-patch runs as one pipeline instead of two with a hand-copied node list. - `Intervene` gains `direction_layers` (`"all"`, `"best"`, or an explicit layer list) for the `direction_key` steering path, so a one-pipeline steer can apply only the best-separating layer's direction instead of every recorded layer, which keeps deep models coherent. - `Sweep` step and `SweepResult` artifact: run a step chain once per item and harvest one or more metric keys. Every component study has this shape ("patch each head and measure what it restores", "zero each head and measure the damage", "steer at each layer"), and every notebook was hand-rolling it as a closure over the model, the task, and a baseline `Results`. `Sweep` forks the incoming `Results` per item, so the shared prefix runs once and the swept steps' writes stay out of the pipeline, and it derives its own read contract from the chain, so a missing upstream key fails pre-flight validation. A sweep over `Node` addresses publishes the same `{Node: float}` map an attribution does, so it feeds `SelectComponents` and `plot_head_matrix` with no adapter: attribute, sweep, select and path-patch now compose in one pipeline. @@ -19,7 +23,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `murano.tasks`: the two toy tasks the tutorials share, defined once and tested. `ioi()` builds the indirect-object-identification task as a `CleanCorruptDataset`; `sentiment()` returns contrastive sentences; `positive_word_rate()` is the crude scorer the steering notebooks use. - `plot_activation_projection` in `murano.plotting`: reduce one component's activations with any scikit-learn-style reducer (PCA, LDA, t-SNE, UMAP) and scatter them by class. It accepts both a contrastive `ActivationStore` and a `LabeledActivationStore`, so a single `Record` can feed both `Probe` and the plot. - `zmid` on `plot_heatmap` and `plot_head_matrix`, anchoring a diverging colorscale at zero so a signed statistic no longer shades zero as if it had a sign. -- `notebooks/getting_started.ipynb`, plus fourteen runnable notebooks under `notebooks/applications/`: steering, probing, logit lens, logit attribution, attention, ablation, activation patching, circuit discovery, metrics, custom pipeline, weight ablation, and the three sparse-autoencoder notebooks. All fifteen share one template, enforced by `tests/test_notebook_structure.py`, which also rejects a step constructed inside a loop (that is a hand-rolled `Sweep`) and a pipeline built inside a function. +- `notebooks/getting_started.ipynb`, plus fifteen runnable notebooks under `notebooks/applications/`: steering, probing, logit lens, logit attribution, attention, ablation, activation patching, circuit discovery, metrics, custom pipeline, weight ablation, gradient interpretability, and the three sparse-autoencoder notebooks. All sixteen share one template, enforced by `tests/test_notebook_structure.py`, which also rejects a step constructed inside a loop (that is a hand-rolled `Sweep`) and a pipeline built inside a function. - The notebooks now render on the documentation site, generated from the executed `.ipynb` files by `docs/scripts/gen_notebook_docs.py` at deploy time. ### Changed diff --git a/docs/scripts/doc_check.py b/docs/scripts/doc_check.py index 158e837..0da5207 100644 --- a/docs/scripts/doc_check.py +++ b/docs/scripts/doc_check.py @@ -28,6 +28,9 @@ "murano.io", "murano.steps.base", "murano.steps.record", + "murano.steps.gradients", + "murano.steps.gfc", + "murano.steps.gsae", "murano.steps.intervene", "murano.steps.train", "murano.steps.probe", diff --git a/docs/scripts/gen_api_docs.py b/docs/scripts/gen_api_docs.py index 34f5595..fdc1fd6 100644 --- a/docs/scripts/gen_api_docs.py +++ b/docs/scripts/gen_api_docs.py @@ -27,6 +27,9 @@ ("murano.io", "io"), ("murano.steps.base", "steps/base"), ("murano.steps.record", "steps/record"), + ("murano.steps.gradients", "steps/gradients"), + ("murano.steps.gfc", "steps/gfc"), + ("murano.steps.gsae", "steps/gsae"), ("murano.steps.intervene", "steps/intervene"), ("murano.steps.train", "steps/train"), ("murano.steps.probe", "steps/probe"), diff --git a/docs/scripts/gen_notebook_docs.py b/docs/scripts/gen_notebook_docs.py index 77e8cf0..3212d7a 100644 --- a/docs/scripts/gen_notebook_docs.py +++ b/docs/scripts/gen_notebook_docs.py @@ -42,6 +42,7 @@ "applications/metrics.ipynb", "applications/custom_pipeline.ipynb", "applications/weight_ablation.ipynb", + "applications/gradient_interpretability.ipynb", "applications/sae_features.ipynb", "applications/sae_steering.ipynb", "applications/sae_enrichment.ipynb", diff --git a/notebooks/applications/gradient_interpretability.ipynb b/notebooks/applications/gradient_interpretability.ipynb new file mode 100644 index 0000000..1c33202 --- /dev/null +++ b/notebooks/applications/gradient_interpretability.ipynb @@ -0,0 +1,719 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "afcad2d4", + "metadata": {}, + "source": [ + "# Gradient interpretability: the murano gradient toolkit\n", + "\n", + "Every other capture step reads what a model *represents*. This tutorial reads\n", + "what training would *move*, and walks the whole gradient toolkit on one tiny\n", + "model you can run on a laptop. It records the gradient of a frozen checkpoint,\n", + "builds the routing operator, a map of which residual-stream directions at a\n", + "late layer feed which coordinates at an earlier layer through the model's own\n", + "backward transport, uses that map as a fine-tune lens to tell reuse from\n", + "rewiring, and finally fits a gradient sparse autoencoder whose features name\n", + "what a fine-tune recruited.\n", + "\n", + "**Key questions**\n", + "\n", + "- What does the gradient of a frozen model look like at a residual layer, and\n", + " how do we record it without an optimizer step?\n", + "- How similar is the credit routing of a base model and a fine-tune of it, and\n", + " what does chance similarity look like?\n", + "- Can the routing map tell a fine-tune that reused the base's paths from one\n", + " that rewired them, apart from how much each moved the loss?\n", + "- Which gradient features does a fine-tune introduce, turn up, or suppress, and\n", + " what do they push the model toward?\n", + "\n", + "**Structure**\n", + "\n", + "1. A base and three updates of it\n", + "2. Record gradients on a shared rollout corpus\n", + "3. The routing operator and its overlap\n", + "4. The gradient-off control\n", + "5. The fine-tune lens: reuse versus rewire\n", + "6. A gradient dictionary\n", + "7. The feature census\n", + "8. What a recruited feature promotes\n", + "\n", + "**Model and data.** A tiny two-layer, 32-dimensional Llama built from scratch\n", + "in this notebook (random weights, word-level tokenizer), plus three updates of\n", + "it: a light fine-tune (reuse), a heavier fine-tune on a different rule\n", + "(rewire), and a same-size sign-flip of the light update (the inert control).\n", + "The \"rollouts\" are fixed random token sequences teacher-forced through every\n", + "checkpoint, so any difference is in the weights alone. Everything runs on CPU\n", + "in under a minute and nothing is downloaded.\n", + "\n", + "**Requirements.** The core install is enough, no extras: `pip install\n", + "murano-interp`." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "6fa93f96", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-22T13:54:57.312316Z", + "iopub.status.busy": "2026-07-22T13:54:57.312084Z", + "iopub.status.idle": "2026-07-22T13:57:34.358102Z", + "shell.execute_reply": "2026-07-22T13:57:34.356997Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "" + ] + }, + "execution_count": 1, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from pathlib import Path\n", + "\n", + "import torch\n", + "\n", + "from murano import (\n", + " MuranoModel,\n", + " Pipeline,\n", + " RolloutBatch,\n", + " GFCOperator,\n", + " GSAE,\n", + " keys,\n", + " permutation_floor,\n", + " pairing_overlap,\n", + ")\n", + "from murano.steps import LoadRollouts, RecordGradients\n", + "from murano.steps.gsae import (\n", + " normalized_gradient_inputs,\n", + " firing_rates,\n", + " classify_features,\n", + " promoted_tokens,\n", + " CENSUS_CLASSES,\n", + ")\n", + "\n", + "OUTPUT_DIR = \"murano_outputs/gradient_interpretability\"\n", + "Path(OUTPUT_DIR).mkdir(parents=True, exist_ok=True)\n", + "torch.manual_seed(0)" + ] + }, + { + "cell_type": "markdown", + "id": "595c6498", + "metadata": {}, + "source": [ + "## 1. A base and three updates of it\n", + "\n", + "We build one tiny random Llama as the base, then make three checkpoints from\n", + "it. The light fine-tune learns \"copy the token two positions back\", a rule the\n", + "base can already almost express, so it has little reason to rewire. The heavy\n", + "fine-tune learns a different, harder rule (map each symbol to a fixed partner)\n", + "for more steps, so it does. The inert arm takes the light fine-tune's weight\n", + "change, flips its sign per coordinate, and adds it back to the base: the same\n", + "support and per-coordinate size as a real update, in a direction unrelated to\n", + "what training did." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "876c9282", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-22T13:57:34.361312Z", + "iopub.status.busy": "2026-07-22T13:57:34.360611Z", + "iopub.status.idle": "2026-07-22T13:57:34.488991Z", + "shell.execute_reply": "2026-07-22T13:57:34.487802Z" + } + }, + "outputs": [], + "source": [ + "from tokenizers import Tokenizer\n", + "from tokenizers.models import WordLevel\n", + "from tokenizers.pre_tokenizers import Whitespace\n", + "from transformers import LlamaConfig, LlamaForCausalLM, PreTrainedTokenizerFast\n", + "\n", + "# Four specials plus a 20-symbol alphabet t0..t19.\n", + "SPECIALS = {\"\": 0, \"\": 1, \"\": 2, \"\": 3}\n", + "VOCAB = dict(SPECIALS)\n", + "VOCAB.update({f\"t{i}\": len(SPECIALS) + i for i in range(20)})\n", + "FIRST = len(SPECIALS)\n", + "\n", + "\n", + "def make_base(path: Path) -> None:\n", + " \"\"\"Save a tiny randomly initialized Llama and its word-level tokenizer.\"\"\"\n", + " tok = Tokenizer(WordLevel(vocab=dict(VOCAB), unk_token=\"\"))\n", + " tok.pre_tokenizer = Whitespace()\n", + " PreTrainedTokenizerFast(\n", + " tokenizer_object=tok,\n", + " unk_token=\"\",\n", + " pad_token=\"\",\n", + " bos_token=\"\",\n", + " eos_token=\"\",\n", + " model_max_length=48,\n", + " ).save_pretrained(path)\n", + " cfg = LlamaConfig(\n", + " vocab_size=len(VOCAB),\n", + " hidden_size=32,\n", + " intermediate_size=64,\n", + " num_hidden_layers=2,\n", + " num_attention_heads=4,\n", + " num_key_value_heads=4,\n", + " max_position_embeddings=48,\n", + " pad_token_id=0,\n", + " bos_token_id=1,\n", + " eos_token_id=2,\n", + " )\n", + " LlamaForCausalLM(cfg).save_pretrained(path)\n", + "\n", + "\n", + "base_dir = Path(OUTPUT_DIR) / \"base\"\n", + "make_base(base_dir)" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "ec28d4b1", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-22T13:57:34.491285Z", + "iopub.status.busy": "2026-07-22T13:57:34.491023Z", + "iopub.status.idle": "2026-07-22T13:57:48.232683Z", + "shell.execute_reply": "2026-07-22T13:57:48.231926Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "reuse fine-tune loss: 0.523\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "rewire fine-tune loss: 0.069\n" + ] + } + ], + "source": [ + "def train_to_targets(src: Path, dst: Path, inputs, targets, steps: int, lr: float) -> float:\n", + " \"\"\"Fine-tune the checkpoint at src so inputs predict targets, save to dst.\n", + "\n", + " The tokenizer files are copied from the base first; only the weights change.\n", + " Returns the final loss so the caller can see the update took.\n", + " \"\"\"\n", + " make_base(dst) # writes the tokenizer; weights are overwritten next\n", + " model = LlamaForCausalLM.from_pretrained(src)\n", + " optimizer = torch.optim.AdamW(model.parameters(), lr=lr)\n", + " labels = inputs.clone()\n", + " labels[:, :-1] = targets[:, 1:]\n", + " for _ in range(steps):\n", + " loss = model(input_ids=inputs, labels=labels).loss\n", + " loss.backward()\n", + " optimizer.step()\n", + " optimizer.zero_grad()\n", + " model.save_pretrained(dst)\n", + " return float(loss.detach())\n", + "\n", + "\n", + "rng = torch.Generator().manual_seed(1)\n", + "data = torch.randint(FIRST, len(VOCAB), (32, 16), generator=rng)\n", + "\n", + "# Reuse rule: copy the token two back. Rewire rule: map each symbol to a fixed\n", + "# partner (symbol i -> symbol (i * 7 + 3) mod 20), a relabeling the base shares\n", + "# no structure with, trained longer.\n", + "reuse_targets = torch.roll(data, shifts=2, dims=1)\n", + "partner = FIRST + (torch.arange(20) * 7 + 3) % 20\n", + "rewire_targets = partner[data - FIRST]\n", + "\n", + "reuse_dir = Path(OUTPUT_DIR) / \"reuse\"\n", + "rewire_dir = Path(OUTPUT_DIR) / \"rewire\"\n", + "print(\"reuse fine-tune loss: %.3f\" % train_to_targets(base_dir, reuse_dir, data, reuse_targets, 40, 3e-3))\n", + "print(\"rewire fine-tune loss: %.3f\" % train_to_targets(base_dir, rewire_dir, data, rewire_targets, 120, 3e-3))" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "08763375", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-22T13:57:48.234882Z", + "iopub.status.busy": "2026-07-22T13:57:48.234672Z", + "iopub.status.idle": "2026-07-22T13:57:48.376232Z", + "shell.execute_reply": "2026-07-22T13:57:48.375310Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "random arm written (sign-flipped reuse delta)\n" + ] + } + ], + "source": [ + "# The inert control: take the reuse fine-tune's weight change, flip its sign per\n", + "# coordinate, and add it back to the base. Same support and per-coordinate size\n", + "# as a real update, but a direction unrelated to what training did.\n", + "random_dir = Path(OUTPUT_DIR) / \"random\"\n", + "make_base(random_dir)\n", + "base_model = LlamaForCausalLM.from_pretrained(base_dir)\n", + "reuse_model = LlamaForCausalLM.from_pretrained(reuse_dir)\n", + "flip_gen = torch.Generator().manual_seed(2)\n", + "base_state = dict(base_model.named_parameters())\n", + "with torch.no_grad():\n", + " for name, reuse_param in reuse_model.named_parameters():\n", + " delta = reuse_param - base_state[name]\n", + " signs = torch.randint(0, 2, delta.shape, generator=flip_gen) * 2.0 - 1.0\n", + " base_state[name].add_(delta * signs)\n", + "base_model.save_pretrained(random_dir)\n", + "print(\"random arm written (sign-flipped reuse delta)\")" + ] + }, + { + "cell_type": "markdown", + "id": "e1a15834", + "metadata": {}, + "source": [ + "## 2. Record gradients on a shared rollout corpus\n", + "\n", + "A rollout is a prompt plus a completion; here we fix eight random token\n", + "sequences with a two-token prompt. `RecordGradients` teacher-forces the frozen\n", + "model through them, sums the log-likelihood of the completion tokens inside the\n", + "window, and takes one backward pass. The stored `[T, d_model]` tensor is the\n", + "gradient at the layer-1 residual stream, the direction training would push that\n", + "activation, read without an optimizer step. The store also carries the\n", + "teacher-forced NLL, used later to see how much each update moved the loss." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "175b82e5", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-22T13:57:48.378063Z", + "iopub.status.busy": "2026-07-22T13:57:48.377871Z", + "iopub.status.idle": "2026-07-22T13:57:48.932645Z", + "shell.execute_reply": "2026-07-22T13:57:48.931753Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "per-rollout gradient shape: (16, 32)\n", + "teacher-forced NLL of the first rollouts: tensor([3.1251, 3.1802, 3.1928])\n" + ] + } + ], + "source": [ + "# Eight fixed rollouts of 20 tokens, first two as the prompt; teacher-forced\n", + "# through every checkpoint so results differ only in weights.\n", + "walks = torch.randint(FIRST, len(VOCAB), (8, 20), generator=torch.Generator().manual_seed(3))\n", + "rollouts = RolloutBatch(\n", + " input_ids=[row for row in walks],\n", + " prompt_lengths=[2] * len(walks),\n", + ")\n", + "\n", + "base = MuranoModel(str(base_dir), device_map=\"cpu\", dtype=torch.float32)\n", + "recorded = Pipeline(\n", + " [\n", + " LoadRollouts(rollouts),\n", + " RecordGradients(base, layer=1, window=16),\n", + " ]\n", + ").run()\n", + "\n", + "store = recorded[keys.GRADIENT_RECORD]\n", + "print(\"per-rollout gradient shape:\", tuple(store.gradients[0].shape))\n", + "print(\"teacher-forced NLL of the first rollouts:\", store.nll[:3])" + ] + }, + { + "cell_type": "markdown", + "id": "36a53858", + "metadata": {}, + "source": [ + "## 3. The routing operator and its overlap\n", + "\n", + "The routing operator projects each position's gradient onto `k` fixed random\n", + "orthonormal directions (drawn once from a seed, shared by every checkpoint) and\n", + "transports each direction's component from layer 1 back to layer 0 with a\n", + "vector-Jacobian product against the frozen forward map. Row `i` of the\n", + "resulting `[k, d_model]` map says which layer-0 coordinates direction `i` feeds.\n", + "`pairing_overlap` compares two maps after removing the rank-one loudness\n", + "background and reducing to unit-strength singular pairs: 1 means the same\n", + "routing, and `permutation_floor` measures what chance agreement looks like\n", + "instead of assuming it is 0. Here we read the base against its light fine-tune." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "b903647d", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-22T13:57:48.934545Z", + "iopub.status.busy": "2026-07-22T13:57:48.934346Z", + "iopub.status.idle": "2026-07-22T13:57:49.465194Z", + "shell.execute_reply": "2026-07-22T13:57:49.464223Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "operator shape: (16, 32)\n", + "self overlap: 1.0\n", + "base vs reuse overlap: 0.868\n", + "permutation floor: -0.004 +- 0.062\n" + ] + } + ], + "source": [ + "reuse = MuranoModel(str(reuse_dir), device_map=\"cpu\", dtype=torch.float32)\n", + "\n", + "# Same corpus, same seeded basis, same layer pair, only the weights differ.\n", + "GF = dict(source_layer=1, target_layer=0, k=16, window=16)\n", + "base_map = Pipeline([LoadRollouts(rollouts), GFCOperator(base, **GF)]).run()[keys.GFC_OPERATOR]\n", + "reuse_map = Pipeline([LoadRollouts(rollouts), GFCOperator(reuse, **GF)]).run()[keys.GFC_OPERATOR]\n", + "\n", + "print(\"operator shape:\", tuple(base_map.operator.shape))\n", + "print(\"self overlap:\", round(pairing_overlap(base_map.operator, base_map.operator), 3))\n", + "score = pairing_overlap(base_map.operator, reuse_map.operator)\n", + "floor_mean, floor_sd = permutation_floor(base_map.operator, reuse_map.operator)\n", + "print(f\"base vs reuse overlap: {score:.3f}\")\n", + "print(f\"permutation floor: {floor_mean:.3f} +- {floor_sd:.3f}\")" + ] + }, + { + "cell_type": "markdown", + "id": "8f3196f7", + "metadata": {}, + "source": [ + "## 4. The gradient-off control\n", + "\n", + "Is the overlap above about the gradient, or about the network it flows through?\n", + "`gradient_off=True` sets every read strength to 1, so the operator reads the\n", + "frozen transport alone and the gradient contributes nothing. Whatever overlap\n", + "survives this switch is carried by the models' paths, not by the gradient read\n", + "along them. A fine-tune can therefore lower the full overlap in two ways: by\n", + "rewiring transport (the gradient-off overlap drops too) or by redistributing\n", + "the gradient over conserved paths (the gradient-off overlap stays high)." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "186c1a8b", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-22T13:57:49.466987Z", + "iopub.status.busy": "2026-07-22T13:57:49.466786Z", + "iopub.status.idle": "2026-07-22T13:57:49.702221Z", + "shell.execute_reply": "2026-07-22T13:57:49.701512Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "gradient-weighted overlap: 0.868\n", + "gradient-off overlap: 0.667\n" + ] + } + ], + "source": [ + "GF_OFF = dict(source_layer=1, target_layer=0, k=16, window=16, gradient_off=True)\n", + "base_off = Pipeline([LoadRollouts(rollouts), GFCOperator(base, **GF_OFF)]).run()[keys.GFC_OPERATOR]\n", + "reuse_off = Pipeline([LoadRollouts(rollouts), GFCOperator(reuse, **GF_OFF)]).run()[keys.GFC_OPERATOR]\n", + "\n", + "off_score = pairing_overlap(base_off.operator, reuse_off.operator)\n", + "print(f\"gradient-weighted overlap: {score:.3f}\")\n", + "print(f\"gradient-off overlap: {off_score:.3f}\")" + ] + }, + { + "cell_type": "markdown", + "id": "82b2b5a7", + "metadata": {}, + "source": [ + "## 5. The fine-tune lens: reuse versus rewire\n", + "\n", + "With the gradient-off operator per checkpoint, routing drift, one minus the\n", + "gradient-off overlap, reads how far each update moved the credit paths, and the\n", + "change in teacher-forced NLL reads how far it moved the loss. The light\n", + "fine-tune reuses the base's paths and drifts little though it trained; the heavy\n", + "fine-tune on a different rule drifts far because it is rewiring rather than\n", + "re-weighting. Two caveats this toy makes honest. The sign-flip arm is only\n", + "roughly inert here, because a sign flip on a 32-dimensional model is a large\n", + "perturbation, where at real scale the same control barely moves the loss. And\n", + "drift-per-nat, the paper's summary that divides drift by loss change, only means\n", + "something when the arms move the loss by comparable amounts, which real\n", + "checkpoints give and this toy does not. The clean matched-loss version is the\n", + "real-model study; the toy shows the mechanism." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "f2e8f6e8", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-22T13:57:49.704516Z", + "iopub.status.busy": "2026-07-22T13:57:49.704275Z", + "iopub.status.idle": "2026-07-22T13:57:50.452081Z", + "shell.execute_reply": "2026-07-22T13:57:50.451172Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "arm overlap drift dNLL\n", + "random 0.689 0.311 0.064\n", + "reuse 0.667 0.333 0.979\n", + "rewire 0.114 0.886 2.641\n" + ] + } + ], + "source": [ + "rewire = MuranoModel(str(rewire_dir), device_map=\"cpu\", dtype=torch.float32)\n", + "random_arm = MuranoModel(str(random_dir), device_map=\"cpu\", dtype=torch.float32)\n", + "rewire_off = Pipeline([LoadRollouts(rollouts), GFCOperator(rewire, **GF_OFF)]).run()[keys.GFC_OPERATOR]\n", + "random_off = Pipeline([LoadRollouts(rollouts), GFCOperator(random_arm, **GF_OFF)]).run()[keys.GFC_OPERATOR]\n", + "\n", + "\n", + "def mean_nll(result) -> float:\n", + " kept = result.kept.bool()\n", + " return float(result.nll[kept].mean())\n", + "\n", + "\n", + "# The headline is routing DRIFT = 1 - overlap: how far the update moved the\n", + "# credit paths. dNLL is how much it moved the loss. Reuse drifts little though\n", + "# it trained; rewire drifts far.\n", + "base_nll = mean_nll(base_off)\n", + "arms = {\"random\": random_off, \"reuse\": reuse_off, \"rewire\": rewire_off}\n", + "\n", + "print(f\"{'arm':8s} {'overlap':>9s} {'drift':>8s} {'dNLL':>9s}\")\n", + "for name, arm_off in arms.items():\n", + " overlap = pairing_overlap(base_off.operator, arm_off.operator, rank=8)\n", + " print(f\"{name:8s} {overlap:9.3f} {1.0 - overlap:8.3f} {mean_nll(arm_off) - base_nll:9.3f}\")" + ] + }, + { + "cell_type": "markdown", + "id": "4c874873", + "metadata": {}, + "source": [ + "## 6. A gradient dictionary\n", + "\n", + "The routing operator asks where credit flows. A gradient sparse autoencoder\n", + "asks what recurring directions the gradient is made of. We pool the\n", + "per-position gradients of the base and the rewiring fine-tune on a larger\n", + "shared corpus, RMS-normalize each position so loud and quiet ones weigh alike,\n", + "and fit a TopK autoencoder whose decoder rows are a dictionary of gradient\n", + "features. `normalized_gradient_inputs` does the pooling; `GSAE.train` fits the\n", + "dictionary and reports `fvu`, the fraction of gradient variance it leaves\n", + "unexplained. The two `RecordGradients` runs stay at cell level so the reader\n", + "sees both captures, not a hidden loop." + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "889d8dc8", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-22T13:57:50.454007Z", + "iopub.status.busy": "2026-07-22T13:57:50.453812Z", + "iopub.status.idle": "2026-07-22T13:57:52.699094Z", + "shell.execute_reply": "2026-07-22T13:57:52.698267Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "pooled positions per checkpoint: 1536\n", + "dictionary: 48 features, 4 active per input, fvu 0.702\n" + ] + } + ], + "source": [ + "# A larger shared corpus so the firing rates below are not read off a handful\n", + "# of positions. Same two-token prompt and window as the routing capture.\n", + "census_walks = torch.randint(\n", + " FIRST, len(VOCAB), (96, 20), generator=torch.Generator().manual_seed(4)\n", + ")\n", + "census_rollouts = RolloutBatch(\n", + " input_ids=[row for row in census_walks],\n", + " prompt_lengths=[2] * len(census_walks),\n", + ")\n", + "\n", + "base_grads = Pipeline(\n", + " [LoadRollouts(census_rollouts), RecordGradients(base, layer=1, window=16)]\n", + ").run()[keys.GRADIENT_RECORD]\n", + "rewire_grads = Pipeline(\n", + " [LoadRollouts(census_rollouts), RecordGradients(rewire, layer=1, window=16)]\n", + ").run()[keys.GRADIENT_RECORD]\n", + "\n", + "base_inputs = normalized_gradient_inputs(base_grads)\n", + "rewire_inputs = normalized_gradient_inputs(rewire_grads)\n", + "\n", + "# One dictionary sees both regimes, so a feature has a well-defined firing rate\n", + "# on each checkpoint.\n", + "gsae = GSAE.train(torch.cat([base_inputs, rewire_inputs]), m=48, k=4, epochs=6)\n", + "print(f\"pooled positions per checkpoint: {base_inputs.shape[0]}\")\n", + "print(f\"dictionary: {gsae.m} features, {gsae.k} active per input, fvu {gsae.fvu:.3f}\")" + ] + }, + { + "cell_type": "markdown", + "id": "052098c9", + "metadata": {}, + "source": [ + "## 7. The feature census\n", + "\n", + "With one dictionary, we can ask how often each feature fires on the base\n", + "checkpoint's gradients versus the rewiring fine-tune's. `firing_rates` counts,\n", + "per feature, the fraction of positions where it is among the active `k`.\n", + "`classify_features` sorts every feature into four classes by how that rate\n", + "shifts: introduced (rare in the base, common after), turned up (already\n", + "present, fires much more), suppressed (fires much less), and stable (little\n", + "change). On this toy the counts are only illustrative; at scale the paper runs\n", + "the same census on a 4096-dimension dictionary over millions of positions." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "22144dbf", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-22T13:57:52.701211Z", + "iopub.status.busy": "2026-07-22T13:57:52.700991Z", + "iopub.status.idle": "2026-07-22T13:57:52.710540Z", + "shell.execute_reply": "2026-07-22T13:57:52.709701Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "feature census (base gradients -> rewired gradients):\n", + " introduced: 0 features\n", + " turned_up: 6 features\n", + " suppressed: 1 features\n", + " stable: 41 features\n" + ] + } + ], + "source": [ + "base_rates = firing_rates(gsae, base_inputs)\n", + "rewire_rates = firing_rates(gsae, rewire_inputs)\n", + "census = classify_features(base_rates, rewire_rates)\n", + "\n", + "print(\"feature census (base gradients -> rewired gradients):\")\n", + "for name in CENSUS_CLASSES:\n", + " print(f\" {name:>11}: {len(census[name])} features\")" + ] + }, + { + "cell_type": "markdown", + "id": "0be60a43", + "metadata": {}, + "source": [ + "## 8. What a recruited feature promotes\n", + "\n", + "A gradient feature is a direction in the residual stream. Sending it through the\n", + "final norm's gain and the unembedding, `W_U (gamma * w_j)`, reads out which\n", + "tokens that direction pushes the model toward, the cheap decode the paper uses\n", + "to name a recruited feature. We take the feature whose firing rose most from\n", + "base to fine-tune and read its top tokens; on a random tiny model these are the\n", + "mechanism, not a meaning, but at scale the same call names what a recruited\n", + "feature would push a real model to say." + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "8064ac70", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-22T13:57:52.712299Z", + "iopub.status.busy": "2026-07-22T13:57:52.712099Z", + "iopub.status.idle": "2026-07-22T13:57:52.716428Z", + "shell.execute_reply": "2026-07-22T13:57:52.715585Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "feature 44: base fires 0.06, rewire fires 0.15\n", + "it promotes: ['t19', 't13', 't2', 't5', 't7']\n" + ] + } + ], + "source": [ + "# The feature the rewiring fine-tune leaned on most, decoded on that checkpoint.\n", + "feature = int((rewire_rates - base_rates).argmax())\n", + "print(f\"feature {feature}: base fires {base_rates[feature]:.2f}, \"\n", + " f\"rewire fires {rewire_rates[feature]:.2f}\")\n", + "print(\"it promotes:\", promoted_tokens(gsae, rewire, feature, top_k=5))" + ] + }, + { + "cell_type": "markdown", + "id": "f62d19bc", + "metadata": {}, + "source": [ + "## What next\n", + "\n", + "- [Logit attribution](logit_attribution.ipynb) decomposes a *forward* pass into per-component contributions, the representational complement of the gradient view here.\n", + "- [Weight ablation](weight_ablation.ipynb) removes components from the weights; pairing it with the fine-tune lens above asks which components carry the routing an update reuses." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/tests/test_notebook_structure.py b/tests/test_notebook_structure.py index 755d2b0..cb020b8 100644 --- a/tests/test_notebook_structure.py +++ b/tests/test_notebook_structure.py @@ -56,6 +56,8 @@ "Probe", "Record", "RecordAttention", + "RecordGradients", + "GFCOperator", "SAEEncode", "SteeringVector", "WeightAblation",