File size: 9,024 Bytes
74fd26a ca791bd 74fd26a ca791bd 74fd26a ca791bd 74fd26a | 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 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | """Single-cell causal patch for the single-token organism.
Overwrite ONE cell value c_i(t) inside the soft token z_t (force cell i = digit d), then CONTINUE
free-running so the change propagates through the learned CA rule, and check the final answer equals
what the *true* CA produces for the altered row. This is a surgical test (one cell, one step) that the
latent at (t,i) causally carries c_i(t) AND that the downstream computation uses it.
Patch (per the card): generate free-running to step t; read z_t = softmax(head(h_t)) in R^{K x 10};
set z_t[i] <- one-hot(d) (other cells untouched); re-embed via the codebook (fed = sum_k z_t[k].C[k])
as the input to z_{t+1}; continue free-running; read the answer from the head at z_T, cell q.
Expected = CA-evolve(c_t with cell i := d) for the remaining T-1-t steps, read at cell q.
python -m latent_threads.eval_single_cellpatch --ckpt .../sg1/diffuse/best --n 400
"""
from __future__ import annotations
import argparse
import json
import os
import random
import dotenv
dotenv.load_dotenv()
_USER = os.environ.get("USER", "jbauer")
os.environ.setdefault("HF_HOME", f"/workspace-vast/{_USER}/hf")
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F
from tqdm.auto import tqdm
from latent_threads import tasks as LT
from latent_threads.single import build_single_batch, build_single_mask
from latent_threads.eval_single_report import load, RES
@torch.no_grad()
def patched_answer(model, head, codebook, q_emb, tok, task, batch, im, device, t_p, i_p, d_p):
"""Free-run; at step t_p[j] overwrite cell i_p[j] of the soft token with one-hot(d_p[j]); continue.
Returns the head's predicted digit for each example's queried cell at z_T."""
K, m = task.K, task.m
input_ids, roles, zs_, ls_, lens_ = build_single_batch(tok, task, batch, im, device, with_answer=False)
B, Lmax = input_ids.shape
attn = build_single_mask(roles, zs_, m)
pos = torch.arange(Lmax, device=device)[None].expand(B, Lmax)
E = model.get_input_embeddings()(input_ids)
bidx = torch.arange(B, device=device); zs = torch.tensor(zs_, device=device)
E = E.clone(); E[bidx, zs] = q_emb
out = None
for t in range(m):
out = model(inputs_embeds=E, attention_mask=attn, position_ids=pos, output_hidden_states=True)
if t == m - 1:
break
z = torch.softmax(head(out.hidden_states[-1][bidx, zs + t]).view(B, K, 10).float(), -1) # [B,K,10]
for j in range(B): # patch examples whose t_p == t
if t_p[j] == t:
z[j, i_p[j]] = F.one_hot(torch.tensor(d_p[j], device=device), 10).float()
fed = torch.einsum("bkv,kvd->bd", z, codebook.float())
E = E.clone(); E[bidx, zs + t + 1] = fed.to(E.dtype)
final = head(out.hidden_states[-1][bidx, zs + m - 1]).view(B, K, 10) # head at z_T
return [final[j, batch[j].q].argmax().item() for j in range(B)]
def expected_answer(task, p, t, i, d):
row = list(p.rows[t]); row[i] = d
for _ in range(task.m - 1 - t):
row = task._step(row)
return row[p.q]
def run(ckpt, model_name, n, device):
model, head, codebook, q_emb, tok, did, K, cfg, tag = load(ckpt, model_name, device)
im = tok.convert_tokens_to_ids("<|im_end|>")
m = cfg["task_kwargs"]["m"]; task = LT.make_task(cfg["task"], k=K, m=m)
rng = random.Random(7)
probs = [task.sample(rng) for _ in range(n)]
# per-example patch: a random step t in 0..m-2, cell i, target digit d != current value
t_p, i_p, d_p = [], [], []
for p in probs:
t = rng.randrange(0, m - 1); i = rng.randrange(K)
cur = p.rows[t][i]; d = rng.choice([x for x in range(10) if x != cur])
t_p.append(t); i_p.append(i); d_p.append(d)
exp = np.array([expected_answer(task, probs[j], t_p[j], i_p[j], d_p[j]) for j in range(n)])
orig = np.array([task.answer(probs[j]) for j in range(n)])
preds = []
bs = 16
for s in tqdm(range(0, n, bs), desc=f"{tag} cell-patch"):
b = probs[s:s + bs]
preds += patched_answer(model, head, codebook, q_emb, tok, task, b, im, device,
t_p[s:s + bs], i_p[s:s + bs], d_p[s:s + bs])
preds = np.array(preds)
follows_ca = float((preds == exp).mean())
changed = exp != orig # patch actually moves the answer
follows_on_changed = float((preds[changed] == exp[changed]).mean())
follows_orig_on_changed = float((preds[changed] == orig[changed]).mean()) # ignores the patch?
dist = np.array([m - 1 - t for t in t_p]) # propagation distance
by_dist = {int(dd): float((preds[dist == dd] == exp[dist == dd]).mean()) for dd in sorted(set(dist.tolist()))}
print("=" * 70 + f"\nSINGLE-CELL PATCH {tag} K={K} m={m} n={n}\n" + "=" * 70, flush=True)
print(f" follows CA-propagated answer (all patches): {follows_ca:.3f}", flush=True)
print(f" on patches that CHANGE the answer (n={int(changed.sum())}): follows patched CA={follows_on_changed:.3f} "
f"follows ORIGINAL (ignores patch)={follows_orig_on_changed:.3f}", flush=True)
print(f" by propagation distance (steps from patch to read): { {k: round(v,3) for k,v in by_dist.items()} }", flush=True)
res = dict(tag=tag, K=K, m=m, n=n, follows_ca=follows_ca, n_changed=int(changed.sum()),
follows_on_changed=follows_on_changed, follows_orig_on_changed=follows_orig_on_changed,
by_distance=by_dist)
json.dump(res, open(f"{RES}/single_cellpatch_{tag}.json", "w"), indent=2)
fig, ax = plt.subplots(figsize=(6, 4))
ds = sorted(by_dist); ax.plot(ds, [by_dist[d] for d in ds], "o-", color="#2ca02c", label="follows patched CA")
ax.axhline(1 / 10, ls="--", color="grey", label="chance 0.10")
ax.set_xlabel("propagation distance (CA steps from patched cell to read-out)")
ax.set_ylabel("P(answer = CA-propagated value)"); ax.set_ylim(0, 1.02); ax.legend()
ax.set_title(f"{tag}: single-cell patch c_i(t):=d propagates correctly")
png = f"{RES}/single_cellpatch_{tag}.png"; fig.savefig(png, bbox_inches="tight", dpi=130)
print(png, flush=True)
return res
def run_sweep(ckpt, model_name, n_probs, device):
"""Systematic spot-check: patch EVERY (step t, cell i) site on a small set of problems, force a
changed value, and print the full per-patch detail (orig answer / CA-expected / model) for review."""
model, head, codebook, q_emb, tok, did, K, cfg, tag = load(ckpt, model_name, device)
im = tok.convert_tokens_to_ids("<|im_end|>")
m = cfg["task_kwargs"]["m"]; task = LT.make_task(cfg["task"], k=K, m=m)
rng = random.Random(11)
probs = [task.sample(rng) for _ in range(n_probs)]
sites = [(t, i) for t in range(m - 1) for i in range(K)] # all positions x all cells
batch, T, I, D, meta = [], [], [], [], []
for pi, p in enumerate(probs):
for (t, i) in sites:
cur = p.rows[t][i]; d = rng.choice([x for x in range(10) if x != cur])
batch.append(p); T.append(t); I.append(i); D.append(d)
meta.append((pi, t, i, cur, d, p.q, task.answer(p), expected_answer(task, p, t, i, d)))
preds = []
for s in range(0, len(batch), 16):
preds += patched_answer(model, head, codebook, q_emb, tok, task, batch[s:s + 16], im, device,
T[s:s + 16], I[s:s + 16], D[s:s + 16])
print("=" * 92 + f"\nSINGLE-CELL PATCH SWEEP {tag} K={K} m={m} {len(sites)} sites x {n_probs} problems = {len(meta)} patches\n" + "=" * 92, flush=True)
print(f"{'prob':>4} {'site':>8} {'c_i(t)->d':>10} {'q':>4} {'orig':>5} {'CA_exp':>7} {'model':>6} match", flush=True)
persite, ok_all = {}, []
for k, (pi, t, i, cur, d, q, orig, exp) in enumerate(meta):
mod = preds[k]; ok = (mod == exp); ok_all.append(ok); persite.setdefault((t, i), []).append(ok)
print(f"{pi:>4} {f's{t+1}.c{i+1}':>8} {f'{cur}->{d}':>10} {f'c{q+1}':>4} {orig:>5} {exp:>7} {mod:>6} {'OK' if ok else 'XX'}", flush=True)
print("-" * 92 + "\nper-site follows-CA (over the problems):", flush=True)
for (t, i) in sites:
v = persite[(t, i)]; print(f" s{t+1}.c{i+1}: {sum(v)}/{len(v)} = {sum(v)/len(v):.2f}", flush=True)
print(f"OVERALL follows CA-propagated: {sum(ok_all)}/{len(ok_all)} = {sum(ok_all)/len(ok_all):.3f}", flush=True)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--ckpt", required=True); ap.add_argument("--model", default="Qwen/Qwen3-8B")
ap.add_argument("--n", type=int, default=400)
ap.add_argument("--sweep", action="store_true"); ap.add_argument("--probs", type=int, default=7)
args = ap.parse_args()
if args.sweep:
run_sweep(args.ckpt, args.model, args.probs, "cuda")
else:
run(args.ckpt, args.model, args.n, "cuda")
if __name__ == "__main__":
main()
|