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

[emirhanboge]

Summary 📝

Adds a gradient-side capability to murano. Every existing capture path answers what a model represents; this answers what training would move. It records the gradient of the completion log-likelihood with respect to the residual stream on a frozen checkpoint, builds a cross-layer map of where that gradient travels between two layers, and trains a sparse dictionary over the recorded gradients so the directions of the training signal can be read and compared across checkpoints.

Three new step modules, one new MuranoModel method, save/load for the three new artifacts, tests, and one runnable notebook. No new dependencies, and pyproject.toml is untouched.

Details

1. MuranoModel.grad_forward — the graph-preserving counterpart of forward_logits. The proxy capture paths all detach, so gradient-based steps need a forward whose activations stay attached to the autograd graph. It runs the base module natively with plain forward hooks on the requested decoder layers, the same raw-module seam the generation interventions already use, and returns the block outputs undetached.

2. murano.steps.gradients — RolloutBatch holds a fixed rollout corpus (token ids plus the prompt boundary, with optional rewards and GRPO group ids). LoadRollouts and RecordGradients teacher-force a frozen checkpoint through those ids and write a GradientStore. Rollouts are processed one at a time rather than in padded batches, because padding would put the completion window at a different sequence offset per row.

3. murano.steps.gfc — GFCOperator reads the gradient at a source layer, projects it onto k orthonormal directions fixed by a seed before any checkpoint is loaded, and transports each direction's seed field back to an earlier target layer with one batched vector-Jacobian product. gradient_off=True holds every read strength at 1 so the operator reads the frozen transport alone. pairing_overlap compares two operators after removing the rank-one marginal background, and permutation_floor measures the chance level of that comparison.

4. murano.steps.gsae — a TopK sparse autoencoder over the recorded gradients. The existing SAE path loads pre-trained dictionaries and none exist for gradients, so this ships the minimal trainer the feature census needs: one dictionary fit once over pooled inputs from several checkpoints, decoder rows renormalised after every step, and dead features re-initialised from reconstruction residuals between epochs.

5. murano.io — save and load for all three artifacts, registered with the existing artifact serializer registry so they persist from a pipeline run like every other store.

6. notebooks/applications/gradient_interpretability.ipynb — runs the whole flow on a tiny model, following the shared template that tests/test_notebook_structure.py enforces. Wired into the docs generators.

Checks

• Tested Changes
• Stakeholder Approval