bundle/scripts/claim5_steer.py

#!/usr/bin/env python
"""
Claim 5: Causal steering. MFA centroid interventions (Eq 14) steer generations toward
a target concept; the paper reports MFA ~2x the median steering score of SAE/DiffMeans
on Gemma. We run a scaled MFA-vs-DiffMeans comparison (SAE requires Gemmascope +
Neuronpedia infra; documented separately).

Protocol (mirrors the paper / AxBench, scaled down):
  - prompt "<BOS> I think that", sweep alpha, sample completions, score with an LLM
    judge on concept alignment (0-2) and fluency (0-2); final = harmonic mean.
  - MFA:      f_mu(x) = (1-alpha) x + alpha * mu_k         (Eq 14, released MFASteerer)
  - DiffMeans: x + alpha * dhat, where d = mean_resid(concept tokens) - mean_resid(all)
               (paper builds DiffMeans from difference of average token representations)
This script only GENERATES completions; judging is done by claim5_judge.py.
"""
import os, sys, json, time, argparse, collections
import torch

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "repo"))
from modeling.model_checkpointing import load_mfa
from intervention.mfa_steering import MFASteerer


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--model", default="gemma-2-2b")
    ap.add_argument("--layer", type=int, default=18)
    ap.add_argument("--n_concepts", type=int, default=12)
    ap.add_argument("--samples", type=int, default=6)
    ap.add_argument("--max_new", type=int, default=36)
    ap.add_argument("--claim2", default="outputs/claim2_gemma_l18.json")
    ap.add_argument("--out", default="outputs/claim5_completions.json")
    args = ap.parse_args()
    device = "cuda"
    torch.manual_seed(0)
    t0 = time.time()

    from transformer_lens import HookedTransformer
    from datasets import load_dataset
    tl = HookedTransformer.from_pretrained(args.model, device=device, dtype=torch.float32)
    mfa = load_mfa(args.ckpt, map_location=device).eval()
    tok = tl.tokenizer
    hook = f"blocks.{args.layer}.hook_resid_post"

    # ---- build per-token-id mean residual for DiffMeans + global mean ----
    D = mfa.D
    sum_by_tok = collections.defaultdict(lambda: torch.zeros(D))
    cnt_by_tok = collections.Counter()
    gsum = torch.zeros(D); gcnt = 0
    ds = load_dataset("NeelNanda/pile-10k", split="train", streaming=True)
    seen = 0
    for ex in ds:
        toks = tl.to_tokens(ex["text"])[:, :64]
        with torch.no_grad():
            _, cache = tl.run_with_cache(toks, names_filter=hook, return_type=None)
        a = cache[hook].reshape(-1, D).float().cpu()
        tflat = toks.reshape(-1).cpu()
        a, tflat = a[1:], tflat[1:]
        gsum += a.sum(0); gcnt += a.shape[0]
        for ti in tflat.unique().tolist():
            m = tflat == ti
            sum_by_tok[ti] += a[m].sum(0); cnt_by_tok[ti] += int(m.sum())
        seen += a.shape[0]
        if seen >= 120000:
            break
    gmean = (gsum / gcnt).to(device)
    print(f"DiffMeans table built over {seen} tokens, {len(cnt_by_tok)} token types")

    # ---- choose concepts: for each target theme, pick the centroid whose (ln_final)
    #      logit-lens best matches the theme's word list -> clean, interpretable regions. ----
    W_U = tl.W_U
    with torch.no_grad():
        lens_top = (tl.ln_final(mfa.mu) @ W_U).topk(10, dim=-1).indices.cpu()
    lens_words = [set(tok.decode([t]).strip().lower() for t in row.tolist()) for row in lens_top]
    THEMES = {
        "football/sport": ["football", "soccer", "hockey", "league", "tennis", "basketball", "sport"],
        "quantum/physics": ["quantum", "physics", "particle", "energy", "atom", "theory"],
        "freedom": ["free", "freedom", "liberty", "libre", "freely", "independence"],
        "web/internet": ["web", "website", "site", "page", "online", "internet"],
        "music": ["music", "song", "guitar", "album", "band", "musical", "melody"],
        "money/finance": ["money", "financial", "bank", "dollar", "price", "cash", "economic"],
        "war/military": ["war", "military", "army", "soldier", "battle", "weapon", "combat"],
        "medical/health": ["medical", "health", "disease", "doctor", "patient", "hospital"],
        "food": ["food", "eat", "cook", "meal", "recipe", "restaurant", "delicious"],
        "religion": ["god", "church", "religious", "faith", "holy", "prayer", "spiritual"],
        "law/legal": ["law", "legal", "court", "judge", "justice", "lawyer", "crime"],
        "emotion": ["happy", "sad", "anger", "fear", "joy", "love", "emotion", "feeling"],
    }
    concepts = []
    for theme, words in THEMES.items():
        wset = set(words)
        best, bestov = None, 0
        for c in range(mfa.K):
            ov = len(lens_words[c] & wset)
            if ov > bestov:
                bestov, best = ov, c
        if best is not None and bestov >= 2:  # require >=2 theme words in the centroid's top tokens
            cw = [tok.decode([t]).strip() for t in lens_top[best].tolist()
                  if tok.decode([t]).strip().lower() in wset][:5]
            concepts.append({"comp": best, "tokens": cw or words[:4], "kind": theme})
        if len(concepts) >= args.n_concepts:
            break
    def diffmeans_dir(concept_tokens):
        # token ids whose logit-lens/text form matches concept tokens: use ids of these words
        ids = set()
        for w in concept_tokens:
            for pre in (" " + w, w, " " + w.capitalize(), w.capitalize()):
                enc = tok.encode(pre, add_special_tokens=False)
                if len(enc) == 1:
                    ids.add(enc[0])
        vecs = [sum_by_tok[i] / cnt_by_tok[i] for i in ids if cnt_by_tok[i] > 0]
        if not vecs:
            return None
        cmean = torch.stack(vecs).mean(0).to(device)
        d = cmean - gmean
        return d / d.norm().clamp_min(1e-6)

    steerer = MFASteerer(tl, mfa, intervention_type="resid_post")

    def gen_diffmeans(dir_unit, alpha, n):
        toks = tl.to_tokens("I think that")
        def h(value, hook):
            return value + alpha * dir_unit.to(value.dtype)
        outs = []
        for _ in range(n):
            o = tl.generate(toks, max_new_tokens=args.max_new, temperature=1.0,
                            do_sample=True, verbose=False, fwd_hooks=[(hook, h)])
            outs.append(tl.to_string(o[0]))
        return outs

    def gen_mfa(k, alpha, n):
        outs = []
        for _ in range(n):
            outs.append(steerer.generate("I think that", layers=[args.layer], alpha=alpha,
                                         k=k, max_new_tokens=args.max_new, temperature=1.0,
                                         do_sample=True))
        return outs

    # baseline (no steer)
    def gen_base(n):
        toks = tl.to_tokens("I think that")
        return [tl.to_string(tl.generate(toks, max_new_tokens=args.max_new, temperature=1.0,
                                         do_sample=True, verbose=False)[0]) for _ in range(n)]

    mfa_alphas = [0.3, 0.5, 0.7, 0.9]
    dm_scale = gmean.norm().item()          # scale additive dir to activation magnitude
    dm_alphas = [2.0*dm_scale, 4.0*dm_scale, 6.0*dm_scale, 8.0*dm_scale]

    records = {"meta": {"model": args.model, "layer": args.layer,
                        "prompt": "I think that", "samples": args.samples,
                        "mfa_alphas": mfa_alphas, "dm_scales": dm_alphas,
                        "judge": "meta-llama/Llama-3.2-3B-Instruct (open sub for GPT-4o-mini)"},
               "baseline": gen_base(args.samples), "concepts": []}
    print("baseline done")

    for ci, c in enumerate(concepts):
        conc_str = ", ".join(c["tokens"])
        entry = {"comp": c["comp"], "kind": c["kind"], "concept_tokens": c["tokens"],
                 "mfa": {}, "diffmeans": {}}
        for a in mfa_alphas:
            entry["mfa"][str(a)] = gen_mfa(c["comp"], a, args.samples)
        dvec = diffmeans_dir(c["tokens"])
        if dvec is not None:
            for a in dm_alphas:
                entry["diffmeans"][f"{a:.1f}"] = gen_diffmeans(dvec, a, args.samples)
        records["concepts"].append(entry)
        print(f"[{ci+1}/{len(concepts)}] comp {c['comp']} ({c['kind']}) concept='{conc_str}' done")

    records["meta"]["wall_s"] = round(time.time() - t0, 1)
    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    json.dump(records, open(args.out, "w"), indent=2)
    print("saved", args.out, "wall", records["meta"]["wall_s"])


if __name__ == "__main__":
    main()