feat: compose component discovery with patching and best-layer steering

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

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

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

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.

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

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"

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

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: