"""Run the preregistered held-out L5-N541 -> L6 branch panel. From the znou tools directory: python probe_541_l6_branch_panel.py """ from __future__ import annotations import csv import os from pathlib import Path os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") import torch from transformers import GPT2LMHeadModel, GPT2TokenizerFast REV = "607a30d783dfa663caf39e06633721c8d4cfcd7e" BOS = 50256 L5, L6, N541 = 5, 6, 541 CANDIDATES = (2712, 1506, 3051, 2662) SCRIPT_DIR = Path(__file__).resolve().parent OUTPUT = SCRIPT_DIR / "results" / "n541_l6_held_out_branch_panel.csv" POSITIVES = [ ("S01", 2712, "They rebuilt the wall brick by brick.", " brick"), ("S02", 2712, "She checked the proof page by page.", " page"), ("S03", 2712, "The guests entered one by one.", " one"), ("S04", 2712, "The signal crossed the valley ridge after ridge.", " ridge"), ("S05", 2712, "The lamps failed room after room.", " room"), ("S06", 2712, "We tested the mechanism step by step.", " step"), ("F01", 1506, "The old receiver crackled now and then.", " then"), ("F02", 1506, "A blue flower appeared here and there.", " there"), ("F03", 1506, "Every now and then the floor trembled.", " then"), ("F04", 1506, "He glanced at the doorway now and then.", " then"), ("F05", 1506, "Small errors surfaced here and there.", " there"), ("F06", 1506, "The buried cable showed through here and there.", " there"), ("C01", 3051, "The wheel turned round and round.", " round"), ("C02", 3051, "She read the warning again and again.", " again"), ("C03", 3051, "The same picture flashed over and over.", " over"), ("C04", 3051, "The argument dragged on and on.", " on"), ("C05", 3051, "Around and around the tiny moon travelled.", " around"), ("C06", 3051, "The bell kept ringing again and again.", " again"), ("A01", 2662, "The two negotiators stood face to face.", " face"), ("A02", 2662, "The boats drifted side by side.", " side"), ("A03", 2662, "The dancers waited shoulder to shoulder.", " shoulder"), ("A04", 2662, "The climber pulled herself hand over hand.", " hand"), ("A05", 2662, "The twins sat back to back.", " back"), ("A06", 2662, "For a moment they stared eye to eye.", " eye"), ] CONTROLS = [ "They rebuilt the wall brick by lantern.", "She checked the proof page by window.", "The guests entered one by accident.", "The signal crossed the valley ridge after sunset.", "The lamps failed room after midnight.", "We tested the mechanism step by rumor.", "The old receiver crackled now and violin.", "A blue flower appeared here and copper.", "Every now and violin the floor trembled.", "He glanced at the doorway now and paper.", "Small errors surfaced here and gravel.", "The buried cable showed through here and winter.", "The wheel turned round and square.", "She read the warning again and tomorrow.", "The same picture flashed over and beneath.", "The argument dragged on and paper.", "Around and beneath the tiny moon travelled.", "The bell kept ringing again and silence.", "The two negotiators stood face to weather.", "The boats drifted side by accident.", "The dancers waited shoulder to music.", "The climber pulled herself hand over payment.", "The twins sat back to work.", "For a moment they stared eye to horizon.", ] CONTROL_COMPLETIONS = [ " lantern", " window", " accident", " sunset", " midnight", " rumor", " violin", " copper", " violin", " paper", " gravel", " winter", " square", " tomorrow", " beneath", " paper", " beneath", " silence", " weather", " accident", " music", " payment", " work", " horizon", ] def top(values): x = torch.topk(values, 2) return int(x.indices[0]), float(x.values[0]), int(x.indices[1]), float(x.values[0] - x.values[1]) class Runner: def __init__(self, model, tok, device): self.model, self.tok, self.device = model, tok, device @torch.inference_mode() def run(self, text, ablate=None): body = self.tok(text, add_special_tokens=False)["input_ids"] ids = [BOS] + body x = torch.tensor([ids], device=self.device) cache = {} def h5(_m, _i, out): y = out.clone() if ablate is not None: y[:, ablate, N541] = 0 cache["l5"] = y.detach().float().cpu()[0] return y def h6(_m, _i, out): cache["l6"] = out.detach().float().cpu()[0] a = self.model.transformer.h[L5].mlp.act.register_forward_hook(h5) b = self.model.transformer.h[L6].mlp.act.register_forward_hook(h6) try: self.model(input_ids=x, use_cache=False) finally: a.remove(); b.remove() tokens = [""] + [self.tok.decode([i]) for i in body] return cache["l5"], cache["l6"], tokens def peak(a, neuron): value, pos = torch.max(a[:, neuron], 0) return float(value), int(pos) def find_completion(tokens, expected): matches = [i for i, token in enumerate(tokens) if token == expected] if not matches: raise RuntimeError(f"Expected completion token {expected!r} not found in {tokens!r}") return matches[-1] def measure(runner, case_id, role, predicted, text, completion_token, pair_id): l5, l6, tokens = runner.run(text) completion = find_completion(tokens, completion_token) l5_value, l5_pos = peak(l5, N541) l5_winner, l5_winner_value, l5_runner, l5_margin = top(l5.amax(0)) l6_winner, l6_winner_value, l6_runner, l6_margin = top(l6.amax(0)) ab5, ab6, _ = runner.run(text, ablate=l5_pos) ab_winner, ab_winner_value, ab_runner, ab_margin = top(ab6.amax(0)) row = { "case_id": case_id, "pair_id": pair_id, "role": role, "predicted_l6": predicted, "text": text, "completion_position": completion, "completion_token": tokens[completion], "l5_541_peak_position": l5_pos, "l5_541_peak_token": tokens[l5_pos], "l5_541_peak_activation": l5_value, "l5_541_peaks_at_completion": l5_pos == completion, "natural_l5_winner": l5_winner, "natural_l5_margin": l5_margin, "natural_l6_winner": l6_winner, "natural_l6_margin": l6_margin, "ablated_l6_winner": ab_winner, "ablated_l6_margin": ab_margin, "winner_changed": l6_winner != ab_winner, } at_completion = [] for n in CANDIDATES: natural = float(l6[completion, n]) ablated = float(ab6[completion, n]) maximum, max_pos = peak(l6, n) row[f"l6_{n}_completion_natural"] = natural row[f"l6_{n}_completion_ablated"] = ablated row[f"l6_{n}_completion_delta"] = ablated - natural row[f"l6_{n}_natural_max"] = maximum row[f"l6_{n}_natural_max_position"] = max_pos at_completion.append((natural, n)) at_completion.sort(reverse=True) row["candidate_winner_at_completion"] = at_completion[0][1] row["prediction_wins_candidates_at_completion"] = ( role == "positive" and at_completion[0][1] == predicted ) return row def main(): if len(POSITIVES) != 24 or len(CONTROLS) != 24 or len(CONTROL_COMPLETIONS) != 24: raise RuntimeError("Panel definition is incomplete") torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False torch.set_float32_matmul_precision("highest") torch.use_deterministic_algorithms(True) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Device: {device}") print("Loading pinned GPT-2 Small in float32...") tok = GPT2TokenizerFast.from_pretrained("gpt2", revision=REV, add_prefix_space=False) model = GPT2LMHeadModel.from_pretrained("gpt2", revision=REV, dtype=torch.float32).to(device).eval() runner = Runner(model, tok, device) rows = [] for i, ((case_id, predicted, positive, completion), control, control_completion) in enumerate(zip(POSITIVES, CONTROLS, CONTROL_COMPLETIONS), 1): pair_id = f"P{i:02d}" rows.append(measure(runner, case_id, "positive", predicted, positive, completion, pair_id)) rows.append(measure(runner, case_id + "K", "control", predicted, control, control_completion, pair_id)) print(f"{i:>2}/24 {case_id} predicted N{predicted}") OUTPUT.parent.mkdir(parents=True, exist_ok=True) with OUTPUT.open("w", newline="", encoding="utf-8") as f: w = csv.DictWriter(f, fieldnames=rows[0].keys()); w.writeheader(); w.writerows(rows) positives = [r for r in rows if r["role"] == "positive"] clean = [r for r in positives if r["l5_541_peaks_at_completion"]] correct = [r for r in positives if r["prediction_wins_candidates_at_completion"]] print(f"Saved: {OUTPUT}") print(f"N541 peaks at declared completion: {len(clean)}/24") print(f"Predicted candidate wins at completion: {len(correct)}/24") if __name__ == "__main__": main()