"""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)