brainpatch 1.2.0__py3-none-any.whl

brainpatch/research/ml/patch_search.py

@@ -0,0 +1,337 @@
+"""Candidate-feature search from behavioural contrast data.
+
+Pipeline::
+
+    contrast pairs -> activations for positive/negative responses
+                   -> SAE encode
+                   -> per-feature mean activation difference
+                   -> ranked candidates              (correlational)
+                   -> causal test with controls      (interventional)
+                   -> strength sweep
+                   -> greedy sparse combination
+                   -> held-out evaluation
+
+Only the first three steps are cheap. Everything after that costs GPU time per
+candidate, which is why :func:`rank_candidate_features` returns a *small*
+shortlist and the causal stages take an explicit budget. A search that fans out
+over hundreds of candidates is how a $10 budget disappears.
+
+The ranking step produces correlational evidence only. A feature that activates
+more on "independent criticism" continuations than on "agreeable" ones is a
+feature that *tracks* that distinction in this fixture -- it is not an
+anti-sycophancy control knob until an intervention says so.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Sequence
+
+import torch
+
+from brainpatch.research.ml.generation import build_chat_prompt
+from brainpatch.research.ml.hooks import ResidualCapture
+from brainpatch.research.ml.model import ModelBundle
+from brainpatch.research.ml.sae import TopKSAE
+from brainpatch.schemas.contrast import ContrastSet
+
+
+@dataclass
+class CandidateFeature:
+    """A feature ranked by its activation difference across a contrast set."""
+
+    feature_id: int
+    mean_diff: float
+    """Mean(positive activation) - Mean(negative activation), normalized space."""
+    mean_positive: float
+    mean_negative: float
+    positive_fire_rate: float
+    negative_fire_rate: float
+    effect_size: float
+    """Standardised difference (Cohen's d style), using pooled per-token std."""
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "feature_id": self.feature_id,
+            "mean_diff": self.mean_diff,
+            "mean_positive": self.mean_positive,
+            "mean_negative": self.mean_negative,
+            "positive_fire_rate": self.positive_fire_rate,
+            "negative_fire_rate": self.negative_fire_rate,
+            "effect_size": self.effect_size,
+            "evidence": "correlational -- activation difference only, no causal test",
+        }
+
+
+@torch.inference_mode()
+def collect_response_activations(
+    bundle: ModelBundle,
+    texts: Sequence[str],
+    *,
+    layer: int,
+    prompts: Sequence[str] | None = None,
+    max_length: int = 384,
+) -> torch.Tensor:
+    """Capture residual activations over the *response* tokens only.
+
+    When ``prompts`` is supplied, activations from the prompt portion are
+    excluded. This matters: the prompt is identical across a contrast pair, so
+    including its activations would dilute the signal with tokens that carry no
+    information about the contrast.
+    """
+    capture = ResidualCapture(to_cpu=True, dtype=torch.float32)
+    handle = capture.attach(bundle.layer_module(layer))
+    collected: list[torch.Tensor] = []
+    try:
+        for i, text in enumerate(texts):
+            prompt = prompts[i] if prompts is not None else None
+            full = (prompt + text) if prompt else text
+            encoded = bundle.tokenizer(
+                full, return_tensors="pt", truncation=True, max_length=max_length
+            ).to(bundle.device)
+            capture.activations = None
+            bundle.model(**encoded, use_cache=False)
+            hidden = capture.activations
+            if hidden is None:
+                raise RuntimeError("capture hook did not fire during contrast collection")
+
+            start = 0
+            if prompt:
+                prompt_len = bundle.tokenizer(
+                    prompt, return_tensors="pt", truncation=True, max_length=max_length
+                ).input_ids.shape[1]
+                start = min(prompt_len, hidden.shape[1] - 1)
+            collected.append(hidden[0, start:, :])
+    finally:
+        handle.remove()
+    return torch.cat(collected, dim=0)
+
+
+def rank_candidate_features(
+    bundle: ModelBundle,
+    sae: TopKSAE,
+    contrast_set: ContrastSet,
+    *,
+    layer: int,
+    input_scale: float,
+    top_n: int = 20,
+    use_chat_template: bool = True,
+    min_fire_rate: float = 0.01,
+    max_fire_rate: float = 0.9,
+) -> list[CandidateFeature]:
+    """Rank features by how differently they fire on positive vs negative responses.
+
+    Parameters
+    ----------
+    min_fire_rate, max_fire_rate:
+        Discard features that essentially never fire (no signal) or fire almost
+        everywhere (modelling something global rather than the contrast).
+
+    Returns
+    -------
+    list[CandidateFeature]
+        Sorted by absolute effect size, longest-first, truncated to ``top_n``.
+        **Correlational evidence only.**
+    """
+    prompts = [
+        build_chat_prompt(bundle.tokenizer, e.prompt) if use_chat_template else e.prompt
+        for e in contrast_set
+    ]
+    positives = [e.positive_response for e in contrast_set]
+    negatives = [e.negative_response for e in contrast_set]
+
+    pos_acts = collect_response_activations(bundle, positives, layer=layer, prompts=prompts)
+    neg_acts = collect_response_activations(bundle, negatives, layer=layer, prompts=prompts)
+
+    sae.eval()
+    device = next(sae.parameters()).device
+    pos_features = _encode_in_batches(sae, pos_acts, input_scale, device)
+    neg_features = _encode_in_batches(sae, neg_acts, input_scale, device)
+
+    pos_mean = pos_features.mean(dim=0)
+    neg_mean = neg_features.mean(dim=0)
+    pos_rate = (pos_features > 0).float().mean(dim=0)
+    neg_rate = (neg_features > 0).float().mean(dim=0)
+
+    pooled_std = torch.sqrt(
+        (pos_features.var(dim=0, unbiased=False) + neg_features.var(dim=0, unbiased=False)) / 2
+    ).clamp_min(1e-6)
+    diff = pos_mean - neg_mean
+    effect = diff / pooled_std
+
+    overall_rate = (pos_rate + neg_rate) / 2
+    eligible = (overall_rate >= min_fire_rate) & (overall_rate <= max_fire_rate)
+    effect = torch.where(eligible, effect, torch.zeros_like(effect))
+
+    order = torch.argsort(effect.abs(), descending=True)[:top_n]
+    return [
+        CandidateFeature(
+            feature_id=int(fid),
+            mean_diff=float(diff[fid].item()),
+            mean_positive=float(pos_mean[fid].item()),
+            mean_negative=float(neg_mean[fid].item()),
+            positive_fire_rate=float(pos_rate[fid].item()),
+            negative_fire_rate=float(neg_rate[fid].item()),
+            effect_size=float(effect[fid].item()),
+        )
+        for fid in order.tolist()
+        if effect[fid] != 0
+    ]
+
+
+def _encode_in_batches(
+    sae: TopKSAE, acts: torch.Tensor, input_scale: float, device: torch.device, batch: int = 4096
+) -> torch.Tensor:
+    chunks: list[torch.Tensor] = []
+    for start in range(0, acts.shape[0], batch):
+        x = acts[start : start + batch].to(device=device, dtype=torch.float32) * input_scale
+        sparse, _, _ = sae.encode(x)
+        chunks.append(sparse.cpu())
+    return torch.cat(chunks, dim=0)
+
+
+@dataclass
+class SweepPoint:
+    """One point of a strength sweep."""
+
+    strength: float
+    mean_divergence: float
+    mean_words: float
+    degeneration_rate: float
+    generations: list[str] = field(default_factory=list)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "strength": self.strength,
+            "mean_divergence": self.mean_divergence,
+            "mean_words": self.mean_words,
+            "degeneration_rate": self.degeneration_rate,
+        }
+
+
+def strength_sweep(
+    model: Any,
+    prompts: Sequence[str],
+    *,
+    feature_id: int,
+    strengths: Sequence[float],
+    generation: Any = None,
+) -> list[SweepPoint]:
+    """Sweep one feature's strength and measure divergence and degeneration.
+
+    The purpose is to find the window where an intervention changes behaviour
+    *before* it destroys fluency. Every sweep point costs
+    ``len(prompts)`` generations, so keep both lists short.
+    """
+    from brainpatch.evaluation.metrics import jaccard_similarity, score_generation
+    from brainpatch.research.ml.causal import _spec_for
+    from brainpatch.research.ml.generation import GenerationConfig
+
+    cfg = generation or GenerationConfig(max_new_tokens=96)
+    saved = dict(model.plan.patches)
+
+    model.plan.patches = {}
+    baselines = [model.generate(p, config=cfg) for p in prompts]
+
+    points: list[SweepPoint] = []
+    for strength in strengths:
+        divergences: list[float] = []
+        words: list[int] = []
+        degenerate = 0
+        texts: list[str] = []
+        for prompt, baseline in zip(prompts, baselines):
+            model.plan.patches = {}
+            model.install(_spec_for(model, feature_id, strength, f"sweep-{feature_id}"))
+            text = model.generate(prompt, config=cfg)
+            metrics = score_generation(text)
+            divergences.append(1.0 - jaccard_similarity(baseline, text, n=3))
+            words.append(metrics.num_words)
+            degenerate += int(metrics.degeneration_flag)
+            texts.append(text)
+        points.append(
+            SweepPoint(
+                strength=float(strength),
+                mean_divergence=sum(divergences) / len(divergences),
+                mean_words=sum(words) / len(words),
+                degeneration_rate=degenerate / len(prompts),
+                generations=texts,
+            )
+        )
+
+    model.plan.patches = saved
+    return points
+
+
+def greedy_feature_selection(
+    model: Any,
+    prompts: Sequence[str],
+    candidates: Sequence[CandidateFeature],
+    *,
+    strength: float,
+    max_features: int = 3,
+    generation: Any = None,
+    degeneration_ceiling: float = 0.25,
+) -> list[tuple[int, float]]:
+    """Greedily build a sparse multi-feature patch.
+
+    At each round, every remaining candidate is added tentatively and the one
+    producing the largest divergence-from-baseline *without* pushing the
+    degeneration rate above ``degeneration_ceiling`` is kept. Selection stops
+    when no candidate improves the objective.
+
+    Cost is ``O(max_features * len(candidates) * len(prompts))`` generations.
+    Keep the candidate list to single digits.
+    """
+    from brainpatch.evaluation.metrics import jaccard_similarity, score_generation
+    from brainpatch.research.ml.causal import _spec_for
+    from brainpatch.research.ml.generation import GenerationConfig
+    from brainpatch.schemas.patch import BrainPatchSpec, FeatureEdit
+
+    cfg = generation or GenerationConfig(max_new_tokens=96)
+    saved = dict(model.plan.patches)
+
+    model.plan.patches = {}
+    baselines = [model.generate(p, config=cfg) for p in prompts]
+
+    selected: list[tuple[int, float]] = []
+    remaining = [c.feature_id for c in candidates]
+    best_score = 0.0
+
+    def evaluate(edits: list[tuple[int, float]]) -> tuple[float, float]:
+        base_spec = _spec_for(model, edits[0][0], edits[0][1], "greedy")
+        spec = BrainPatchSpec(
+            name="greedy",
+            base_model=base_spec.base_model,
+            model_revision=base_spec.model_revision,
+            sae=base_spec.sae,
+            features=[FeatureEdit(feature_id=f, strength=s) for f, s in edits],
+            description="Transient greedy-search candidate.",
+        )
+        divergences: list[float] = []
+        degenerate = 0
+        for prompt, baseline in zip(prompts, baselines):
+            model.plan.patches = {}
+            model.install(spec)
+            text = model.generate(prompt, config=cfg)
+            divergences.append(1.0 - jaccard_similarity(baseline, text, n=3))
+            degenerate += int(score_generation(text).degeneration_flag)
+        return sum(divergences) / len(divergences), degenerate / len(prompts)
+
+    for _ in range(max_features):
+        best_feature: int | None = None
+        best_round = best_score
+        for feature_id in list(remaining):
+            score, degeneration = evaluate([*selected, (feature_id, strength)])
+            if degeneration > degeneration_ceiling:
+                continue
+            if score > best_round:
+                best_round = score
+                best_feature = feature_id
+        if best_feature is None:
+            break
+        selected.append((best_feature, strength))
+        remaining.remove(best_feature)
+        best_score = best_round
+
+    model.plan.patches = saved
+    return selected

brainpatch/research/ml/runtime.py

@@ -0,0 +1,343 @@
+"""``BrainPatchedModel`` -- the user-facing runtime.
+
+::
+
+    from brainpatch import BrainPatchedModel
+
+    model = BrainPatchedModel.from_pretrained("Qwen/Qwen2.5-1.5B-Instruct")
+    model.load_sae("/vol/sae/smoke_v0/sae_latest.pt", reference="smoke_v0")
+    model.install("patches/experimental-feature-1207.json")
+    model.set_patch_strength("experimental-feature-1207", 1.5)
+    print(model.generate("Solve this problem..."))
+
+The base model is frozen. Every behavioural change comes from a hook that adds
+a vector to one layer's residual stream, installed for the duration of a
+generation and removed afterwards.
+
+Two invariants this class exists to guarantee:
+
+1. **No patches installed, or all strengths zero, is exactly baseline.** With
+   nothing to apply, the hook returns ``None`` and the tensor is never touched.
+   :meth:`assert_zero_strength_is_baseline` verifies this empirically rather
+   than trusting the argument.
+
+2. **A patch cannot be applied to the wrong model.** Every install runs
+   :meth:`~brainpatch.schemas.patch.BrainPatchSpec.check_compatibility` against
+   the loaded weights and SAE.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any
+
+import torch
+
+from brainpatch.research.ml.generation import GenerationConfig, build_chat_prompt
+from brainpatch.research.ml.hooks import HookSet
+from brainpatch.research.ml.intervention import FeatureSteerer, make_steerer
+from brainpatch.research.ml.model import DEFAULT_MODEL, ModelBundle, load_model
+from brainpatch.research.ml.sae import TopKSAE
+from brainpatch.schemas.patch_io import load_patch
+from brainpatch.schemas.patch import BrainPatchSpec, FeatureEdit, SAEReference
+from brainpatch.steering.plan import InterventionPlan
+from brainpatch.steering.schedule import StrengthSchedule
+
+
+class BrainPatchedModel:
+    """A frozen language model with installable activation-space patches."""
+
+    def __init__(self, bundle: ModelBundle) -> None:
+        self.bundle = bundle
+        self.plan = InterventionPlan()
+        self.sae: TopKSAE | None = None
+        self.sae_reference: str | None = None
+        self.input_scale: float | None = None
+        self._last_stats: dict[str, Any] = {}
+        self._last_trace: list[tuple[int, float]] = []
+
+    # -- construction ----------------------------------------------------------
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_id: str = DEFAULT_MODEL,
+        *,
+        revision: str | None = None,
+        dtype: str = "bfloat16",
+        device: str = "cuda",
+    ) -> "BrainPatchedModel":
+        """Load a frozen base model. Downloads go to the HF cache, never locally."""
+        return cls(load_model(model_id, revision=revision, dtype=dtype, device=device))
+
+    def load_sae(
+        self,
+        checkpoint_path: str | os.PathLike[str],
+        *,
+        reference: str,
+        input_scale: float | None = None,
+    ) -> TopKSAE:
+        """Attach a trained SAE, whose decoder columns become patch directions.
+
+        Raises
+        ------
+        ValueError
+            If the SAE's input width does not match the model's residual width,
+            which would mean the two were never trained on the same activations.
+        """
+        path = Path(checkpoint_path)
+        if not path.is_file():
+            raise FileNotFoundError(f"SAE checkpoint not found: {path}")
+        checkpoint = torch.load(path, map_location="cpu", weights_only=False)
+        sae = TopKSAE.from_checkpoint(checkpoint, device=str(self.bundle.device))
+
+        if sae.d_in != self.bundle.hidden_size:
+            raise ValueError(
+                f"SAE was trained on width {sae.d_in} but the model's residual "
+                f"stream is {self.bundle.hidden_size} wide -- these are not the "
+                "same activations"
+            )
+
+        scale = input_scale if input_scale is not None else sae.config.input_scale
+        if scale is None:
+            raise ValueError(
+                "SAE checkpoint has no recorded input_scale, so a strength value "
+                "would have no defined magnitude. Pass input_scale explicitly."
+            )
+
+        self.sae = sae
+        self.sae_reference = reference
+        self.input_scale = float(scale)
+        return sae
+
+    # -- patch management ------------------------------------------------------
+
+    def install(
+        self,
+        patch: str | os.PathLike[str] | BrainPatchSpec,
+        *,
+        strength: float = 1.0,
+        strict_revision: bool = False,
+    ) -> BrainPatchSpec:
+        """Install a patch from a file path or an in-memory spec.
+
+        The compatibility check runs before anything is registered, so a
+        rejected patch leaves the runtime untouched.
+        """
+        spec = patch if isinstance(patch, BrainPatchSpec) else load_patch(patch)
+        spec.check_compatibility(
+            model=self.bundle.model_id,
+            hidden_size=self.bundle.hidden_size,
+            num_layers=self.bundle.num_layers,
+            model_revision=self.bundle.revision,
+            sae_reference=self.sae_reference,
+            sae_d_sae=self.sae.d_sae if self.sae is not None else None,
+            strict_revision=strict_revision,
+        )
+        self.plan.install(spec, strength=strength)
+        return spec
+
+    def uninstall(self, name: str) -> None:
+        self.plan.uninstall(name)
+
+    def list_patches(self) -> list[str]:
+        return list(self.plan.patches)
+
+    def set_patch_strength(self, name: str, strength: float) -> None:
+        """Change a patch's strength. Takes effect on the next generation."""
+        self.plan.set_strength(name, strength)
+
+    def set_patch_enabled(self, name: str, enabled: bool) -> None:
+        self.plan.set_enabled(name, enabled)
+
+    def set_patch_schedule(self, name: str, schedule: dict[int, float] | StrengthSchedule | None) -> None:
+        """Install a token-indexed strength schedule for dynamic steering."""
+        if isinstance(schedule, dict):
+            schedule = StrengthSchedule(schedule)
+        self.plan.set_schedule(name, schedule)
+
+    def add_feature(
+        self,
+        *,
+        layer: int,
+        feature_id: int,
+        strength: float,
+        name: str | None = None,
+        mode: str = "add",
+    ) -> BrainPatchSpec:
+        """Install an ad-hoc single-feature intervention.
+
+        The convenience path for exploration. It builds a real
+        :class:`BrainPatchSpec` under the hood, so an interactive experiment and
+        a shipped patch go through identical machinery -- and the resulting spec
+        can be saved directly.
+        """
+        if self.sae is None or self.sae_reference is None or self.input_scale is None:
+            raise RuntimeError("load_sae() must be called before adding features")
+        if layer != self.bundle_layer_for_sae():
+            # Not fatal, but the SAE only means anything at the layer it was fitted on.
+            raise ValueError(
+                f"SAE {self.sae_reference!r} was trained at layer "
+                f"{self.sae.config.layer}; refusing to inject its directions at layer {layer}"
+            )
+        spec = BrainPatchSpec(
+            name=name or f"adhoc-feature-{feature_id}",
+            base_model=self.bundle.model_id,
+            model_revision=self.bundle.revision,
+            sae=SAEReference(
+                reference=self.sae_reference,
+                layer=layer,
+                hook=self.sae.config.hook or "residual_post",
+                d_in=self.sae.d_in,
+                d_sae=self.sae.d_sae,
+                input_scale=self.input_scale,
+            ),
+            features=[FeatureEdit(feature_id=feature_id, strength=strength, mode=mode)],
+            description="Ad-hoc exploratory intervention. No evidence of any behavioural effect.",
+            evidence_level="none",
+        )
+        return self.install(spec)
+
+    def bundle_layer_for_sae(self) -> int:
+        """The layer the attached SAE was trained on."""
+        if self.sae is None:
+            raise RuntimeError("no SAE loaded")
+        return int(self.sae.config.layer)
+
+    # -- generation ------------------------------------------------------------
+
+    def generate(
+        self,
+        prompt: str,
+        *,
+        config: GenerationConfig | None = None,
+        system: str | None = None,
+        use_chat_template: bool = True,
+        control: str = "none",
+        control_seed: int = 1234,
+        apply_to_prompt: bool = True,
+    ) -> str:
+        """Generate a completion with all installed patches active.
+
+        Parameters
+        ----------
+        control:
+            ``"random"`` swaps every feature direction for a scale-matched
+            random one, keeping all else identical. This is the control
+            condition, run through exactly the same code path as the real
+            intervention.
+        """
+        cfg = config or GenerationConfig()
+        text = build_chat_prompt(self.bundle.tokenizer, prompt, system) if use_chat_template else prompt
+        inputs = self.bundle.tokenizer(text, return_tensors="pt").to(self.bundle.device)
+
+        if cfg.do_sample:
+            torch.manual_seed(cfg.seed)
+
+        steerer = self._build_steerer(control=control, control_seed=control_seed,
+                                      apply_to_prompt=apply_to_prompt)
+
+        with HookSet() as hooks:
+            if steerer is not None:
+                steerer.reset()
+                hooks.add(steerer.make_injector(), self.bundle.layer_module(steerer.layer))
+            with torch.inference_mode():
+                output = self.bundle.model.generate(
+                    **inputs, **cfg.to_kwargs(self.bundle.tokenizer)
+                )
+
+        self._last_stats = steerer.stats.to_dict() if steerer is not None else {}
+        self._last_trace = list(steerer.stats.per_token_norms) if steerer is not None else []
+        generated = output[0, inputs["input_ids"].shape[1] :]
+        return self.bundle.tokenizer.decode(generated, skip_special_tokens=True)
+
+    @property
+    def last_steering_stats(self) -> dict[str, Any]:
+        """Delta norms and application counts from the most recent generation.
+
+        The empirical answer to "did the intervention actually fire?" -- an
+        ``applied_passes`` of 0 means the plan resolved to nothing.
+        """
+        return dict(self._last_stats)
+
+    @property
+    def last_steering_trace(self) -> list[tuple[int, float]]:
+        """``(generated_token_index, delta_norm)`` for every pass of the last run.
+
+        The evidence that a dynamic schedule fired where it was supposed to:
+        the norm should be zero before the keyframe and non-zero after.
+        """
+        return list(self._last_trace)
+
+    def _build_steerer(
+        self, *, control: str, control_seed: int, apply_to_prompt: bool
+    ) -> FeatureSteerer | None:
+        """Construct a steerer, or ``None`` when there is nothing to apply."""
+        layers = self.plan.layers()
+        if not layers:
+            return None
+        if self.sae is None or self.input_scale is None:
+            raise RuntimeError("patches are installed but no SAE is loaded")
+        if len(layers) > 1:
+            raise NotImplementedError(
+                f"patches span layers {layers}; multi-layer steering needs one SAE "
+                "per layer and is not supported in v0"
+            )
+        return make_steerer(
+            self.sae,
+            self.plan,
+            layer=layers[0],
+            input_scale=self.input_scale,
+            device=self.bundle.device,
+            control=control,
+            control_seed=control_seed,
+            apply_to_prompt=apply_to_prompt,
+        )
+
+    # -- verification ----------------------------------------------------------
+
+    def assert_zero_strength_is_baseline(
+        self, prompt: str, *, config: GenerationConfig | None = None
+    ) -> dict[str, Any]:
+        """Empirically verify that zeroed patches reproduce baseline exactly.
+
+        Generates with every patch uninstalled, then with them installed at
+        strength 0, and compares the strings. This is a correctness test of the
+        hook machinery: if it ever fails, every "baseline" in every experiment
+        is suspect.
+        """
+        cfg = config or GenerationConfig(max_new_tokens=48)
+        saved = {name: p.strength for name, p in self.plan.patches.items()}
+
+        installed = dict(self.plan.patches)
+        self.plan.patches = {}
+        baseline = self.generate(prompt, config=cfg)
+
+        self.plan.patches = installed
+        for name in self.plan.patches:
+            self.plan.set_strength(name, 0.0)
+        zeroed = self.generate(prompt, config=cfg)
+        zero_stats = self.last_steering_stats
+
+        for name, strength in saved.items():
+            self.plan.set_strength(name, strength)
+
+        return {
+            "identical": baseline == zeroed,
+            "baseline": baseline,
+            "zero_strength": zeroed,
+            "applied_passes_at_zero": zero_stats.get("applied_passes", 0),
+            "num_patches": len(saved),
+        }
+
+    def describe(self) -> dict[str, Any]:
+        return {
+            **self.bundle.describe(),
+            "sae_reference": self.sae_reference,
+            "sae_d_sae": self.sae.d_sae if self.sae else None,
+            "sae_k": self.sae.k if self.sae else None,
+            "sae_layer": self.sae.config.layer if self.sae else None,
+            "input_scale": self.input_scale,
+            "installed_patches": self.plan.describe(),
+        }