179 lines
8.4 KiB
Python
179 lines
8.4 KiB
Python
"""Step 3d (clean metric) — end-to-end re-route validated by NLL, not generation.
|
|
|
|
Teacher-force the gold answer; at each answer position choose the ACTIVE block set, build the
|
|
attention mask, and measure the gold token's NLL. No generation loop -> immune to the rambly-output /
|
|
memory-fallback noise that broke the accuracy metric. Modes (gold NLL, lower = better):
|
|
|
|
dense : all blocks active (ceiling)
|
|
static : top-k blocks frozen from the first answer position's real attention
|
|
landmark-reroute : top-k by attention to each block's resident LANDMARK token, re-picked per position
|
|
oracle-reroute : top-k by full-block attention mass, re-picked per position (best-case selection)
|
|
|
|
If landmark-reroute NLL ~= dense ~= oracle-reroute, and static NLL is high, on-demand re-route works.
|
|
Counterfactual facts (fictional capitals) force reading the docs; gold blocks are randomly placed.
|
|
|
|
Run: source env.sh && python scripts/e2e_nll.py --model Qwen/Qwen3-8B --n-hops 3
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import random
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
NEG = -1e9
|
|
SYS = "You are a helpful assistant. Use the documents to answer. Give only the capitals, comma-separated, in order."
|
|
PAIRS = [("France", "Paris"), ("Japan", "Tokyo"), ("Egypt", "Cairo"), ("Brazil", "Brasilia"),
|
|
("Canada", "Ottawa"), ("Kenya", "Nairobi"), ("Norway", "Oslo"), ("Peru", "Lima"),
|
|
("India", "Delhi"), ("Spain", "Madrid"), ("Greece", "Athens"), ("Cuba", "Havana")]
|
|
FICT = ["Zandar", "Qeropolis", "Vunbry", "Marnex", "Trellin", "Osquith", "Balmoor", "Kestol",
|
|
"Yarrow", "Drennik", "Fenwick", "Lorvath", "Sundeep", "Ashgard", "Perrin", "Wexley"]
|
|
|
|
|
|
def build(tok, rng, n_blocks, n_hops, sink_ids):
|
|
docs = rng.sample(range(len(PAIRS)), n_blocks) # which countries are documents
|
|
asked = rng.sample(docs, n_hops) # asked subset (gold blocks, random placement)
|
|
slot = {d: bi for bi, d in enumerate(docs)}
|
|
fic = FICT[:]; rng.shuffle(fic)
|
|
cap = {d: fic[i] for i, d in enumerate(docs)} # fictional capital per country
|
|
|
|
ids, blk, sink, gold = [], [], [], [] # gold[t] = active gold block per answer token
|
|
|
|
def add(text, b, is_sink=False, special=False, gblock=None, is_ans=False):
|
|
t = tok(text, add_special_tokens=special)["input_ids"]
|
|
ids.extend(t); blk.extend([b] * len(t)); sink.extend([is_sink] * len(t))
|
|
if is_ans:
|
|
gold.extend([gblock] * len(t))
|
|
|
|
add(SYS + "\n", -1, special=True)
|
|
for d in docs:
|
|
for s in sink_ids:
|
|
ids.append(s); blk.append(slot[d]); sink.append(True)
|
|
add(f"[Document {slot[d]+1}] The capital of {PAIRS[d][0]} is {cap[d]}.\n", slot[d])
|
|
add(f"List the capitals of: {', '.join(PAIRS[d][0] for d in asked)}.\nAnswer:", -2)
|
|
a0 = len(ids)
|
|
for j, d in enumerate(asked):
|
|
prefix = " " + cap[d] if j == 0 else ", " + cap[d]
|
|
add(prefix, -2, gblock=slot[d], is_ans=True)
|
|
return ids, blk, sink, a0, gold
|
|
|
|
|
|
def landmarks_content(blk, sink, n_blocks):
|
|
lm, isc = {}, [False] * len(blk)
|
|
for j, (b, s) in enumerate(zip(blk, sink)):
|
|
if b >= 0 and not s:
|
|
isc[j] = True; lm[b] = j
|
|
return [lm[b] for b in range(n_blocks)], isc
|
|
|
|
|
|
def mask_for(blk, sink, a0, n_blocks, active_rows, dev):
|
|
n = len(blk)
|
|
b = torch.tensor(blk, device=dev)
|
|
lm, isc = landmarks_content(blk, sink, n_blocks)
|
|
causal = torch.tril(torch.ones(n, n, dtype=torch.bool, device=dev))
|
|
is_static = (b == -1) | (b == -2) # system + question/answer visible to answer rows
|
|
is_lm = torch.zeros(n, dtype=torch.bool, device=dev); is_lm[torch.tensor(lm, device=dev)] = True
|
|
static_key = is_static | is_lm
|
|
isc_t = torch.tensor(isc, device=dev)
|
|
allowed = torch.zeros(n, n, dtype=torch.bool, device=dev)
|
|
bi, bj = b.view(n, 1), b.view(1, n)
|
|
allowed[:a0] = ((bj == -1) | (bi == bj) | (bi == -2))[:a0] # prompt: block attention
|
|
for i in range(a0, n):
|
|
act = active_rows[i - a0]
|
|
vis = static_key.clone()
|
|
if act:
|
|
vis = vis | (isc_t & torch.isin(bj.view(n)[:], torch.tensor(sorted(act), device=dev)))
|
|
allowed[i] = vis
|
|
allowed &= causal
|
|
return torch.where(allowed, 0.0, NEG).view(1, 1, n, n).float()
|
|
|
|
|
|
@torch.no_grad()
|
|
def gold_nll(model, ids, mask, a0, dev):
|
|
t = torch.tensor([ids], device=dev)
|
|
lp = F.log_softmax(model(input_ids=t, attention_mask=mask).logits[0].float(), -1)
|
|
tgt = t[0, a0:]
|
|
return -lp[a0 - 1:-1].gather(-1, tgt.view(-1, 1)).squeeze(-1) # per-answer-token NLL
|
|
|
|
|
|
@torch.no_grad()
|
|
def signals(model, ids, blk, sink, a0, n_blocks, dev):
|
|
"""Dense pass -> per-answer-position full-block mass and landmark attention."""
|
|
n = len(ids)
|
|
dense_rows = {i: set(range(n_blocks)) for i in range(n - a0)}
|
|
m = mask_for(blk, sink, a0, n_blocks, dense_rows, dev)
|
|
out = model(input_ids=torch.tensor([ids], device=dev), attention_mask=m, output_attentions=True)
|
|
a = torch.stack(out.attentions, 0)[:, 0].mean(dim=(0, 1)).float() # [n,n]
|
|
lm, isc = landmarks_content(blk, sink, n_blocks)
|
|
content = torch.zeros(n_blocks, n, device=dev)
|
|
for j, c in enumerate(isc):
|
|
if c:
|
|
content[blk[j], j] = 1.0
|
|
full_mass = a[a0:n] @ content.T # [A, nb]
|
|
lm_mass = a[a0:n][:, torch.tensor(lm, device=dev)] # [A, nb]
|
|
return full_mass.cpu(), lm_mass.cpu(), m
|
|
|
|
|
|
def topk_rows(scores, k, A, static=False):
|
|
if static:
|
|
sel = set(scores[0].argsort(descending=True)[:k].tolist())
|
|
return {t: set(sel) for t in range(A)}
|
|
return {t: set(scores[t].argsort(descending=True)[:k].tolist()) for t in range(A)}
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--model", default="Qwen/Qwen3-8B")
|
|
ap.add_argument("--n-blocks", type=int, default=8)
|
|
ap.add_argument("--n-hops", type=int, default=3)
|
|
ap.add_argument("--k", type=int, default=1)
|
|
ap.add_argument("--n-examples", type=int, default=30)
|
|
ap.add_argument("--n-sinks", type=int, default=4)
|
|
ap.add_argument("--seed", type=int, default=1)
|
|
args = ap.parse_args()
|
|
dev = "cuda"; rng = random.Random(args.seed)
|
|
|
|
print(f"loading {args.model} ...")
|
|
tok = AutoTokenizer.from_pretrained(args.model)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
args.model, dtype=torch.bfloat16, attn_implementation="eager").to(dev).eval()
|
|
sink_ids = tok("\n", add_special_tokens=False)["input_ids"] * args.n_sinks
|
|
|
|
agg = {m: [] for m in ["dense", "static", "landmark", "oracle"]}
|
|
cov = {m: [] for m in ["static", "landmark", "oracle"]}
|
|
loads = {m: [] for m in ["dense", "static", "landmark", "oracle"]}
|
|
for _ in range(args.n_examples):
|
|
ids, blk, sink, a0, gold = build(tok, rng, args.n_blocks, args.n_hops, sink_ids)
|
|
A = len(ids) - a0
|
|
full_mass, lm_mass, dense_mask = signals(model, ids, blk, sink, a0, args.n_blocks, dev)
|
|
sel = {
|
|
"dense": {t: set(range(args.n_blocks)) for t in range(A)},
|
|
"static": topk_rows(full_mass, args.k, A, static=True),
|
|
"landmark": topk_rows(lm_mass, args.k, A),
|
|
"oracle": topk_rows(full_mass, args.k, A),
|
|
}
|
|
for mode, rows in sel.items():
|
|
mask = dense_mask if mode == "dense" else mask_for(blk, sink, a0, args.n_blocks, rows, dev)
|
|
nll = gold_nll(model, ids, mask, a0, dev)
|
|
agg[mode].append(nll.mean().item())
|
|
loads[mode].append(sum(len(rows[t]) for t in range(A)) / A)
|
|
if mode != "dense":
|
|
cov[mode].append(sum(gold[t] in rows[t] for t in range(A)) / A)
|
|
|
|
def mean(x):
|
|
return sum(x) / len(x)
|
|
print(f"\nmodel={args.model} n_blocks={args.n_blocks} n_hops={args.n_hops} k={args.k} n={args.n_examples}")
|
|
print(f"{'mode':>10} | {'gold NLL':>9} | {'avg active':>10} | {'gold-block coverage':>19}")
|
|
for mode in ["dense", "static", "landmark", "oracle"]:
|
|
c = f"{mean(cov[mode]):.2f}" if mode in cov else "-"
|
|
print(f"{mode:>10} | {mean(agg[mode]):>9.3f} | {mean(loads[mode]):>10.2f} | {c:>19}")
|
|
d, s, l = mean(agg['dense']), mean(agg['static']), mean(agg['landmark'])
|
|
print(f"\nlandmark-reroute closes {(s-l)/(s-d)*100 if s>d else 0:.0f}% of the static->dense NLL gap "
|
|
f"at {mean(loads['landmark']):.1f}/{args.n_blocks} blocks active.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|