spb/wpai0001_uqo

1"""Within-model causal localization by activation patching and ablation on toy models.23``patch_sweep`` implements the standard clean/corrupted interchange intervention: run the model on4a *clean* sequence and on a *corrupted* sequence (same program with the target operator replaced5by a different operator, identical teacher-forced prefix), then for every (hook, layer, position)6overwrite the corrupted run's activation with the clean run's activation and measure how much of7the clean next-token answer at the read position ``pos`` is restored.8The effect is reported as the fraction of items whose argmax flips to the clean answer and as the9mean normalized logit-difference recovery. Larger = the patched location carries the information.1011``ablation_sweep`` zeroes (or noises) the activation at each location on clean prompts and reports12the accuracy drop, i.e. necessity. Both sweeps are deterministic given model and data.13"""14from __future__ import annotations1516from dataclasses import dataclass1718import numpy as np19import torch2021from wp_ai_0001.models.hooks import Capture, GaussianNoise, Replace, Zero22from wp_ai_0001.models.toy_transformer import ToyTransformer232425@dataclass26class PatchResult:27    hook_family: str28    layers: list[int]29    positions: list[int]30    flip_rate: np.ndarray  # [L, P]31    logit_recovery: np.ndarray  # [L, P]32    n_items: int333435def _run(model: ToyTransformer, ids: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:36    with torch.no_grad():37        return model.answer_logits(ids, pos)383940def capture_all(model: ToyTransformer, ids: torch.Tensor, pos: torch.Tensor, family: str) -> dict[str, torch.Tensor]:41    cap = Capture()42    names = [f"{family}.{l}" for l in range(model.cfg.n_layers)]43    for n in names:44        model.hooks.add(n, cap)45    try:46        _run(model, ids, pos)47    finally:48        for n in names:49            model.hooks.remove(n, cap)50    return cap.store515253def patch_sweep(model: ToyTransformer, clean_ids: torch.Tensor, corrupt_ids: torch.Tensor, pos: torch.Tensor,54                clean_ans: torch.Tensor, corrupt_ans: torch.Tensor, family: str = "resid_post", positions: list[int] | None = None) -> PatchResult:55    """Patch clean activations into the corrupted run, one (layer, position) at a time."""56    T = clean_ids.shape[1]57    positions = positions if positions is not None else list(range(T))58    clean_acts = capture_all(model, clean_ids, pos, family)59    base = _run(model, corrupt_ids, pos)60    clean_logits = _run(model, clean_ids, pos)61    ar = torch.arange(len(pos), device=clean_ids.device)62    # logit difference clean-answer minus corrupt-answer, in the corrupted run (baseline) and clean run (ceiling)63    ld_base = base[ar, clean_ans] - base[ar, corrupt_ans]64    ld_clean = clean_logits[ar, clean_ans] - clean_logits[ar, corrupt_ans]65    denom = (ld_clean - ld_base).clamp_min(1e-6)66    L = model.cfg.n_layers67    flip = np.zeros((L, len(positions)))68    rec = np.zeros((L, len(positions)))69    for l in range(L):70        name = f"{family}.{l}"71        for j, p in enumerate(positions):72            fn = Replace(clean_acts[name], positions=[p])73            model.hooks.add(name, fn)74            try:75                out = _run(model, corrupt_ids, pos)76            finally:77                model.hooks.remove(name, fn)78            ld = out[ar, clean_ans] - out[ar, corrupt_ans]79            rec[l, j] = float(((ld - ld_base) / denom).clamp(-1, 2).mean())80            flip[l, j] = float((out.argmax(-1) == clean_ans).float().mean())81    return PatchResult(family, list(range(L)), positions, flip, rec, int(len(pos)))828384def ablation_sweep(model: ToyTransformer, ids: torch.Tensor, pos: torch.Tensor, ans: torch.Tensor, family: str = "resid_post",85                   positions: list[int] | None = None, mode: str = "zero", noise_std: float = 1.0, seed: int = 0) -> np.ndarray:86    """Accuracy after ablating (layer, position). Returns [L, P] accuracies."""87    T = ids.shape[1]88    positions = positions if positions is not None else list(range(T))89    L = model.cfg.n_layers90    acc = np.zeros((L, len(positions)))91    for l in range(L):92        name = f"{family}.{l}"93        for j, p in enumerate(positions):94            fn = Zero(positions=[p]) if mode == "zero" else GaussianNoise(noise_std, positions=[p], seed=seed)95            model.hooks.add(name, fn)96            try:97                out = _run(model, ids, pos)98            finally:99                model.hooks.remove(name, fn)100            acc[l, j] = float((out.argmax(-1) == ans).float().mean())101    return acc102103104def head_patch_sweep(model: ToyTransformer, clean_ids: torch.Tensor, corrupt_ids: torch.Tensor, pos: torch.Tensor,105                     clean_ans: torch.Tensor, corrupt_ans: torch.Tensor, position: int) -> np.ndarray:106    """Per-(layer, head) patching of the attention head output at one token position. Returns [L, H] flip rates."""107    clean_acts = capture_all(model, clean_ids, pos, "attn_heads")108    L, H = model.cfg.n_layers, model.cfg.n_heads109    flip = np.zeros((L, H))110    for l in range(L):111        name = f"attn_heads.{l}"112        for h in range(H):113            fn = Replace(clean_acts[name], positions=[position], heads=[h])114            model.hooks.add(name, fn)115            try:116                out = _run(model, corrupt_ids, pos)117            finally:118                model.hooks.remove(name, fn)119            flip[l, h] = float((out.argmax(-1) == clean_ans).float().mean())120    return flip121