File size: 4,972 Bytes
8f46582 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | #!/usr/bin/env python3
"""Which exact positions receive the attention mass? Per-head top-k diagnostic.
Disambiguates "the model attends to separator tokens as a class" from "the model
parks on one specific position that happens to be a separator" (a positional sink).
Also reports norm-weighted attention (||alpha_i * v_i|| in the style of Kobayashi et
al. 2020), since a raw-attention sink can be a no-op if its value vector is small.
"""
import argparse
import json
import os
import sys
import numpy as np
import torch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from attention_atlas import build_prompt, load_model, graph_roles, arm_labels
def role_str(roles, p, node_role):
kind, payload = roles[p]
if payload is None:
return kind
comp, depth, _ = node_role.get(payload, ("off", -99, -1))
if kind == "latent":
return f"latent{payload}"
return f"{kind}[{payload}]{comp}{depth}"
@torch.no_grad()
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--ckpt", default=None)
ap.add_argument("--tag", required=True)
ap.add_argument("--val", default="data/star_2arm_L10_valid_fo_bfs.json")
ap.add_argument("--model_id", default="configs/symbol-2layer-8head-768dim-L20.json")
ap.add_argument("--L", type=int, default=10)
ap.add_argument("--n", type=int, default=64)
ap.add_argument("--hops", type=int, nargs="+", default=[1, 5, 10])
ap.add_argument("--topk", type=int, default=5)
ap.add_argument("--device", default="cuda:0")
args = ap.parse_args()
model, tok = load_model(args.ckpt, args.model_id, args.device)
data = json.load(open(args.val))[: args.n]
L = args.L
seqs, metas = [], []
order = np.arange(len(data[0]["edges"]))
for s in data:
ids, roles, root_pos = build_prompt(s, tok, L, order, False)
seqs.append(ids)
metas.append((s, roles, root_pos))
input_ids = torch.tensor(seqs, device=args.device)
B, T = input_ids.shape
attn_mask = torch.ones_like(input_ids)
pos = torch.arange(T, device=args.device).unsqueeze(0).expand(B, -1)
out = model.forward(input_ids, attn_mask, input_ids.clone(), pos)
rep = model.base_causallm(inputs_embeds=out.inputs_embeds, attention_mask=attn_mask,
position_ids=pos, output_attentions=True)
att = torch.stack(rep.attentions, 0).float().cpu().numpy() # (Lyr,B,H,T,T)
# value-vector norms per layer/head, for norm-weighted attention
vnorm = value_norms(model, out.inputs_embeds, attn_mask, pos) # (Lyr,B,H,T)
n_layer, _, n_head = att.shape[0], att.shape[1], att.shape[2]
print(f"\n########## {args.tag} ##########")
for m in args.hops:
print(f"\n--- hop {m} (query = root_pos+{m-1}) ---")
for l in range(n_layer):
for h in range(n_head):
# average the attention row over graphs, in *role* space is hard because
# node ids differ; instead average in absolute-position space, which is
# valid here because edge order is held fixed across these graphs.
rows = np.stack([att[l, bi, h, mp[2] + m - 1, : mp[2] + m]
for bi, mp in enumerate(metas)])
mean_row = rows.mean(0)
top = np.argsort(-mean_row)[: args.topk]
s0, roles0, _ = metas[0]
nr0 = graph_roles(s0, L)
desc = " ".join(f"{p}:{mean_row[p]:.2f}({role_str(roles0, p, nr0)})"
for p in top)
# norm-weighted: alpha * ||v||, renormalized
nw = mean_row * np.stack([vnorm[l, bi, h, : mp[2] + m]
for bi, mp in enumerate(metas)]).mean(0)
nw = nw / max(nw.sum(), 1e-9)
topn = np.argsort(-nw)[: args.topk]
descn = " ".join(f"{p}:{nw[p]:.2f}({role_str(roles0, p, nr0)})"
for p in topn)
print(f" L{l}H{h} raw | {desc}")
print(f" nw | {descn}")
def value_norms(model, embeds, attn_mask, pos):
"""||v_i|| per layer/head, captured by hooking each block's attention module."""
store = {}
def mk_hook(idx):
def hook(mod, inp, out):
hidden = inp[0]
_, _, v = mod.c_attn(hidden).split(mod.split_size, dim=2)
B, T, _ = v.shape
v = v.view(B, T, mod.num_heads, -1)
store[idx] = v.norm(dim=-1).permute(0, 2, 1).float().cpu().numpy()
return hook
handles = [blk.attn.register_forward_hook(mk_hook(i))
for i, blk in enumerate(model.base_causallm.transformer.h)]
model.base_causallm(inputs_embeds=embeds, attention_mask=attn_mask, position_ids=pos)
for hd in handles:
hd.remove()
return np.stack([store[i] for i in sorted(store)], 0) # (Lyr,B,H,T)
if __name__ == "__main__":
main()
|