open-interpretability 0.0.20__py3-none-any.whl

openinterp/pole_effects.py

@@ -0,0 +1,102 @@
+"""Each MLP pole's effect on the loss, estimated by attribution patching (issue #3486).
+
+Mean-ablating one pole replaces its part of the activation with that part's mean over the
+fit windows. The part is the positive part or the negative part. This is the ablation that
+``openinterp/coverage.py`` applies to every pole at once. Measuring it pole by pole costs
+one forward pass per pole: 73,728 for GPT-2 small. To first order, the change in mean
+loss is the sum over tokens of ``dL/dh × (mean − part)``. One backward pass per batch
+gives it for every pole of every layer. :func:`exact_effect` measures single poles, to
+check the estimate.
+
+The planner draws some explore units in proportion to this effect (``planner_effects``).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+
+from openinterp.coverage import mean_loss, pole_means
+from openinterp.harness.attribution import _corrupted_grads
+from openinterp.harness.hooks import NeuronHookRunner
+
+POLES = ("positive", "negative")
+
+
+def _loss_sum(ids: torch.Tensor):
+    """A metric for ``_corrupted_grads``: the summed next-token cross-entropy on ``ids``."""
+
+    def metric(logits: torch.Tensor) -> torch.Tensor:
+        target = ids[:, 1:].to(logits.device).flatten()
+        return torch.nn.functional.cross_entropy(
+            logits[:, :-1].flatten(0, 1).float(), target, reduction="sum"
+        )
+
+    return metric
+
+
+def estimate(
+    runner: NeuronHookRunner,
+    fit: torch.Tensor,
+    test: torch.Tensor,
+    layers: Sequence[int],
+    batch: int,
+) -> dict[tuple[int, str], torch.Tensor]:
+    """``{(layer, pole): tensor[width]}``: each pole's estimated rise in mean loss on ``test``."""
+    means = {layer: pole_means(runner, fit, layer, batch) for layer in layers}
+    totals: dict[tuple[int, str], torch.Tensor] = {}
+    for chunk in test.split(batch):
+        acts, grads = _corrupted_grads(runner, chunk, list(layers), _loss_sum(chunk))
+        for layer, act, grad in zip(layers, acts, grads, strict=True):
+            act, grad = act.detach().float(), grad.float()
+            m_pos, m_neg = (m.to(act.device) for m in means[layer])
+            parts = {"positive": (m_pos - act.clamp(min=0)), "negative": (m_neg - act.clamp(max=0))}
+            for pole, diff in parts.items():
+                value = (diff * grad).sum(dim=(0, 1)).cpu()
+                key = (layer, pole)
+                totals[key] = value if key not in totals else totals[key] + value
+    tokens = test[:, 1:].numel()
+    return {key: value / tokens for key, value in totals.items()}
+
+
+def single_pole(means, idx: int, pole: str):
+    """An intervention that mean-ablates only pole ``pole`` of neuron ``idx``."""
+    m_pos, m_neg = means
+
+    def intervene(act: torch.Tensor) -> torch.Tensor:
+        out = act.clone()
+        column = act[..., idx]
+        if pole == "positive":
+            out[..., idx] = m_pos[idx].to(act.device, act.dtype) + column.clamp(max=0)
+        else:
+            out[..., idx] = column.clamp(min=0) + m_neg[idx].to(act.device, act.dtype)
+        return out
+
+    return intervene
+
+
+def exact_effect(
+    runner: NeuronHookRunner,
+    fit: torch.Tensor,
+    test: torch.Tensor,
+    layer: int,
+    idx: int,
+    pole: str,
+    clean: float,
+    batch: int,
+) -> float:
+    """The measured rise in mean loss on ``test`` when one pole is mean-ablated."""
+    means = pole_means(runner, fit, layer, batch)
+    return mean_loss(runner, test, batch, {layer: single_pole(means, idx, pole)}) - clean
+
+
+def top_poles(effects: dict[tuple[int, str], torch.Tensor], k: int) -> list[dict]:
+    """The ``k`` poles with the largest estimated effect, largest first."""
+    rows = [
+        (float(v), layer, idx, pole)
+        for (layer, pole), values in effects.items()
+        for idx, v in enumerate(values.tolist())
+    ]
+    rows.sort(reverse=True)
+    return [{"layer": la, "neuron": i, "pole": p, "effect": v} for v, la, i, p in rows[:k]]

openinterp/prior_art.py

@@ -0,0 +1,143 @@
+"""Find prior art for a unit after the proposer answers (`open-interpretability audit`).
+
+A dictionary latent that Neuronpedia hosts is looked up directly. A neuron pole is
+aligned with each dictionary preset that Neuronpedia hosts at its layer
+(:func:`openinterp.harness.crossref.align_layer`). Rows from ``open-interpretability crossref`` are
+used where they exist. The latents above chance on each side are then looked up, the best
+``TOP`` first.
+
+Lookups go through the cached, polite fetch of :mod:`openinterp.crossref`. A failed
+lookup or alignment is left out, so prior art never stops an audit.
+``$OPENINTERP_PRIOR_ART=0`` turns the readout off.
+"""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Callable
+
+from openinterp.core.fence import fence
+from openinterp.core.prior_art import PriorArt
+
+TOP = 3  # latents per side of a pole
+ENV = "OPENINTERP_PRIOR_ART"
+
+
+def enabled() -> bool:
+    """False when ``$OPENINTERP_PRIOR_ART`` turns the readout off."""
+    return os.environ.get(ENV, "1").lower() not in ("", "0", "off", "false")
+
+
+def find_prior_art(submission, runner, *, rows: dict | None = None, fetch: Callable | None = None):
+    """Return the unit's prior art; ``None`` for a unit kind with none, or when turned off.
+
+    ``rows`` are crossref rows by slug (:func:`openinterp.crossref.load_rows`). ``fetch``
+    returns a latent's explanations (default :func:`openinterp.crossref.latent_explanations`).
+    """
+    from openinterp.crossref import latent_explanations
+
+    if not enabled():
+        return None
+    fetch = fetch or latent_explanations
+    ident = submission.identifier
+    if ident.kind == "dictionary_feature":
+        entry = _entry(fetch, ident.dictionary, ident.layer, ident.feature_idx, "self", None)
+        return [entry] if entry else []
+    if ident.kind != "mlp_neuron_pole":
+        return None
+    found = []
+    for name, side, matches in _aligned(ident, runner, rows):
+        for feature_idx, cosine in matches:
+            entry = _entry(fetch, name, ident.layer, feature_idx, side, cosine)
+            if entry:
+                found.append(entry)
+    return found
+
+
+def _entry(fetch, dictionary: str, layer: int, feature_idx: int, relation: str, cosine):
+    """The latent's first machine-written explanation, else its first; ``None`` if none."""
+    from openinterp.crossref import latent_url, machine_explanation
+
+    url = latent_url(dictionary, layer, feature_idx)
+    if url is None:
+        return None
+    explanations = fetch(dictionary, layer, feature_idx) or []
+    best = machine_explanation(explanations) or next(iter(explanations), None)
+    if best is None:
+        return None
+    return PriorArt(
+        dictionary=dictionary,
+        feature_idx=feature_idx,
+        relation=relation,
+        cosine=cosine,
+        description=best.description,
+        explainer=best.explainer,
+        url=url,
+    )
+
+
+def _aligned(ident, runner, rows: dict | None):
+    """Yield ``(dictionary, side, [(feature_idx, cosine)])`` for each hosted dictionary."""
+    from openinterp.crossref import dictionaries_for, load_rows
+    from openinterp.harness.dictionaries import DICTIONARIES
+
+    rows = load_rows() if rows is None else rows
+    known = {r.dictionary: r for r in rows.get(ident.slug, [])}
+    for name in dictionaries_for(ident.model_name, ident.layer):
+        if not DICTIONARIES[name].neuronpedia_source(ident.layer):
+            continue
+        sides = _sides(known[name]) if name in known else _align(runner, name, ident)
+        for side, (matches, chance) in sides.items():
+            above = [m for m in matches if abs(m[1]) > chance]
+            yield name, side, sorted(above, key=lambda m: -abs(m[1]))[:TOP]
+
+
+def _sides(row) -> dict:
+    """``{side: ([(feature_idx, cosine)], chance)}`` of a stored crossref row."""
+    out = {}
+    for side in ("writes", "reads"):
+        s = getattr(row, side)
+        if s is not None:
+            out[side] = ([(m.feature_idx, m.cosine) for m in s.matches], s.chance)
+    return out
+
+
+def _align(runner, name: str, ident) -> dict:
+    """Align the pole with dictionary ``name`` now; empty if that fails."""
+    from openinterp.core.exceptions import ArchitectureNotSupportedError
+    from openinterp.harness.crossref import align_layer
+    from openinterp.harness.dictionaries import dictionary_spec
+
+    try:
+        found = align_layer(
+            runner, dictionary_spec(name), ident.layer, [(ident.neuron_idx, ident.pole)]
+        )
+    except (ArchitectureNotSupportedError, OSError, ValueError):
+        return {}
+    write, read = found[(ident.neuron_idx, ident.pole)]
+    out = {"writes": ([(m.feature_idx, m.cosine) for m in write.matches], write.chance)}
+    if read is not None:
+        out["reads"] = ([(m.feature_idx, m.cosine) for m in read.matches], read.chance)
+    return out
+
+
+def format_for_reviewer(prior_art: list[PriorArt] | None) -> str:
+    """The prior art as a block of the reviewer's prompt; empty if there is none."""
+    if not prior_art:
+        return ""
+    lines = []
+    for p in prior_art:
+        where = (
+            "this latent itself"
+            if p.relation == "self"
+            else f"{p.dictionary} latent {p.feature_idx}, which {p.relation} along this pole "
+            f"(cosine {p.cosine:+.2f})"
+        )
+        lines.append(f"- {where}: {p.description} [{p.source}, {p.explainer}]")
+    return (
+        "Prior art: published explanations of this unit or of latents aligned with it. They "
+        "are other people's claims, not measurements, and the proposer did not see them. "
+        "Where one disagrees with the hypothesis, write probes that tell them apart.\n"
+        + fence("\n".join(lines))
+        + "\n"
+    )