From ac2b070e0310f833b1b80321fe2289bb14de2bcc Mon Sep 17 00:00:00 2001 From: Alireza Bayat Date: Wed, 8 Jul 2026 22:09:31 +0200 Subject: [PATCH] feat: compose component discovery with patching and best-layer steering --- CHANGELOG.md | 5 + README.md | 4 +- src/murano/__init__.py | 6 + src/murano/artifacts.py | 41 +++++ src/murano/io.py | 57 +++++++ src/murano/keys.py | 1 + src/murano/steps/__init__.py | 2 + src/murano/steps/ablate.py | 58 +++++-- src/murano/steps/intervene.py | 67 ++++++++- src/murano/steps/patch.py | 13 +- src/murano/steps/path_patch.py | 70 +++++++-- src/murano/steps/select.py | 180 ++++++++++++++++++++++ tests/test_intervene.py | 111 ++++++++++++++ tests/test_select.py | 268 +++++++++++++++++++++++++++++++++ 14 files changed, 854 insertions(+), 29 deletions(-) create mode 100644 src/murano/steps/select.py create mode 100644 tests/test_select.py diff --git a/CHANGELOG.md b/CHANGELOG.md index a63c915..e3ed6d1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- `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. + ## [0.1.0a2] - 2026-07-08 ### Added diff --git a/README.md b/README.md index fe445f7..7db4d67 100644 --- a/README.md +++ b/README.md @@ -215,8 +215,8 @@ then, please cite this repository: ## Contact & Maintainers -Murano is developed and maintained by the -[Ubiquitous Knowledge Processing (UKP) Lab](https://www.ukp.tu-darmstadt.de/) at +Murano is developed and maintained by the Mechanistic Interpretability Team of +[Ubiquitous Knowledge Processing (UKP) Lab](https://www.informatik.tu-darmstadt.de/ukp/ukp_home/index.en.jsp) at the [Technische Universität Darmstadt](https://www.tu-darmstadt.de/). For questions, bug reports, and feature requests, please open an issue on the diff --git a/src/murano/__init__.py b/src/murano/__init__.py index 3c94281..5cd3a52 100644 --- a/src/murano/__init__.py +++ b/src/murano/__init__.py @@ -15,6 +15,7 @@ if TYPE_CHECKING: from murano.artifacts import ( + ComponentSelection, MetricScore, GenerationComparison, MetricComparison, @@ -55,6 +56,7 @@ ) from murano.steps.logit_lens import LogitLens, LogitLensResult from murano.steps.logits import Logits + from murano.steps.select import SelectComponents from murano.steps.sae import ( SAEActivationStore, SAEEncode, @@ -82,6 +84,8 @@ "GenerationComparison": ("murano.artifacts", "GenerationComparison"), "MetricComparison": ("murano.artifacts", "MetricComparison"), "MetricScore": ("murano.artifacts", "MetricScore"), + "ComponentSelection": ("murano.artifacts", "ComponentSelection"), + "SelectComponents": ("murano.steps.select", "SelectComponents"), "MuranoDataset": ("murano.dataset", "MuranoDataset"), "LabeledDataset": ("murano.dataset", "LabeledDataset"), "CleanCorruptDataset": ("murano.dataset", "CleanCorruptDataset"), @@ -135,6 +139,8 @@ "GenerationComparison", "MetricComparison", "MetricScore", + "ComponentSelection", + "SelectComponents", "MuranoDataset", "LabeledDataset", "CleanCorruptDataset", diff --git a/src/murano/artifacts.py b/src/murano/artifacts.py index 662a60d..96110b4 100644 --- a/src/murano/artifacts.py +++ b/src/murano/artifacts.py @@ -5,6 +5,8 @@ from dataclasses import dataclass, field from typing import Any +from murano.nodes import Node, NodeDict + @dataclass class PromptBatch: @@ -62,6 +64,45 @@ def __len__(self) -> int: return len(self.baseline_generations) +@dataclass +class ComponentSelection: + """An ordered, scored set of component addresses chosen by a discovery step. + + A discovery step (for example ranking heads by direct logit attribution) + writes this so a downstream :class:`~murano.steps.patch.Patch` or + :class:`~murano.steps.path_patch.PathPatch` can read its targets at run time, + letting attribute-then-patch compose in a single pipeline. It iterates over + its :attr:`nodes` (best first), so it drops straight into the same target + coercion the patch steps use for a hand-written address list. + + Attributes: + nodes: Selected component addresses, best first. + scores: ``{Node: float}`` score behind each selected node (the ranking + value that put it in the selection). + metadata: Provenance (source key, ranking criterion, cutoff). + """ + + nodes: list[Node] + scores: dict[Node, float] = field(default_factory=dict) + metadata: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + self.nodes = [Node.coerce(node) for node in self.nodes] + self.scores = NodeDict(self.scores) + + def __iter__(self): + return iter(self.nodes) + + def __len__(self) -> int: + return len(self.nodes) + + def __contains__(self, node: object) -> bool: + try: + return Node.coerce(node) in self.nodes # type: ignore[arg-type] + except (TypeError, ValueError): + return False + + @dataclass class MetricComparison: """Aggregate metric comparing baseline vs modified generations. diff --git a/src/murano/io.py b/src/murano/io.py index 7e10854..1583edb 100644 --- a/src/murano/io.py +++ b/src/murano/io.py @@ -420,6 +420,45 @@ def _restore(value: Any) -> Any: ) +def save_component_selection(selection: Any, path: Path) -> None: + """Save a ComponentSelection to a JSON file. + + Nodes and score keys are stored by their string address; the loader coerces + them back to Node objects. + + Args: + selection: ComponentSelection to serialize. + path: Output path. Parent directory is created if missing. + """ + path.parent.mkdir(parents=True, exist_ok=True) + data = { + "nodes": [str(node) for node in selection.nodes], + "scores": {str(k): v for k, v in selection.scores.items()}, + "metadata": selection.metadata, + } + _write_json(path, data) + logger.info("Saved component selection to %s", path) + + +def load_component_selection(path: str | Path) -> Any: + """Load a ComponentSelection from a file written by :func:`save_component_selection`. + + Args: + path: Path to the selection.json file. + + Returns: + ComponentSelection with string addresses coerced back to Node keys. + """ + from murano.artifacts import ComponentSelection + + data = json.loads(Path(path).read_text()) + return ComponentSelection( + nodes=data["nodes"], + scores=data.get("scores", {}), + metadata=data.get("metadata", {}), + ) + + def save_activation_store(activation_store: Any, path: Path) -> None: """Save an ActivationStore to a .pt file. @@ -738,6 +777,7 @@ def register_artifact_serializer( def _serializer_registry() -> list[tuple[type, ArtifactSerializer]]: + from murano.artifacts import ComponentSelection from murano.steps.attention import AttentionResult from murano.steps.logit_attribution import LogitAttributionResult from murano.steps.logit_lens import LogitLensResult @@ -906,6 +946,20 @@ def serialize_logit_attribution( "completeness_error": attribution.completeness_error, } + def serialize_component_selection( + key: str, + selection: Any, + out: Path, + _results: Any, + metadata: dict[str, Any], + ) -> None: + filename = "selection.json" if key == keys.SELECTION else f"{key}.json" + save_component_selection(selection, out / "selection" / filename) + metadata[key] = { + "nodes": [str(node) for node in selection.nodes], + "metadata": selection.metadata, + } + def serialize_activation_store( key: str, store: Any, @@ -1010,6 +1064,9 @@ def serialize_sae_labels( register_artifact_serializer( registry, LogitAttributionResult, serialize_logit_attribution ) + register_artifact_serializer( + registry, ComponentSelection, serialize_component_selection + ) register_artifact_serializer(registry, ActivationStore, serialize_activation_store) register_artifact_serializer( registry, LabeledActivationStore, serialize_labeled_activation_store diff --git a/src/murano/keys.py b/src/murano/keys.py index 71ef11d..303dfd0 100644 --- a/src/murano/keys.py +++ b/src/murano/keys.py @@ -22,6 +22,7 @@ LOGIT_ATTRIBUTION: Final = "logit_attribution" SAE_RECORD: Final = "sae_record" FEATURE_EXAMPLES: Final = "feature_examples" +SELECTION: Final = "selection" METRIC: Final = "metric" EVAL: Final = "eval" WEIGHT_ABLATION: Final = "weight_ablation" diff --git a/src/murano/steps/__init__.py b/src/murano/steps/__init__.py index 45e53d7..4cbdf3a 100644 --- a/src/murano/steps/__init__.py +++ b/src/murano/steps/__init__.py @@ -11,6 +11,7 @@ from murano.steps.ablate import Ablate from murano.steps.patch import Patch from murano.steps.path_patch import PathPatch +from murano.steps.select import SelectComponents from murano.steps.attention import ( AblateAttention, AttentionResult, @@ -59,6 +60,7 @@ "Ablate", "Patch", "PathPatch", + "SelectComponents", "RecordAttention", "AblateAttention", "AttentionResult", diff --git a/src/murano/steps/ablate.py b/src/murano/steps/ablate.py index 221d0a3..97655e6 100644 --- a/src/murano/steps/ablate.py +++ b/src/murano/steps/ablate.py @@ -33,7 +33,7 @@ from torch import randperm, tensor, zeros # pyright: ignore[reportPrivateImportUsage] from murano import keys -from murano.artifacts import PromptBatch +from murano.artifacts import ComponentSelection, PromptBatch from murano.logging import logger from murano.nodes import ( RESID_MID, @@ -272,7 +272,11 @@ class Ablate(Step): an iterable of addresses. A whole-component target ablates the module output; a head target (``Node(layer, "self_attn", head=h)``) ablates that head. One call is a single mode (all whole-component or all - per-head). + per-head). Pass this or ``targets_key``, not both. + targets_key: Results key holding a + :class:`~murano.artifacts.ComponentSelection` a discovery step wrote, + read at run time so attribute-then-ablate composes in one pipeline. + Pass this or ``targets``, not both. method: ``"zero"``, ``"mean"``, or ``"resample"``. mean_over: For ``method="mean"``, whether the mean pools over batch and positions (``"all"``) or is taken per position (``"position"``). @@ -306,10 +310,10 @@ class Ablate(Step): mask_key: Results key to write the attention mask under. Raises: - ValueError: If ``method``/``mean_over`` is unknown, ``targets`` is - empty or mixes modes, a method-specific argument is misused, or a - precomputed ``means`` table does not cover every target with a - right-sized vector. + ValueError: If neither or both of ``targets`` and ``targets_key`` are + given, ``method``/``mean_over`` is unknown, ``targets`` is empty or + mixes modes, a method-specific argument is misused, or a precomputed + ``means`` table does not cover every target with a right-sized vector. NotImplementedError: If per-head capture is requested for a module whose architecture's attention output projection is not recognized. """ @@ -320,9 +324,10 @@ class Ablate(Step): def __init__( self, model: ModelBackend, - targets: NodeSet | AddressLike | Iterable[AddressLike], + targets: NodeSet | AddressLike | Iterable[AddressLike] | None = None, method: Literal["zero", "mean", "resample"] = "zero", *, + targets_key: str | None = None, mean_over: Literal["all", "position"] = "all", means: dict[AddressLike, torch.Tensor] | None = None, source: Sequence[str] | None = None, @@ -334,6 +339,16 @@ def __init__( logits_key: str = keys.ABLATED_LOGITS, mask_key: str = keys.ATTENTION_MASK, ): + if (targets is None) == (targets_key is None): + raise ValueError( + "Pass exactly one of targets= (addresses) or targets_key= (a " + "ComponentSelection a discovery step wrote to Results)." + ) + if means is not None and targets_key is not None: + raise ValueError( + "means= needs the targets known up front, so pass explicit " + "targets=, not targets_key=." + ) if method not in _METHODS: raise ValueError(f"method must be one of {_METHODS}, got {method!r}") if mean_over not in _MEAN_OVER: @@ -368,7 +383,13 @@ def __init__( "apply when source=/source_key= supplies the replacement batch." ) - self.sites, self.per_head = _group_targets(_coerce_targets(targets)) + # With targets_key the sites are resolved from Results in __call__; hold + # valid-but-empty placeholders so the type stays a plain dict/bool. + if targets is not None: + self.sites, self.per_head = _group_targets(_coerce_targets(targets)) + else: + self.sites, self.per_head = {}, False + self.targets_key = targets_key self.model = model self.method = method @@ -385,11 +406,17 @@ def __init__( self.logits_key = logits_key self.mask_key = mask_key # Read the batch to run from prompts_key; a source_key, when given, names - # a second batch already in results whose activations are resampled in. - self.reads = [prompts_key] + ([source_key] if source_key is not None else []) + # a second batch already in results whose activations are resampled in; a + # targets_key, when given, names the ComponentSelection to ablate. + self.reads = ( + [prompts_key] + + ([source_key] if source_key is not None else []) + + ([targets_key] if targets_key is not None else []) + ) self.read_types = { prompts_key: PromptBatch, **({source_key: PromptBatch} if source_key is not None else {}), + **({targets_key: ComponentSelection} if targets_key is not None else {}), } self.writes = [logits_key, mask_key] self.write_types = {logits_key: torch.Tensor, mask_key: torch.Tensor} @@ -434,6 +461,13 @@ def _validate_means(self, means: NodeDict) -> None: ) def __call__(self, results: Results) -> Results: + if self.targets_key is not None: + # Resolve the targets a discovery step wrote before running, so + # attribute-then-ablate is one pipeline. The mode (per-head vs + # whole-component) is inferred here, the same as for explicit targets. + self.sites, self.per_head = _group_targets( + _coerce_targets(results[self.targets_key]) + ) prompt_batch: PromptBatch = results[self.prompts_key] prompts = prompt_batch.prompts tokens = self.model.tokenizer( @@ -558,6 +592,10 @@ def _capture(self, tokens: Any) -> dict[Site, torch.Tensor]: captured: dict[Site, torch.Tensor] = {} for module, layers in layers_by_module.items(): + # nnsight requires the per-layer saves in one trace to follow execution + # (ascending) order; self.sites keeps target-insertion order, which for + # a ranked ComponentSelection is not ascending, so sort here. + layers = sorted(layers) saved = {} with self.model.trace(tokens): for layer in layers: diff --git a/src/murano/steps/intervene.py b/src/murano/steps/intervene.py index a904e37..ded519f 100644 --- a/src/murano/steps/intervene.py +++ b/src/murano/steps/intervene.py @@ -119,6 +119,13 @@ class Intervene(Step): steer or ablate intervention, so deriving a direction and applying it compose in a single pipeline. + With ``direction_key``, ``direction_layers`` chooses which of the recorded + per-layer directions to apply. A ``SteeringResult`` carries one direction per + recorded layer; adding all of them at once over-steers deep models into + degenerate text, so ``"best"`` (the single best-separating layer) or an + explicit layer list keeps the flagship one-pipeline steer coherent, while + ``"all"`` preserves the every-layer behavior. + Reads from results: results['prompts']: PromptBatch results[direction_key]: SteeringResult, when ``direction_key`` is set. @@ -135,13 +142,19 @@ class Intervene(Step): mode: With ``direction_key``, ``"steer"`` adds the direction and ``"ablate"`` projects it out. Ignored when ``fn`` is given. alpha: Steering strength for ``mode="steer"``. + direction_layers: With ``direction_key``, which recorded directions to + apply: ``"all"`` every layer, ``"best"`` only the best-separating + layer (``SteeringResult.best_layer``), or an explicit list of layer + indices. Only valid with ``direction_key``. layers: Layers to apply the intervention at, or ``"all"``. modules: Module name(s) to apply the intervention at each layer. gen_kwargs: Keyword arguments forwarded to generation. Raises: ValueError: If neither or both of ``fn`` and ``direction_key`` are given, - or if ``mode`` is not ``"steer"`` or ``"ablate"``. + if ``mode`` is not ``"steer"`` or ``"ablate"``, or if + ``direction_layers`` is given without ``direction_key`` or is not + ``"all"``, ``"best"``, or a list of layer indices. """ writes = [keys.INTERVENE] @@ -166,6 +179,7 @@ def __init__( direction_key: str | None = None, mode: str = "steer", alpha: float = 1.0, + direction_layers: str | list[int] = "all", layers: list[int] | str = "all", modules: str | list[str] = "residual", gen_kwargs: dict | None = None, @@ -178,6 +192,21 @@ def __init__( ) if direction_key is not None and mode not in ("steer", "ablate"): raise ValueError(f"mode must be 'steer' or 'ablate', got {mode!r}") + if fn is not None and direction_layers != "all": + raise ValueError( + "direction_layers only applies with direction_key=; a prebuilt " + "fn= already fixes which layers it touches." + ) + if isinstance(direction_layers, str): + if direction_layers not in ("all", "best"): + raise ValueError( + f"direction_layers as a string must be 'all' or 'best', got " + f"{direction_layers!r}; pass a list of layer indices to choose " + f"specific layers." + ) + self.direction_layers: str | list[int] = direction_layers + else: + self.direction_layers = list(direction_layers) self.model = model self.fn = fn self.direction_key = direction_key @@ -235,11 +264,45 @@ def _resolve_fn(self, results: Results) -> Callable: # so the derive-then-apply arc is one pipeline. steer_direction / # ablate_direction normalize the direction keys to canonical Nodes. steering = results[self.direction_key] - directions = getattr(steering, "direction_per_layer", steering) + directions = self._select_directions(steering) if self.mode == "steer": return steer_direction(directions, self.alpha) return ablate_direction(directions) + def _select_directions(self, steering: object) -> dict[Node, Tensor]: + """Return the subset of recorded directions ``direction_layers`` selects. + + ``steering`` is the ``SteeringResult`` read from Results (its + ``direction_per_layer`` is a ``{Node: tensor}`` map); a bare mapping is + also accepted so a caller can write directions directly under the key. + + Raises: + ValueError: If ``direction_layers="best"`` but ``steering`` carries no + ``best_layer``, or an explicit layer list matches no recorded layer. + """ + directions = NodeDict(getattr(steering, "direction_per_layer", steering)) + if self.direction_layers == "all": + return directions + if self.direction_layers == "best": + best = getattr(steering, "best_layer", None) + if best is None: + raise ValueError( + "direction_layers='best' needs a SteeringResult with a " + "best_layer; the direction source has none." + ) + return NodeDict({best: directions[best]}) + wanted = set(self.direction_layers) + subset = NodeDict( + {node: vec for node, vec in directions.items() if node.layer in wanted} + ) + if not subset: + available = sorted({node.layer for node in directions}) + raise ValueError( + f"direction_layers={self.direction_layers} selected no recorded " + f"layer; the direction source covers layers {available}." + ) + return subset + def _generate_clean(self, text: str) -> str: """Generate without any intervention.""" return self.model.generate_with_hooks(text, gen_kwargs=self.gen_kwargs) diff --git a/src/murano/steps/patch.py b/src/murano/steps/patch.py index 35b5199..c62bcad 100644 --- a/src/murano/steps/patch.py +++ b/src/murano/steps/patch.py @@ -62,7 +62,11 @@ class Patch(Ablate): address, or an iterable of addresses. A whole-component target patches the module output; a head target (``Node(layer, "self_attn", head=h)``) patches that head. One call is a single mode (all whole-component or - all per-head). + all per-head). Pass this or ``targets_key``, not both. + targets_key: Results key holding a + :class:`~murano.artifacts.ComponentSelection` a discovery step wrote, + read at run time so attribute-then-patch composes in one pipeline. Pass + this or ``targets``, not both. base_key: Results key of the batch to patch into (default ``corrupt_prompts``). source_key: Results key of the batch supplying replacement activations @@ -75,15 +79,17 @@ class Patch(Ablate): mask_key: Results key to write the base attention mask under. Raises: - ValueError: If ``targets`` is empty or mixes modes, or the base and source + ValueError: If neither or both of ``targets`` and ``targets_key`` are + given, ``targets`` is empty or mixes modes, or the base and source prompts are not token-length-matched per pair. """ def __init__( self, model: ModelBackend, - targets: NodeSet | AddressLike | Iterable[AddressLike], + targets: NodeSet | AddressLike | Iterable[AddressLike] | None = None, *, + targets_key: str | None = None, base_key: str = keys.CORRUPT_PROMPTS, source_key: str = keys.PROMPTS, positions: int | Sequence[int] | torch.Tensor | None = None, @@ -94,6 +100,7 @@ def __init__( model, targets, method="resample", + targets_key=targets_key, source_key=source_key, positions=positions, prompts_key=base_key, diff --git a/src/murano/steps/path_patch.py b/src/murano/steps/path_patch.py index c1e2103..7e89d1c 100644 --- a/src/murano/steps/path_patch.py +++ b/src/murano/steps/path_patch.py @@ -46,7 +46,7 @@ from torch import Tensor, arange, tensor, zeros # pyright: ignore[reportPrivateImportUsage] from murano import keys -from murano.artifacts import PromptBatch +from murano.artifacts import ComponentSelection, PromptBatch from murano.logging import logger from murano.nodes import ( MLP, @@ -198,7 +198,12 @@ class PathPatch(Step): address, an iterable of addresses, or an :class:`~murano.nodes.Edge` (its source is the sender, its dest the receiver). Each sender is an attention head (``Node(layer, "self_attn", head=h)``) or an MLP - (``Node(layer, "mlp")``). + (``Node(layer, "mlp")``). Pass this or ``senders_key``, not both. + senders_key: Results key holding a + :class:`~murano.artifacts.ComponentSelection` a discovery step wrote, + read at run time so attribute-then-path-patch composes in one pipeline. + The ``receiver`` is still fixed at construction. Pass this or + ``senders``, not both. receiver: The receiving site. A residual-stream ``Node(layer, "resid_post")`` (default: the final residual, the last layer), or a head's query/key/value input ``Node(layer, "self_attn", head=h, @@ -218,10 +223,12 @@ class PathPatch(Step): mask_key: Results key to write the base attention mask under. Raises: - ValueError: If ``senders`` is empty or malformed, a sender/receiver layer - or sender/receiver head is out of range, a per-node receiver position is - given, both an ``Edge`` and a ``receiver`` are given, or the base and - source prompts are not token-length-matched per pair. + ValueError: If neither or both of ``senders`` and ``senders_key`` are + given, ``senders`` is empty or malformed, a sender/receiver layer or + sender/receiver head is out of range, a per-node receiver position is + given, both an ``Edge`` and a ``receiver`` are given, an ``Edge`` is + combined with ``senders_key``, or the base and source prompts are not + token-length-matched per pair. NotImplementedError: If ``receiver`` is neither a residual-stream site nor a head's Q/K/V input, or if a head's Q/K/V receiver is requested on an architecture with interleaved fused q/k/v (e.g. GPT-NeoX). @@ -230,9 +237,10 @@ class PathPatch(Step): def __init__( self, model: ModelBackend, - senders: NodeSet | AddressLike | Iterable[AddressLike] | Edge, + senders: NodeSet | AddressLike | Iterable[AddressLike] | Edge | None = None, receiver: AddressLike | None = None, *, + senders_key: str | None = None, base_key: str = keys.PROMPTS, source_key: str = keys.CORRUPT_PROMPTS, positions: int | Sequence[int] | Tensor | None = None, @@ -240,15 +248,33 @@ def __init__( logits_key: str = keys.PATH_PATCHED_LOGITS, mask_key: str = keys.PATH_PATCHED_MASK, ): + if (senders is None) == (senders_key is None): + raise ValueError( + "Pass exactly one of senders= (addresses or an Edge) or " + "senders_key= (a ComponentSelection a discovery step wrote)." + ) if isinstance(senders, Edge): + # An Edge is a senders= form, so senders_key= is already rejected by the + # exactly-one guard above; here only the receiver can conflict. if receiver is not None: raise ValueError("pass an Edge OR (senders, receiver), not both.") receiver = senders.dest senders = senders.source - self.attn_senders, self.mlp_senders = _group_senders(_coerce_targets(senders)) - if not self.attn_senders and not self.mlp_senders: - raise ValueError("PathPatch needs at least one sender (a head or an MLP).") + # With senders_key the senders are resolved from Results in __call__; hold + # valid-but-empty placeholders so _check_bounds validates the receiver here + # and the sender bounds are checked once the selection is read. + if senders is not None: + self.attn_senders, self.mlp_senders = _group_senders( + _coerce_targets(senders) + ) + if not self.attn_senders and not self.mlp_senders: + raise ValueError( + "PathPatch needs at least one sender (a head or an MLP)." + ) + else: + self.attn_senders, self.mlp_senders = {}, set() + self.senders_key = senders_key receiver_node = ( Node(model.n_layers - 1, RESID_POST) @@ -287,8 +313,14 @@ def __init__( self.freeze_mlps = freeze_mlps self.logits_key = logits_key self.mask_key = mask_key - self.reads = [base_key, source_key] - self.read_types = {base_key: PromptBatch, source_key: PromptBatch} + self.reads = [base_key, source_key] + ( + [senders_key] if senders_key is not None else [] + ) + self.read_types = { + base_key: PromptBatch, + source_key: PromptBatch, + **({senders_key: ComponentSelection} if senders_key is not None else {}), + } self.writes = [logits_key, mask_key] self.write_types = {logits_key: Tensor, mask_key: Tensor} @@ -358,6 +390,20 @@ def _check_bounds(self, model: ModelBackend, receiver: Node) -> None: ) def __call__(self, results: Results) -> Results: + if self.senders_key is not None: + # Resolve the senders a discovery step wrote before running, so + # attribute-then-path-patch is one pipeline. Group them, then check + # the sender bounds now that they are known (the receiver was already + # bounds-checked at construction). + self.attn_senders, self.mlp_senders = _group_senders( + _coerce_targets(results[self.senders_key]) + ) + if not self.attn_senders and not self.mlp_senders: + raise ValueError( + "PathPatch needs at least one sender (a head or an MLP); the " + f"selection at '{self.senders_key}' was empty." + ) + self._check_bounds(self.model, self.receiver) base_prompts = results[self.base_key].prompts source_prompts = results[self.source_key].prompts base_tokens, source_tokens, attention_mask = self._tokenize_pair( diff --git a/src/murano/steps/select.py b/src/murano/steps/select.py new file mode 100644 index 0000000..c387a55 --- /dev/null +++ b/src/murano/steps/select.py @@ -0,0 +1,180 @@ +"""SelectComponents step: rank an attribution result and pick the top components. + +Turns a per-component importance readout into a concrete target set. It reads a +:class:`~murano.steps.logit_attribution.LogitAttributionResult` (its +``contributions`` map each head and MLP to a signed contribution), ranks the +components, keeps the strongest, and writes a +:class:`~murano.artifacts.ComponentSelection`. A downstream +:class:`~murano.steps.patch.Patch` or :class:`~murano.steps.path_patch.PathPatch` +reads that selection at run time, so "attribute the important heads, then patch +them" runs as one pipeline instead of two with a hand-copied node list in between. +""" + +from __future__ import annotations + +from murano import keys +from murano.artifacts import ComponentSelection +from murano.logging import logger +from murano.nodes import Node, NodeDict, canonical_module +from murano.results import Results +from murano.steps.base import Step +from murano.steps.logit_attribution import LogitAttributionResult + +_BY = ("abs", "signed", "negative") + + +def _extract_scores(source: object) -> NodeDict: + """Return a component's ``{Node: float}`` scores from an attribution result. + + Accepts anything exposing a ``contributions`` mapping (a + :class:`LogitAttributionResult`), or a bare ``{Node: float}`` mapping so a + caller can rank scores it computed itself. + + Raises: + TypeError: If ``source`` exposes neither a ``contributions`` map nor a + mapping interface. + """ + contributions = getattr(source, "contributions", None) + if contributions is not None: + return NodeDict(contributions) + if hasattr(source, "items"): + return NodeDict(source) + raise TypeError( + f"SelectComponents needs a result with per-component scores (a " + f"LogitAttributionResult, or a {{Node: float}} mapping); got " + f"{type(source).__name__}." + ) + + +class SelectComponents(Step): + """Rank per-component scores and write the strongest as a ComponentSelection. + + Reads an attribution result's per-component scores, ranks them by magnitude + (``by="abs"``, the default), signed value, or most-negative, and keeps either + the top ``top_k`` or every component past ``threshold``. An optional + ``modules`` filter restricts the pool before ranking, which is how you pick + heads only for a :class:`~murano.steps.patch.Patch` (whose targets must be a + single mode, all heads or all whole-components). + + Reads from results: + results[source_key]: LogitAttributionResult (default ``logit_attribution``). + + Writes to results: + results[output_key]: ComponentSelection (default ``selection``). + + Args: + source_key: Results key of the attribution result to rank (default + ``logit_attribution``). + top_k: Keep the ``top_k`` highest-ranked components. Pass this or + ``threshold``, not both. + threshold: Keep every component whose score passes the cutoff: ``by="abs"`` + keeps ``abs(score) >= threshold``, ``"signed"`` keeps + ``score >= threshold``, ``"negative"`` keeps ``score <= threshold``. + by: Ranking criterion: ``"abs"`` (largest magnitude, either sign), + ``"signed"`` (most positive), or ``"negative"`` (most negative). + modules: Optional module name or list of names to keep before ranking + (e.g. ``"self_attn"`` for heads only); ``None`` ranks every component. + output_key: Results key to write the ComponentSelection under. + + Raises: + ValueError: If neither or both of ``top_k`` and ``threshold`` are given, + ``top_k`` is not positive, or ``by`` is not one of the allowed values. + """ + + def __init__( + self, + source_key: str = keys.LOGIT_ATTRIBUTION, + *, + top_k: int | None = None, + threshold: float | None = None, + by: str = "abs", + modules: str | list[str] | None = None, + output_key: str = keys.SELECTION, + ): + if (top_k is None) == (threshold is None): + raise ValueError( + "Pass exactly one of top_k= (keep the k strongest) or threshold= " + "(keep everything past a cutoff)." + ) + if top_k is not None and top_k <= 0: + raise ValueError(f"top_k must be positive, got {top_k}.") + if by not in _BY: + raise ValueError(f"by must be one of {_BY}, got {by!r}.") + + self.source_key = source_key + self.top_k = top_k + self.threshold = threshold + self.by = by + if modules is None: + self.modules: set[str] | None = None + else: + names = [modules] if isinstance(modules, str) else list(modules) + self.modules = {canonical_module(name) for name in names} + self.output_key = output_key + self.reads = [source_key] + self.read_types = {source_key: LogitAttributionResult} + self.writes = [output_key] + self.write_types = {output_key: ComponentSelection} + + def __call__(self, results: Results) -> Results: + scores = _extract_scores(results[self.source_key]) + pool = { + node: value + for node, value in scores.items() + if self.modules is None or node.module in self.modules + } + ranked = sorted( + pool.items(), key=lambda item: self._rank(item[1]), reverse=True + ) + + if self.top_k is not None: + chosen = ranked[: self.top_k] + else: + chosen = [(node, value) for node, value in ranked if self._passes(value)] + + if not chosen: + raise ValueError( + f"SelectComponents chose no component from '{self.source_key}' " + f"(by={self.by!r}, top_k={self.top_k}, threshold={self.threshold}, " + f"modules={sorted(self.modules) if self.modules else None}); " + f"loosen the cutoff or widen the module filter." + ) + + nodes: list[Node] = [node for node, _ in chosen] + logger.info( + "SelectComponents: kept %d of %d component(s) from '%s' by %s", + len(nodes), + len(pool), + self.source_key, + self.by, + ) + results[self.output_key] = ComponentSelection( + nodes=nodes, + scores={node: float(value) for node, value in chosen}, + metadata={ + "source_key": self.source_key, + "by": self.by, + "top_k": self.top_k, + "threshold": self.threshold, + "modules": sorted(self.modules) if self.modules else None, + "n_available": len(pool), + }, + ) + return results + + def _rank(self, value: float) -> float: + """Return the sort key for a score under the active criterion (higher wins).""" + if self.by == "abs": + return abs(value) + if self.by == "signed": + return value + return -value + + def _passes(self, value: float) -> bool: + """Return whether a score clears ``threshold`` under the active criterion.""" + assert self.threshold is not None + if self.by == "abs": + return abs(value) >= self.threshold + if self.by == "signed": + return value >= self.threshold + return value <= self.threshold diff --git a/tests/test_intervene.py b/tests/test_intervene.py index 61ea537..a37f9ba 100644 --- a/tests/test_intervene.py +++ b/tests/test_intervene.py @@ -256,6 +256,117 @@ def test_full_arc_validates(self, fixture, request): assert keys.INTERVENE in produced +class TestDirectionLayers: + """direction_layers chooses which recorded directions the steer applies, so a + single best layer (or an explicit subset) replaces the over-steering + every-layer default without changing how the direction was derived.""" + + def _steering(self, model) -> SteeringResult: + # Three recorded layers with layer 1 the best-separating one, so "best" + # and the explicit-list paths select a different subset than "all". + directions = { + Node(layer, "residual"): torch.ones(model.d_model) for layer in range(3) + } + return SteeringResult( + direction_per_layer=directions, + separation_scores={ + Node(0, "residual"): 0.1, + Node(1, "residual"): 0.9, + Node(2, "residual"): 0.2, + }, + best_layer=Node(1, "residual"), + ) + + def test_all_selects_every_recorded_layer(self, murano_model): + step = Intervene(murano_model, direction_key=keys.STEERING, mode="steer") + selected = step._select_directions(self._steering(murano_model)) + assert {node.layer for node in selected} == {0, 1, 2} + + def test_best_selects_only_the_best_layer(self, murano_model): + step = Intervene( + murano_model, + direction_key=keys.STEERING, + mode="steer", + direction_layers="best", + ) + selected = step._select_directions(self._steering(murano_model)) + assert list(selected) == [Node(1, "residual")] + + def test_explicit_list_selects_named_layers(self, murano_model): + step = Intervene( + murano_model, + direction_key=keys.STEERING, + mode="steer", + direction_layers=[0, 2], + ) + selected = step._select_directions(self._steering(murano_model)) + assert {node.layer for node in selected} == {0, 2} + + def test_best_only_perturbs_at_the_best_layer(self, murano_model): + # The built fn is identity at any layer whose direction was dropped, so a + # best-only steer touches only the best layer's activation. + model = murano_model + step = Intervene( + model, + direction_key=keys.STEERING, + mode="steer", + alpha=5.0, + direction_layers="best", + ) + results = Results() + results[keys.STEERING] = self._steering(model) + fn = step._resolve_fn(results) + + activation = torch.zeros(1, 1, model.d_model) + assert not torch.allclose(fn(activation, Node(1, "residual")), activation) + assert torch.allclose(fn(activation, Node(0, "residual")), activation) + + def test_best_runs_in_one_pipeline(self, murano_model): + model = murano_model + out = Pipeline( + [ + Load(_contrastive()), + Record(model, layers="all", position="mean"), + SteeringVector(normalize=True), + Intervene( + model, + direction_key=keys.STEERING, + mode="steer", + alpha=4.0, + direction_layers="best", + gen_kwargs={"max_new_tokens": 3, "do_sample": False}, + ), + ] + ).run() + assert keys.INTERVENE in out + + def test_rejects_unknown_string(self, murano_model): + with pytest.raises(ValueError, match="'all' or 'best'"): + Intervene( + murano_model, direction_key=keys.STEERING, direction_layers="worst" + ) + + def test_rejects_with_fn_path(self, murano_model): + with pytest.raises(ValueError, match="direction_layers only applies"): + Intervene(murano_model, fn=lambda a, k: a, direction_layers="best") + + def test_empty_selection_raises(self, murano_model): + step = Intervene( + murano_model, direction_key=keys.STEERING, direction_layers=[9] + ) + with pytest.raises(ValueError, match="selected no recorded layer"): + step._select_directions(self._steering(murano_model)) + + def test_best_without_best_layer_raises(self, murano_model): + # A bare direction mapping (no SteeringResult) has no best_layer, so "best" + # cannot resolve and says so. + step = Intervene( + murano_model, direction_key=keys.STEERING, direction_layers="best" + ) + with pytest.raises(ValueError, match="needs a SteeringResult"): + step._select_directions({Node(0, "residual"): torch.ones(4)}) + + class TestInterveneConstructorGuards: @pytest.mark.parametrize("fixture", FIXTURES) def test_requires_exactly_one_source(self, fixture, request): diff --git a/tests/test_select.py b/tests/test_select.py new file mode 100644 index 0000000..6fb41b1 --- /dev/null +++ b/tests/test_select.py @@ -0,0 +1,268 @@ +"""Tests for SelectComponents and the discover-then-patch composition. + +SelectComponents turns a per-component attribution readout into a +ComponentSelection, and Patch/PathPatch/Ablate can read that selection at run +time. The unit tests pin the ranking, filtering, and cutoff logic on a +hand-built LogitAttributionResult (weight-independent); the composition tests run +the tiny ``murano_model`` fixture on CPU to prove attribute-then-patch is one +pipeline and that a selection resolves to the same result as the explicit target +list it names. +""" + +from __future__ import annotations + +import pytest +import torch + +from murano import Pipeline, keys +from murano.artifacts import ComponentSelection +from murano.dataset import CleanCorruptDataset +from murano.nodes import MLP, SELF_ATTN, Node +from murano.results import Results +from murano.steps.ablate import Ablate +from murano.steps.logit_attribution import LogitAttribution, LogitAttributionResult +from murano.steps.paired import LoadPaired +from murano.steps.patch import Patch +from murano.steps.path_patch import PathPatch +from murano.steps.select import SelectComponents + +# Equal token length per pair so patch positions align (see test_patch.py). +CLEAN = ["hello world", "good world"] +CORRUPT = ["good world", "bad world"] + + +def _attribution() -> LogitAttributionResult: + """A small attribution result: two heads and one MLP with known scores.""" + return LogitAttributionResult( + contributions={ + Node(0, SELF_ATTN, head=0): 2.0, + Node(1, SELF_ATTN, head=1): -3.0, + Node(0, MLP): 0.5, + }, + embed_contribution=0.0, + other_contribution=0.0, + target="logit_diff", + total=0.0, + completeness_error=0.0, + ) + + +def _with_attribution() -> Results: + results = Results() + results[keys.LOGIT_ATTRIBUTION] = _attribution() + return results + + +# ── SelectComponents ranking / filtering ────────────────────────────── + + +class TestSelectComponents: + def test_top_k_by_abs_ranks_by_magnitude(self): + out = SelectComponents(top_k=2)(_with_attribution()) + selection = out[keys.SELECTION] + assert isinstance(selection, ComponentSelection) + # |-3| > |2| > |0.5|, so the two strongest are the layer-1 head then the + # layer-0 head, best first. + assert selection.nodes == [ + Node(1, SELF_ATTN, head=1), + Node(0, SELF_ATTN, head=0), + ] + + def test_signed_prefers_most_positive(self): + out = SelectComponents(top_k=1, by="signed")(_with_attribution()) + assert out[keys.SELECTION].nodes == [Node(0, SELF_ATTN, head=0)] + + def test_negative_prefers_most_negative(self): + out = SelectComponents(top_k=1, by="negative")(_with_attribution()) + assert out[keys.SELECTION].nodes == [Node(1, SELF_ATTN, head=1)] + + def test_threshold_abs_keeps_everything_past_cutoff(self): + out = SelectComponents(threshold=1.0)(_with_attribution()) + nodes = out[keys.SELECTION].nodes + assert set(nodes) == {Node(0, SELF_ATTN, head=0), Node(1, SELF_ATTN, head=1)} + + def test_modules_filter_keeps_heads_only(self): + # Filtering to self_attn drops the MLP, so a downstream Patch sees a single + # (per-head) mode. "attn_out" is an accepted alias for the canonical name. + out = SelectComponents(top_k=5, modules="attn_out")(_with_attribution()) + assert all(node.module == SELF_ATTN for node in out[keys.SELECTION].nodes) + assert Node(0, MLP) not in out[keys.SELECTION].nodes + + def test_scores_recorded_on_selection(self): + selection = SelectComponents(top_k=2)(_with_attribution())[keys.SELECTION] + assert selection.scores[Node(1, SELF_ATTN, head=1)] == -3.0 + + def test_empty_selection_raises(self): + with pytest.raises(ValueError, match="chose no component"): + SelectComponents(threshold=100.0)(_with_attribution()) + + def test_writes_selection_type(self): + step = SelectComponents(top_k=1) + assert step.reads == [keys.LOGIT_ATTRIBUTION] + assert step.writes == [keys.SELECTION] + assert step.write_types == {keys.SELECTION: ComponentSelection} + + def test_constructor_guards(self): + with pytest.raises(ValueError, match="exactly one"): + SelectComponents() + with pytest.raises(ValueError, match="exactly one"): + SelectComponents(top_k=1, threshold=0.5) + with pytest.raises(ValueError, match="top_k must be positive"): + SelectComponents(top_k=0) + with pytest.raises(ValueError, match="by must be one of"): + SelectComponents(top_k=1, by="huge") + + +# ── Discover-then-patch composition ──────────────────────────────────── + + +def _loaded(clean=CLEAN, corrupt=CORRUPT, correct=5, incorrect=6): + ds = CleanCorruptDataset( + clean=clean, corrupt=corrupt, correct=correct, incorrect=incorrect + ) + return LoadPaired(ds)(Results()) + + +class TestDiscoverThenPatch: + def test_patch_targets_key_matches_explicit_targets(self, murano_model): + # A selection resolves at run time to the same sites the explicit list + # names, so the patched logits are identical. + targets = [Node(0, SELF_ATTN, head=0)] + selection = ComponentSelection(nodes=list(targets)) + + explicit = Patch(murano_model, targets)(_loaded())[keys.PATCHED_LOGITS] + + results = _loaded() + results[keys.SELECTION] = selection + via_key = Patch(murano_model, targets_key=keys.SELECTION)(results)[ + keys.PATCHED_LOGITS + ] + assert torch.equal(explicit, via_key) + + def test_patch_targets_key_declares_the_read(self, murano_model): + step = Patch(murano_model, targets_key=keys.SELECTION) + assert keys.SELECTION in step.reads + assert step.read_types[keys.SELECTION] is ComponentSelection + + def test_missing_selection_fails_preflight(self, murano_model): + pipe = Pipeline( + [ + LoadPaired(CleanCorruptDataset(clean=CLEAN, corrupt=CORRUPT)), + Patch(murano_model, targets_key=keys.SELECTION), + ] + ) + with pytest.raises(KeyError, match=keys.SELECTION): + pipe.validate() + + def test_full_attribute_then_patch_arc_validates(self, murano_model): + produced = Pipeline( + [ + LoadPaired( + CleanCorruptDataset( + clean=CLEAN, corrupt=CORRUPT, correct=5, incorrect=6 + ) + ), + LogitAttribution(murano_model, correct=5, incorrect=6), + SelectComponents(top_k=2, modules="self_attn"), + Patch(murano_model, targets_key=keys.SELECTION), + ] + ).validate() + assert keys.PATCHED_LOGITS in produced + + def test_full_attribute_then_patch_arc_runs(self, murano_model): + out = Pipeline( + [ + LoadPaired( + CleanCorruptDataset( + clean=CLEAN, corrupt=CORRUPT, correct=5, incorrect=6 + ) + ), + LogitAttribution(murano_model, correct=5, incorrect=6), + SelectComponents(top_k=2, modules="self_attn"), + Patch(murano_model, targets_key=keys.SELECTION), + ] + ).run() + assert torch.isfinite(out[keys.PATCHED_LOGITS]).all() + # The selection is per-head, so Patch inferred per-head mode from it. + assert out[keys.SELECTION].nodes and all( + node.head is not None for node in out[keys.SELECTION].nodes + ) + + +class TestDiscoverThenPathPatch: + def test_senders_key_matches_explicit_senders(self, murano_model): + senders = [Node(0, SELF_ATTN, head=0)] + selection = ComponentSelection(nodes=list(senders)) + + explicit = PathPatch(murano_model, senders)(_loaded())[keys.PATH_PATCHED_LOGITS] + + results = _loaded() + results[keys.SELECTION] = selection + via_key = PathPatch(murano_model, senders_key=keys.SELECTION)(results)[ + keys.PATH_PATCHED_LOGITS + ] + assert torch.equal(explicit, via_key) + + def test_senders_key_declares_the_read(self, murano_model): + step = PathPatch(murano_model, senders_key=keys.SELECTION) + assert keys.SELECTION in step.reads + assert step.read_types[keys.SELECTION] is ComponentSelection + + def test_missing_selection_fails_preflight(self, murano_model): + pipe = Pipeline( + [ + LoadPaired(CleanCorruptDataset(clean=CLEAN, corrupt=CORRUPT)), + PathPatch(murano_model, senders_key=keys.SELECTION), + ] + ) + with pytest.raises(KeyError, match=keys.SELECTION): + pipe.validate() + + def test_empty_selection_at_runtime_raises(self, murano_model): + results = _loaded() + results[keys.SELECTION] = ComponentSelection(nodes=[]) + with pytest.raises(ValueError, match="at least one sender"): + PathPatch(murano_model, senders_key=keys.SELECTION)(results) + + +# ── Deferred-target constructor guards ───────────────────────────────── + + +class TestDeferredTargetGuards: + def test_ablate_requires_exactly_one_target_source(self, murano_model): + with pytest.raises(ValueError, match="exactly one"): + Ablate(murano_model) + with pytest.raises(ValueError, match="exactly one"): + Ablate(murano_model, 0, targets_key=keys.SELECTION) + + def test_patch_requires_exactly_one_target_source(self, murano_model): + with pytest.raises(ValueError, match="exactly one"): + Patch(murano_model) + with pytest.raises(ValueError, match="exactly one"): + Patch(murano_model, 0, targets_key=keys.SELECTION) + + def test_means_with_targets_key_rejected(self, murano_model): + with pytest.raises(ValueError, match="means="): + Ablate( + murano_model, + targets_key=keys.SELECTION, + method="mean", + means={Node(0, "residual"): torch.zeros(murano_model.d_model)}, + ) + + def test_pathpatch_requires_exactly_one_sender_source(self, murano_model): + with pytest.raises(ValueError, match="exactly one"): + PathPatch(murano_model) + with pytest.raises(ValueError, match="exactly one"): + PathPatch(murano_model, 0, senders_key=keys.SELECTION) + + def test_pathpatch_edge_with_senders_key_rejected(self, murano_model): + from murano.nodes import Edge + + edge = Edge( + Node(0, SELF_ATTN, head=0), Node(murano_model.n_layers - 1, "residual") + ) + # An Edge is a senders= form, so combining it with senders_key= trips the + # exactly-one guard. + with pytest.raises(ValueError, match="exactly one"): + PathPatch(murano_model, edge, senders_key=keys.SELECTION)