archscope 0.2.2__py3-none-any.whl

archscope/neurons.py

@@ -0,0 +1,118 @@
+"""Targeted Neuron Modulation via Contrastive Pair Search (2605.12290).
+
+5-step algorithm reimplemented from paper:
+1. Run harmful + benign prompts through model
+2. Record MLP activations at final token position
+3. Compute per-neuron mean activation difference
+4. Select top k (default 0.1%) neurons by |Δ|
+5. At inference: multiply selected neuron activations by scalar m
+"""
+from __future__ import annotations
+from dataclasses import dataclass
+import torch
+from .backends import Backend
+from ._utils import resolve_layer_module
+
+
+@dataclass
+class NeuronEditConfig:
+    top_frac: float = 0.001         # top 0.1% by default
+    layer_filter: str | None = None # e.g., "mlp" to restrict to MLP neurons
+    mode: str = "scalar"            # "scalar" (multiply by m) or "ablate" (m=0)
+
+
+@dataclass
+class NeuronEdit:
+    """Result of contrastive search — which neurons to edit + how."""
+    layer_to_indices: dict[str, torch.Tensor]   # layer_name → neuron indices
+    layer_to_deltas: dict[str, torch.Tensor]    # layer_name → mean diff values
+    config: NeuronEditConfig
+    multiplier: float = 0.0   # 0 = ablate; >1 = amplify; <1 = dampen
+
+    def apply_hook(self, model, backend: Backend | None = None):
+        """Register forward hooks that modulate the selected neurons.
+
+        Returns a context manager that auto-removes hooks on exit.
+        """
+        hooks = []
+        for layer_name, indices in self.layer_to_indices.items():
+            module = resolve_layer_module(model, layer_name)
+            if module is None:
+                continue
+
+            indices_local = indices
+
+            def hook(module, input, output, idxs=indices_local, m=self.multiplier):
+                # Return MODIFIED output (don't rely on in-place — some HF layers
+                # produce fresh tensors that won't propagate in-place changes).
+                if isinstance(output, tuple):
+                    h = output[0].clone()
+                    if h.dim() >= 2:
+                        h[..., idxs] = h[..., idxs] * m
+                    return (h,) + output[1:]
+                h = output.clone()
+                if h.dim() >= 2:
+                    h[..., idxs] = h[..., idxs] * m
+                return h
+            hooks.append(module.register_forward_hook(hook))
+
+        return _HookContext(hooks)
+
+
+class _HookContext:
+    """Auto-removes forward hooks when used as a `with` block."""
+
+    def __init__(self, hooks):
+        self.hooks = hooks
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        for h in self.hooks:
+            h.remove()
+
+
+def find_neurons(
+    model,
+    inputs_harmful: list,
+    inputs_benign: list,
+    config: NeuronEditConfig | None = None,
+    backend_hint: str | None = None,
+) -> NeuronEdit:
+    """Algorithm 1 of Targeted Neuron Modulation paper.
+
+    Returns a NeuronEdit that can be applied as a hook during inference.
+    """
+    config = config or NeuronEditConfig()
+    backend = Backend.for_model(model, hint=backend_hint)
+
+    # Get all layers (will filter to MLP later if requested)
+    all_layers = backend.layer_names()
+
+    # Forward both classes, collect final-token activations
+    harm_acts = backend.extract(inputs_harmful, layers=all_layers)
+    ben_acts = backend.extract(inputs_benign, layers=all_layers)
+
+    layer_to_indices = {}
+    layer_to_deltas = {}
+    for h_rec, b_rec in zip(harm_acts, ben_acts):
+        # Final token: (batch, -1, hidden_dim) -> (batch, hidden_dim)
+        h_final = h_rec.activations[:, -1, :]
+        b_final = b_rec.activations[:, -1, :]
+        # Per-neuron mean diff
+        delta = h_final.mean(dim=0) - b_final.mean(dim=0)
+        # Top k by absolute value
+        k = max(1, int(config.top_frac * len(delta)))
+        topk = torch.topk(delta.abs(), k=k)
+        layer_to_indices[h_rec.layer_name] = topk.indices
+        layer_to_deltas[h_rec.layer_name] = delta[topk.indices]
+
+    return NeuronEdit(
+        layer_to_indices=layer_to_indices,
+        layer_to_deltas=layer_to_deltas,
+        config=config,
+        multiplier=0.0,
+    )
+
+

archscope/probes.py

@@ -0,0 +1,160 @@
+"""Probe-Filtered RL — linear/MLP probes over hidden states (Drop the Act, 2605.11467).
+
+Core idea: train a frozen-base probe over residual/hidden states that predicts
+some property (faithfulness, refusal, deception). Use probe scores to filter
+trajectories during RL training, or as inference-time signals.
+
+Reimplemented from paper (code anonymized for review).
+"""
+from __future__ import annotations
+from dataclasses import dataclass
+import torch
+import torch.nn as nn
+from .backends import Backend
+
+
+@dataclass
+class ProbeConfig:
+    """Configuration for a single probe."""
+    layer_name: str
+    probe_type: str = "linear"   # "linear" or "mlp"
+    hidden_dim: int = 64         # only used if mlp
+    target: str = "performativity"
+
+
+class Probe(nn.Module):
+    """Linear or MLP probe over a layer's hidden states.
+
+    Trained to predict a scalar property from activations at one layer.
+    Frozen base model: probe is the only trainable component.
+    """
+
+    def __init__(self, input_dim: int, config: ProbeConfig):
+        super().__init__()
+        self.config = config
+        if config.probe_type == "linear":
+            self.net = nn.Linear(input_dim, 1)
+        elif config.probe_type == "mlp":
+            self.net = nn.Sequential(
+                nn.Linear(input_dim, config.hidden_dim),
+                nn.GELU(),
+                nn.Linear(config.hidden_dim, 1),
+            )
+        else:
+            raise ValueError(f"Unknown probe_type: {config.probe_type}")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply probe. Accepts (N, hidden_dim) or (N, seq, hidden_dim); returns
+        same leading dims with the final hidden_dim collapsed to scalar logits."""
+        return self.net(x).squeeze(-1)
+
+
+class ProbeFit:
+    """Fit a probe given (activations, labels) pairs."""
+
+    def __init__(self, config: ProbeConfig, input_dim: int, device: str = "cpu"):
+        self.config = config
+        self.probe = Probe(input_dim, config).to(device)
+        self.device = device
+
+    def train(
+        self,
+        activations: torch.Tensor,   # (N, hidden_dim) or (N, seq_len, hidden_dim) pooled
+        labels: torch.Tensor,         # (N,)
+        epochs: int = 50,
+        lr: float = 1e-3,
+        batch_size: int = 64,
+        val_split: float = 0.2,
+    ) -> dict:
+        """Standard supervised fit. Returns train/val AUROC + final loss."""
+        if activations.dim() == 3:
+            # Pool seq dim by mean (could be configurable)
+            activations = activations.mean(dim=1)
+        activations = activations.to(self.device)
+        labels = labels.float().to(self.device)
+
+        n = len(activations)
+        idx = torch.randperm(n)
+        n_val = int(n * val_split)
+        train_idx, val_idx = idx[n_val:], idx[:n_val]
+
+        opt = torch.optim.AdamW(self.probe.parameters(), lr=lr)
+        loss_fn = nn.BCEWithLogitsLoss()
+
+        for _ in range(epochs):
+            self.probe.train()
+            perm = train_idx[torch.randperm(len(train_idx))]
+            for b in range(0, len(perm), batch_size):
+                batch = perm[b:b+batch_size]
+                logits = self.probe(activations[batch])
+                loss = loss_fn(logits, labels[batch])
+                opt.zero_grad()
+                loss.backward()
+                opt.step()
+
+        self.probe.eval()
+        with torch.no_grad():
+            train_logits = self.probe(activations[train_idx])
+            val_logits = self.probe(activations[val_idx])
+        return {
+            "train_auroc": _auroc(train_logits, labels[train_idx]),
+            "val_auroc": _auroc(val_logits, labels[val_idx]),
+            "train_loss": loss_fn(train_logits, labels[train_idx]).item(),
+        }
+
+    def score(self, activations: torch.Tensor) -> torch.Tensor:
+        """Apply probe — returns per-token (or per-example if pooled) scores."""
+        self.probe.eval()
+        with torch.no_grad():
+            return torch.sigmoid(self.probe(activations.to(self.device)))
+
+
+def _auroc(logits: torch.Tensor, labels: torch.Tensor) -> float:
+    """Simple AUROC from logits + binary labels."""
+    from sklearn.metrics import roc_auc_score
+    scores = torch.sigmoid(logits).cpu().numpy()
+    y = labels.cpu().numpy()
+    try:
+        return float(roc_auc_score(y, scores))
+    except ValueError:
+        return float("nan")   # happens when only one class present
+
+
+# High-level API matching paper
+
+def fit_probe(
+    model,
+    inputs_pos: list,    # examples where target=1 (e.g., faithful)
+    inputs_neg: list,    # examples where target=0 (e.g., reasoning theater)
+    layer_name: str,
+    backend_hint: str | None = None,
+    config: ProbeConfig | None = None,
+    device: str = "cpu",
+) -> ProbeFit:
+    """End-to-end: extract activations from model, fit probe."""
+    backend = Backend.for_model(model, hint=backend_hint)
+    config = config or ProbeConfig(layer_name=layer_name)
+
+    # Extract activations for both classes — detach from autograd graph
+    # (probe is trained separately; base model is frozen)
+    with torch.no_grad():
+        acts_pos = backend.extract(inputs_pos, layers=[layer_name])[0].activations.detach()
+        acts_neg = backend.extract(inputs_neg, layers=[layer_name])[0].activations.detach()
+
+    # Pool across sequence dim — handles independent padding across batches
+    if acts_pos.dim() == 3:
+        acts_pos = acts_pos.mean(dim=1)   # (batch, hidden)
+    if acts_neg.dim() == 3:
+        acts_neg = acts_neg.mean(dim=1)
+
+    labels_pos = torch.ones(len(acts_pos))
+    labels_neg = torch.zeros(len(acts_neg))
+
+    all_acts = torch.cat([acts_pos, acts_neg], dim=0)
+    all_labels = torch.cat([labels_pos, labels_neg])
+
+    hidden_dim = backend.hidden_dim(layer_name)
+    probe_fit = ProbeFit(config, hidden_dim, device=device)
+    metrics = probe_fit.train(all_acts, all_labels)
+    probe_fit.metrics = metrics
+    return probe_fit

archscope/sae.py

@@ -0,0 +1,127 @@
+"""Sparse Autoencoders for hidden states (WriteSAE, 2605.12770).
+
+Two SAE variants:
+- DenseSAE: standard SAE with L1 sparsity penalty
+- Rank1FactoredSAE: WriteSAE's contribution — atoms = v_i w_i^T (rank-1 outer
+  product), applied to recurrent cache writes specifically.
+
+Both work on transformer residual streams AND recurrent hidden states.
+"""
+from __future__ import annotations
+from dataclasses import dataclass
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@dataclass
+class SAEConfig:
+    input_dim: int
+    n_features: int             # dictionary size
+    sparsity: float = 1e-3      # L1 coefficient
+    sae_type: str = "dense"     # "dense" or "rank1"
+    learning_rate: float = 1e-3
+    batch_size: int = 64
+
+
+class DenseSAE(nn.Module):
+    """Standard SAE: encoder + decoder, L1 sparsity on hidden code."""
+
+    def __init__(self, config: SAEConfig):
+        super().__init__()
+        self.config = config
+        self.encoder = nn.Linear(config.input_dim, config.n_features)
+        self.decoder = nn.Linear(config.n_features, config.input_dim, bias=False)
+        # Initialize decoder weights to encoder.T transpose (helpful for SAEs)
+        with torch.no_grad():
+            self.decoder.weight.data = self.encoder.weight.data.T.clone()
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        return F.relu(self.encoder(x))
+
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        return self.decoder(z)
+
+    def forward(self, x: torch.Tensor):
+        z = self.encode(x)
+        x_hat = self.decode(z)
+        return x_hat, z
+
+    def loss(self, x: torch.Tensor):
+        x_hat, z = self.forward(x)
+        recon = F.mse_loss(x_hat, x)
+        l1 = z.abs().mean() * self.config.sparsity
+        return recon + l1, {"recon": recon.item(), "l1": l1.item(), "n_active": (z > 0).float().mean().item()}
+
+
+class Rank1FactoredSAE(nn.Module):
+    """WriteSAE rank-1 factored atoms: each feature i = v_i w_i^T outer product.
+
+    Designed for recurrent CACHE WRITES specifically — atoms substitute for
+    native write contributions at matched Frobenius norm.
+
+    Reference: Eq. 3-factor closed form Δℓ ≈ G·⟨w_i,q_t⟩·⟨v_i,W_u[tok]⟩
+    """
+
+    def __init__(self, config: SAEConfig):
+        super().__init__()
+        assert config.sae_type == "rank1"
+        self.config = config
+        # Atoms parameterized as v (output) and w (input) vectors
+        self.v = nn.Parameter(torch.randn(config.n_features, config.input_dim) * 0.02)
+        self.w = nn.Parameter(torch.randn(config.n_features, config.input_dim) * 0.02)
+
+    def fire(self, x: torch.Tensor) -> torch.Tensor:
+        """Firing strengths per atom: <w_i, x>, kept positive."""
+        # x: (..., input_dim), w: (n_features, input_dim)
+        return F.relu(x @ self.w.T)
+
+    def reconstruct(self, x: torch.Tensor) -> torch.Tensor:
+        fires = self.fire(x)
+        # Output: sum_i fires_i * v_i
+        return fires @ self.v
+
+    def forward(self, x: torch.Tensor):
+        x_hat = self.reconstruct(x)
+        z = self.fire(x)
+        return x_hat, z
+
+    def loss(self, x: torch.Tensor):
+        x_hat, z = self.forward(x)
+        recon = F.mse_loss(x_hat, x)
+        l1 = z.abs().mean() * self.config.sparsity
+        return recon + l1, {"recon": recon.item(), "l1": l1.item(), "n_active": (z > 0).float().mean().item()}
+
+
+def build_sae(config: SAEConfig):
+    if config.sae_type == "dense":
+        return DenseSAE(config)
+    if config.sae_type == "rank1":
+        return Rank1FactoredSAE(config)
+    raise ValueError(f"Unknown sae_type: {config.sae_type}")
+
+
+def fit_sae(
+    activations: torch.Tensor,   # (N, input_dim) flattened
+    config: SAEConfig,
+    epochs: int = 100,
+    device: str = "cpu",
+) -> nn.Module:
+    """Train SAE on flattened activations."""
+    sae = build_sae(config).to(device)
+    opt = torch.optim.AdamW(sae.parameters(), lr=config.learning_rate)
+    activations = activations.to(device)
+
+    n = len(activations)
+    last_metrics: dict = {}
+    for _ in range(epochs):
+        perm = torch.randperm(n)
+        for b in range(0, n, config.batch_size):
+            batch = activations[perm[b:b + config.batch_size]]
+            loss, metrics = sae.loss(batch)
+            opt.zero_grad()
+            loss.backward()
+            opt.step()
+            last_metrics = metrics
+    sae.last_metrics = last_metrics
+    return sae

archscope/transfer.py

@@ -0,0 +1,188 @@
+"""Cross-architecture probe transfer via paired-activation linear alignment.
+
+Setup:
+- Two models with different hidden dims (e.g., Pythia 768 vs Kazdov 512)
+- Train a linear projection M: kazdov_space (512) → pythia_space (768) using
+  paired (text → activation pair) data
+- A Pythia probe with weights w_py ∈ R^768 transferred to kazdov-space becomes
+  w_kz = M^T · w_py ∈ R^512
+
+This tests: do probe directions from one architecture transfer to another?
+"""
+from __future__ import annotations
+import torch
+import torch.nn as nn
+from dataclasses import dataclass
+
+
+@dataclass
+class TransferResult:
+    """One transfer experiment's metrics."""
+    source_arch: str           # e.g., "pythia"
+    target_arch: str           # e.g., "kazdov"
+    source_layer: str
+    target_layer: str
+    n_align_pairs: int
+
+    baseline_source_auroc: float  # source probe on source data (no transfer)
+    baseline_target_auroc: float  # target probe on target data (in-arch reference)
+    transfer_auroc: float         # source probe via M applied to target data
+    transfer_drop: float           # baseline_target - transfer (how much we lose)
+
+
+def learn_alignment(
+    src_acts: torch.Tensor,    # (N, d_src)
+    tgt_acts: torch.Tensor,    # (N, d_tgt)
+    ridge: float = 1e-3,
+) -> torch.Tensor:
+    """Ridge regression: learn M such that M @ tgt_acts.T ≈ src_acts.T.
+
+    M has shape (d_src, d_tgt). After fit, src_act ≈ M @ tgt_act.
+
+    Ridge term prevents overfitting on small N.
+    """
+    assert src_acts.shape[0] == tgt_acts.shape[0]
+    src_acts = src_acts.float()
+    tgt_acts = tgt_acts.float()
+    # M = src @ tgt.T @ (tgt @ tgt.T + λI)^-1
+    # But the standard formulation: we want X @ M.T = Y where X is target, Y is source
+    # So M.T = (X.T X + λI)^-1 X.T Y
+    # Then M = ((X.T X + λI)^-1 X.T Y).T
+    X = tgt_acts  # (N, d_tgt)
+    Y = src_acts  # (N, d_src)
+    XtX = X.T @ X
+    reg = torch.eye(XtX.shape[0], device=XtX.device) * ridge
+    M_T = torch.linalg.solve(XtX + reg, X.T @ Y)
+    M = M_T.T  # (d_src, d_tgt)
+    return M
+
+
+def transfer_probe(
+    probe_weights_source: torch.Tensor,  # (d_src,) — Pythia probe direction
+    probe_bias_source: torch.Tensor,     # scalar
+    alignment: torch.Tensor,             # M: (d_src, d_tgt)
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Transport a probe from source arch to target arch via alignment matrix.
+
+    Given Pythia probe f(x_py) = w_py · x_py + b, and alignment x_py ≈ M @ x_kz,
+    the kazdov-space probe is w_kz = M.T @ w_py, with same bias.
+    """
+    w_target = alignment.T @ probe_weights_source
+    return w_target, probe_bias_source
+
+
+def auroc_from_scores(scores: torch.Tensor, labels: torch.Tensor) -> float:
+    """Quick AUROC."""
+    from sklearn.metrics import roc_auc_score
+    try:
+        return float(roc_auc_score(labels.cpu().numpy(), scores.cpu().numpy()))
+    except ValueError:
+        return float("nan")
+
+
+def evaluate_transfer(
+    source_model, target_model,
+    source_backend, target_backend,
+    source_tokenize, target_tokenize,
+    align_texts: list[str],      # texts for learning M (no labels needed)
+    train_pos: list[str], train_neg: list[str],
+    test_pos: list[str], test_neg: list[str],
+    source_layer: str,
+    target_layer: str,
+    source_arch_name: str = "source",
+    target_arch_name: str = "target",
+) -> TransferResult:
+    """Full transfer experiment.
+
+    1. Get paired source/target activations on `align_texts` → learn alignment M
+    2. Train source probe on source activations of train_pos/neg
+    3. Train target probe on target activations of train_pos/neg (in-arch baseline)
+    4. Transfer source probe via M, apply to target activations of test_pos/neg
+    5. Compare AUROCs.
+    """
+    def _extract_pooled(backend, tokenize, texts: list[str], layer: str) -> torch.Tensor:
+        """Extract activations and pool to (N, hidden_dim) by averaging seq dim."""
+        with torch.no_grad():
+            rec = backend.extract(tokenize(texts), layers=[layer])[0]
+            acts = rec.activations.detach()
+        return acts.mean(dim=1) if acts.dim() == 3 else acts
+
+    # -- 1. Learn alignment from paired activations on the same texts.
+    src_align = _extract_pooled(source_backend, source_tokenize, align_texts, source_layer)
+    tgt_align = _extract_pooled(target_backend, target_tokenize, align_texts, target_layer)
+    M = learn_alignment(src_align, tgt_align)
+
+    # -- 2. Train source probe (Pythia-style) on source activations
+    from . import probes
+    src_train_inputs_pos = source_tokenize(train_pos)
+    src_train_inputs_neg = source_tokenize(train_neg)
+    pf_src = probes.fit_probe(
+        source_model, src_train_inputs_pos, src_train_inputs_neg,
+        layer_name=source_layer, backend_hint=source_arch_name,
+    )
+
+    # Extract source probe weights (linear case)
+    src_probe_net = pf_src.probe.net
+    if not isinstance(src_probe_net, nn.Linear):
+        raise NotImplementedError("transfer only works for linear probes for now")
+    w_src = src_probe_net.weight.data.squeeze(0)
+    b_src = src_probe_net.bias.data
+
+    # -- 3. Train target probe (in-arch baseline)
+    tgt_train_inputs_pos = target_tokenize(train_pos)
+    tgt_train_inputs_neg = target_tokenize(train_neg)
+    pf_tgt = probes.fit_probe(
+        target_model, tgt_train_inputs_pos, tgt_train_inputs_neg,
+        layer_name=target_layer, backend_hint=target_arch_name,
+    )
+
+    # -- 4. Transfer: project w_src into target space via M
+    w_transferred, b_transferred = transfer_probe(w_src, b_src, M)
+
+    # -- 5. Evaluate on the test set.
+    with torch.no_grad():
+        # Source baseline: source probe on source test data.
+        src_test_acts_pos = _extract_pooled(source_backend, source_tokenize, test_pos, source_layer)
+        src_test_acts_neg = _extract_pooled(source_backend, source_tokenize, test_neg, source_layer)
+        src_scores = torch.cat([
+            src_probe_net(src_test_acts_pos).squeeze(-1),
+            src_probe_net(src_test_acts_neg).squeeze(-1),
+        ])
+        src_labels = torch.cat([
+            torch.ones(len(src_test_acts_pos)),
+            torch.zeros(len(src_test_acts_neg)),
+        ])
+        baseline_source = auroc_from_scores(src_scores, src_labels)
+
+        # Target baseline: target probe on target test data.
+        tgt_test_acts_pos = _extract_pooled(target_backend, target_tokenize, test_pos, target_layer)
+        tgt_test_acts_neg = _extract_pooled(target_backend, target_tokenize, test_neg, target_layer)
+        tgt_probe_net = pf_tgt.probe.net
+        tgt_scores = torch.cat([
+            tgt_probe_net(tgt_test_acts_pos).squeeze(-1),
+            tgt_probe_net(tgt_test_acts_neg).squeeze(-1),
+        ])
+        tgt_labels = torch.cat([
+            torch.ones(len(tgt_test_acts_pos)),
+            torch.zeros(len(tgt_test_acts_neg)),
+        ])
+        baseline_target = auroc_from_scores(tgt_scores, tgt_labels)
+
+        # Transfer: transferred probe on target activations
+        transfer_scores = torch.cat([
+            (tgt_test_acts_pos @ w_transferred + b_transferred),
+            (tgt_test_acts_neg @ w_transferred + b_transferred),
+        ])
+        transfer_auroc = auroc_from_scores(transfer_scores, tgt_labels)
+
+    return TransferResult(
+        source_arch=source_arch_name,
+        target_arch=target_arch_name,
+        source_layer=source_layer,
+        target_layer=target_layer,
+        n_align_pairs=len(align_texts),
+        baseline_source_auroc=baseline_source,
+        baseline_target_auroc=baseline_target,
+        transfer_auroc=transfer_auroc,
+        transfer_drop=baseline_target - transfer_auroc,
+    )