Gradient Interpretability Toolkit (gradient recording, GFC operator, GSAE)

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

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",

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

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",