code/thinking_ws_ablate10.py

"""Top-k gradient-pursuit direction ablation (the paper's original mode,
vs ws_ablate.py's single-direction demo mode) on the flexible-generalization
prompt set: h <- h - Q Q^T h, where Q's <=10 columns are the orthonormalised
gradient-pursuit atoms (ws_pursuit_dump's greedy J-lens dictionary) selected
at that exact (layer, position) cell.

Two-pass, fixed basis: pass 1 does one recorded forward over the union of
all condition layers and runs pursuit (k_max=10) at every position of every
layer, producing a per-position orthonormal basis; pass 2 re-runs the
forward with hooks that project out that frozen basis. Positions align
because the prompt is re-tokenised identically and there is no generation.

control10 ablates the same way but with a basis built from 10 fixed random
vocab tokens (seed 0, shared across positions and trials) - it controls for
the effect of removing any 10 directions, not specifically the pursuit ones.

Usage: python ws_ablate10.py {ouro,qwen,huginn} [--limit N]
"""

import argparse
import json
import random

import torch

import ws_common as W
from ws_extra import DATA, boot_ci, groups, next_logits
from ws_pursuit_dump import pursuit_layer

torch.set_grad_enabled(False)

K = 10
CONTROL_SEED = 0


def atoms_for(ctx, v, ids):
    """Unit-normalised atoms for token ids at layer v: same construction as
    ws_pursuit_dump.pursuit_layer's per-step atom, batched over ids."""
    d = ctx.U[ids] @ ctx.Jt[v].T
    return d / d.norm(dim=1, keepdim=True)


def basis_from_atoms(atoms, d):
    """atoms [k,d] unit vectors (k<=K) -> zero-padded orthonormal Q [d,K].

    Zero-padding is safe: a zero column of Q contributes nothing to Q Q^T h."""
    Q = torch.zeros(d, K, device=atoms.device, dtype=atoms.dtype)
    if atoms.shape[0]:
        q, _ = torch.linalg.qr(atoms.T, mode="reduced")
        Q[:, :q.shape[1]] = q
    return Q


def pursuit_bases(ctx, v, h):
    """Per-position orthonormal basis of the top-K pursuit atoms.

    h [T,d] float32 cuda. Returns Q [T,d,K]."""
    T, d = h.shape
    sel, _, _, _ = pursuit_layer(ctx, v, h, k_max=K, k_snap=(K,))
    bases = h.new_zeros(T, d, K)
    for p in range(T):
        idx = sel[p]
        idx = idx[idx >= 0].to(h.device)
        if idx.numel():
            bases[p] = basis_from_atoms(atoms_for(ctx, v, idx), d)
    return bases


def control_bases(ctx, layers):
    """Fixed 10-random-token basis per layer, shared across all positions."""
    rng = random.Random(CONTROL_SEED)
    ids = torch.tensor(rng.sample(range(ctx.U.shape[0]), K), device="cuda")
    out = {}
    for v in layers:
        atoms = atoms_for(ctx, v, ids)
        out[v] = basis_from_atoms(atoms, atoms.shape[1])
    return out


class PursuitAblateHooks:
    """h <- h - Q Q^T h per position, on a frozen pass-1 (or control) basis.

    bases: {v: [T,d,K]} (pursuit, per-position) or {v: [d,K]} (control,
    shared across positions)."""

    def __init__(self, ctx, layers, bases):
        self.ctx, self.layers, self.bases, self.handles = ctx, layers, bases, []

    def _hook(self, v):
        Q = self.bases[v]
        shared = Q.dim() == 2

        def fn(module, inputs, output):
            h = output[0] if isinstance(output, tuple) else output
            hf32 = h.float()
            if shared:
                c = hf32 @ Q
                corr = c @ Q.T
            else:
                c = torch.einsum("btd,tdk->btk", hf32, Q)
                corr = torch.einsum("btk,tdk->btd", c, Q)
            h[:] = (hf32 - corr).to(h.dtype)

        return fn

    def __enter__(self):
        for v in self.layers:
            self.handles.append(
                self.ctx.model.layers[v].register_forward_hook(self._hook(v)))
        return self

    def __exit__(self, *exc):
        for h in self.handles:
            h.remove()


def ablate10(ctx, limit):
    fg = json.load(open(f"{DATA}/flexible-generalization.json"))["categories"]
    g = groups(ctx)
    ks = sorted(g)
    conds = {"pursuit10 all_int": ctx.int_layers}
    if len(ks) >= 3:
        conds["pursuit10 g2 only"] = g[ks[1]]
        conds["pursuit10 last only"] = g[ks[-1]]
    union = sorted(set().union(*conds.values()))
    ctrl_layers = ctx.int_layers
    for v in set(union) | set(ctrl_layers):
        assert v in ctx.Jt, f"missing Jt at layer {v}"
    ctrl_bases = control_bases(ctx, ctrl_layers)

    all_conds = list(conds) + ["control10 all_int"]
    stats = {c: {"kept": 0, "n": 0, "log_ranks": []} for c in all_conds}
    base_n = 0
    for cat in fg[:limit]:
        for func in cat["funcs"]:
            for arg in cat["args"]:
                prompt = func["template"].format(arg=arg)
                ans_tok = W.first_token(ctx, func["answers"][arg])
                base = next_logits(ctx, prompt)
                if int(base.argmax()) != ans_tok:
                    continue
                base_n += 1

                ids = W.encode(ctx, prompt)
                _, acts = W.forward(ctx, ids, layers=union)
                bases = {v: pursuit_bases(ctx, v, acts[v]) for v in union}
                del acts

                trial_conds = {cname: (layers, {v: bases[v] for v in layers})
                               for cname, layers in conds.items()}
                trial_conds["control10 all_int"] = (ctrl_layers, ctrl_bases)
                for cname, (layers, cbases) in trial_conds.items():
                    with PursuitAblateHooks(ctx, layers, cbases):
                        logits = next_logits(ctx, prompt)
                    rank = int((logits > logits[ans_tok]).sum()) + 1
                    stats[cname]["kept"] += rank == 1
                    stats[cname]["n"] += 1
                    stats[cname]["log_ranks"].append(
                        float(torch.log10(torch.tensor(float(rank)))))
                del bases
        print(f"  {cat['name']} done ({base_n} correct-baseline trials so far)",
              flush=True)

    print(f"\nablate10: {base_n} correct-baseline trials", flush=True)
    out = {"base_n": base_n,
           "design": ("k=10 gradient-pursuit atoms per (layer, position), "
                      "fixed two-pass basis: pass 1 records the union of "
                      "condition layers and runs pursuit once per trial; "
                      "pass 2 re-hooks with the frozen QR basis. "
                      "control10 uses 10 fixed random vocab tokens (seed 0) "
                      "building the same per-layer basis, shared across "
                      "positions."),
           "conditions": {}}
    for cname, s in stats.items():
        if not s["n"]:
            continue
        lo, hi = boot_ci(s["kept"], s["n"])
        mlr = sum(s["log_ranks"]) / len(s["log_ranks"])
        print(f"  {cname:20s}: kept {s['kept']}/{s['n']} "
              f"({100*s['kept']/s['n']:.0f}%, CI {lo:.0f}-{hi:.0f}), "
              f"mean log10 rank {mlr:.2f}", flush=True)
        out["conditions"][cname] = {"kept": s["kept"], "n": s["n"],
                                    "mean_log10_rank": mlr,
                                    "log_ranks": s["log_ranks"]}
    return out


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("model", choices=["ouro", "qwen", "huginn"])
    ap.add_argument("--limit", type=int, default=None)
    args = ap.parse_args()
    ctx = W.load_ctx(args.model, full_depth=True)
    print(f"loaded {args.model}: groups "
          f"{[(k, len(v)) for k, v in groups(ctx).items()]}", flush=True)
    result = ablate10(ctx, args.limit)
    with open(f"ws_{args.model}_ablate10.json", "w") as f:
        json.dump(result, f, default=float)
    print("DONE_WS_ABLATE10", flush=True)