import os from pathlib import Path os.environ.setdefault("TRANSFORMERLENS_ALLOW_MPS", "1") import math from collections import Counter import torch from dataclasses import dataclass from itertools import combinations from typing import Dict, List, Optional, Set, Tuple from transformer_lens import HookedTransformer from benchmark_specs import BENCHMARK_SUITE, BenchmarkSpec import frontier_lab import milestone_interp def preferred_device() -> str: if torch.cuda.is_available(): return "cuda" if torch.backends.mps.is_available(): return "mps" return "cpu" DEVICE = preferred_device() # Half precision on accelerator backends; CPU keeps float32 for speed and compatibility. DTYPE = torch.float16 if DEVICE in ("cuda", "mps") else torch.float32 MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct" def _dump_run_source_snapshots(repo_dir: Path) -> None: """Echo key Python modules to stdout before any model work (repro ledger).""" files = ( ("benchmark_specs.py", repo_dir / "benchmark_specs.py"), ("main.py", repo_dir / "main.py"), ("frontier_lab.py", repo_dir / "frontier_lab.py"), ("milestone_interp.py", repo_dir / "milestone_interp.py"), ) for label, path in files: bar = "=" * 79 print(f"\n{bar}\n# BEGIN FILE DUMP: {label} ({path})\n{bar}\n") src = path.read_text(encoding="utf-8") print(src, end="" if src.endswith("\n") else "\n") print(f"{bar}\n# END FILE DUMP: {label}\n{bar}\n") _REPO_ROOT = Path(__file__).resolve().parent _skip_snapshot = ( os.environ.get("PG_SKIP_SOURCE_DUMP", "").strip().lower() in ("1", "true", "yes", "on") ) if not _skip_snapshot: _dump_run_source_snapshots(_REPO_ROOT) # --------------------------------------------------------- # Load model # --------------------------------------------------------- model = HookedTransformer.from_pretrained_no_processing( MODEL_NAME, device=DEVICE, dtype=DTYPE, ) # --------------------------------------------------------- # Experiment config # --------------------------------------------------------- # What intervention localizes: # divergence between *this* clean vs corrupt trajectory at the hook (not "Paris knowledge" carte blanche). CLEAN_PROMPT = "The capital of France is" CORRUPT_PROMPT = "The capital of France is Rome. The correct answer is" TARGET_TOKEN = " Paris" # Wrong-capital probe for margins (paired with factual target). Tune per falsification variant when spelling changes. DISTRACTOR_TOKEN = " Rome" # Surfaces with tiny |clean_tgt−corrupt_tgt| blow up normalized metrics and distort Jaccard — exclude them. # ~1 logit separates "usable contrast" from borderline denominators where norm is still jittery (~1 gaps). MIN_ABS_TARGET_GAP_FOR_ROBUSTNESS = 1.5 # Attention head sweep uses single-head replacements — underestimates pairs / small ensembles (known limit). # Scripted template mutants (same distractor scaffolding) — cheap automated stress beyond hand-picked multilingual rows. AUTO_STRESS_PAIRS: List[Tuple[str, str, str]] = [ ( "stress_double_space", "The capital of France is", "The capital of France is Rome. The correct answer is", ), ( "stress_lower", "the capital of france is", "the capital of france is rome. the correct answer is", ), ( "stress_dash_distractor", "The capital of France is", "The capital of France is Rome — wrong choice. The correct answer is", ), ( "stress_synonym_fact", "The nation France's administrative capital is", "The nation France's administrative capital is Rome. The correct answer is", ), ] def robustness_surface_eligible(clean_tgt: float, corrupt_tgt: float) -> bool: """Narrow-band or inverted surfaces pollute normalized / Jaccard summaries.""" if corrupt_tgt >= clean_tgt - 1e-9: return False gap = abs(clean_tgt - corrupt_tgt) if gap < MIN_ABS_TARGET_GAP_FOR_ROBUSTNESS: return False return True def nt_nm_sign_conflict(nt: float, nm: float) -> bool: """Finite nt and nm with opposite sign: logit-trajectory disagrees with tgt−distractor margin trajectory.""" if math.isnan(nt) or math.isnan(nm) or math.isinf(nt) or math.isinf(nm): return False if nt == 0.0 or nm == 0.0: return False return (nt > 0.0) != (nm > 0.0) @dataclass class PatchResult: site: str layer: int target_logit: float distractor_logit: float logit_margin: float normalized_effect: float #: Same geometry as normalized_effect but denom = Δ(clean_corr margin); often stabler thin tgt gaps. normalized_margin_effect: float head: int | None = None def final_residual_layer_idx() -> int: """Last transformer block residual before unembedding (= trivial replay into LM head when patched alone).""" return model.cfg.n_layers - 1 def excluding_readout_residual_rows(results: List[PatchResult]) -> List[PatchResult]: """Drop hook_resid_post at final layer — replacement feeds donor residual directly into unembedding.""" hi = final_residual_layer_idx() return [r for r in results if not (r.layer == hi and "resid_post" in r.site)] def residual_readout_bound_row(rows: List[PatchResult]) -> PatchResult | None: """The hook_resid_post row at final block (analytical upper norm bound), if present.""" hi = final_residual_layer_idx() for r in rows: if r.layer == hi and "resid_post" in r.site: return r return None @dataclass(frozen=True) class FalsificationPrompt: clean: str corrupt: str target_token: str distractor_token: str = DISTRACTOR_TOKEN label: str = "" # English surfaces + multilingual parallels. Each row carries distractor spelling aligned with the corrupt wrong city token when it differs from " Rome". PROMPT_VARIANTS: List[FalsificationPrompt] = [ FalsificationPrompt( CLEAN_PROMPT, CORRUPT_PROMPT, TARGET_TOKEN, DISTRACTOR_TOKEN, label="default", ), FalsificationPrompt( "France's capital is", "France's capital is Rome. The correct answer is", TARGET_TOKEN, DISTRACTOR_TOKEN, label="france_alt", ), FalsificationPrompt( "The city governing France is", "The city governing France is Rome. The correct answer is", TARGET_TOKEN, DISTRACTOR_TOKEN, label="governing_city", ), FalsificationPrompt( "What is the capital of France?", "What is the capital of France? It is Rome. The correct answer is", TARGET_TOKEN, DISTRACTOR_TOKEN, label="wh_question", ), FalsificationPrompt( "The French capital is", "The French capital is Rome. The correct answer is", TARGET_TOKEN, DISTRACTOR_TOKEN, label="french_synonym_np", ), FalsificationPrompt( CLEAN_PROMPT, "The capital of France is London. The correct answer is", TARGET_TOKEN, " London", label="distractor_london", ), FalsificationPrompt( "La capitale de la France est", "La capitale de la France est Rome. La bonne réponse est", TARGET_TOKEN, DISTRACTOR_TOKEN, label="french_sentence", ), FalsificationPrompt( "La capital de Francia es", "La capital de Francia es Roma. La respuesta correcta es", TARGET_TOKEN, " Roma", label="spanish_sentence", ), FalsificationPrompt( "Die Hauptstadt von Frankreich ist", "Die Hauptstadt von Frankreich ist Rom. Die richtige Antwort ist", TARGET_TOKEN, " Rom", label="german_sentence", ), ] def _token_logit(logits: torch.Tensor, token: str) -> float: tid = model.to_single_token(token) return logits[0, -1, tid].item() def normalized_denoise( patched_tgt: float, corrupt_tgt: float, clean_tgt: float, ) -> float: denom = clean_tgt - corrupt_tgt if abs(denom) < 1e-9 or denom < 0: # corrupt >= clean — scripted "recovery toward clean" is ill-posed / inverted return float("nan") return (patched_tgt - corrupt_tgt) / denom def normalized_noise( patched_tgt: float, corrupt_tgt: float, clean_tgt: float, ) -> float: denom = clean_tgt - corrupt_tgt if abs(denom) < 1e-9 or denom < 0: return float("nan") return (clean_tgt - patched_tgt) / denom def normalized_denoise_margin( patched_margin: float, corrupt_margin: float, clean_margin: float, ) -> float: """Fraction of (clean−corr) **target−distractor margin** closed by patching (contrast-normalized analogue).""" denom = clean_margin - corrupt_margin if abs(denom) < 1e-9 or denom <= 0: return float("nan") return (patched_margin - corrupt_margin) / denom def normalized_noise_margin( patched_margin: float, corrupt_margin: float, clean_margin: float, ) -> float: denom = clean_margin - corrupt_margin if abs(denom) < 1e-9 or denom <= 0: return float("nan") return (clean_margin - patched_margin) / denom def patch_metrics( patched_logits: torch.Tensor, target_token: str, distractor_token: str, corrupt_tgt_baseline: float, clean_tgt_baseline: float, corrupt_margin_baseline: float, clean_margin_baseline: float, *, direction: str, ) -> Tuple[float, float, float, float, float]: tgt = _token_logit(patched_logits, target_token) dist = _token_logit(patched_logits, distractor_token) margin = tgt - dist if direction == "denoise": ne = normalized_denoise(tgt, corrupt_tgt_baseline, clean_tgt_baseline) ne_margin = normalized_denoise_margin( margin, corrupt_margin_baseline, clean_margin_baseline, ) elif direction == "noise": ne = normalized_noise(tgt, corrupt_tgt_baseline, clean_tgt_baseline) ne_margin = normalized_noise_margin( margin, corrupt_margin_baseline, clean_margin_baseline, ) else: raise ValueError(direction) return tgt, dist, margin, ne, ne_margin def patch_last_pos_inject_from_cache( run_act: torch.Tensor, hook, donor_cache, ) -> torch.Tensor: """Replace final position activation with donor cache slice (same hook name).""" donor = donor_cache[hook.name] run_act[:, -1, :] = donor[:, -1, :] return run_act def zero_last_pos_2d(t: torch.Tensor, *, hook) -> torch.Tensor: t[:, -1, :] = 0 return t def patch_last_pos_attn_z_head_from_cache( run_act: torch.Tensor, hook, donor_cache, head: int, ) -> torch.Tensor: donor = donor_cache[hook.name] run_act[:, -1, head, :] = donor[:, -1, head, :] return run_act def sweep_resid_post_denoise( corrupt_prompt: str, clean_activations, *, target_token: str, distractor_token: str, clean_tgt: float, corrupt_tgt: float, clean_margin: float, corrupt_margin: float, ) -> List[PatchResult]: results: List[PatchResult] = [] for layer in range(model.cfg.n_layers): hook_name = f"blocks.{layer}.hook_resid_post" patched_logits = model.run_with_hooks( corrupt_prompt, fwd_hooks=[ ( hook_name, lambda act, *, hook, ca=clean_activations: patch_last_pos_inject_from_cache( act, hook, ca, ), ), ], ) tg, dd, marg, norm, norm_m = patch_metrics( patched_logits, target_token, distractor_token, corrupt_tgt, clean_tgt, corrupt_margin, clean_margin, direction="denoise", ) results.append( PatchResult( site="hook_resid_post_den", layer=layer, target_logit=tg, distractor_logit=dd, logit_margin=marg, normalized_effect=norm, normalized_margin_effect=norm_m, ), ) return results def sweep_resid_post_noise( clean_prompt: str, corrupt_activations, *, target_token: str, distractor_token: str, clean_tgt: float, corrupt_tgt: float, clean_margin: float, corrupt_margin: float, ) -> List[PatchResult]: results: List[PatchResult] = [] for layer in range(model.cfg.n_layers): hook_name = f"blocks.{layer}.hook_resid_post" patched_logits = model.run_with_hooks( clean_prompt, fwd_hooks=[ ( hook_name, lambda act, *, hook, cc=corrupt_activations: patch_last_pos_inject_from_cache( act, hook, cc, ), ), ], ) tg, dd, marg, norm, norm_m = patch_metrics( patched_logits, target_token, distractor_token, corrupt_tgt, clean_tgt, corrupt_margin, clean_margin, direction="noise", ) results.append( PatchResult( site="hook_resid_post_noise", layer=layer, target_logit=tg, distractor_logit=dd, logit_margin=marg, normalized_effect=norm, normalized_margin_effect=norm_m, ), ) return results def sweep_hook_mlp_out_denoise( corrupt_prompt: str, clean_activations, *, target_token: str, distractor_token: str, clean_tgt: float, corrupt_tgt: float, clean_margin: float, corrupt_margin: float, ) -> List[PatchResult]: results: List[PatchResult] = [] for layer in range(model.cfg.n_layers): hook_name = f"blocks.{layer}.hook_mlp_out" patched_logits = model.run_with_hooks( corrupt_prompt, fwd_hooks=[ ( hook_name, lambda act, *, hook, ca=clean_activations: patch_last_pos_inject_from_cache( act, hook, ca, ), ), ], ) tg, dd, marg, norm, norm_m = patch_metrics( patched_logits, target_token, distractor_token, corrupt_tgt, clean_tgt, corrupt_margin, clean_margin, direction="denoise", ) results.append( PatchResult( site="hook_mlp_out_den", layer=layer, target_logit=tg, distractor_logit=dd, logit_margin=marg, normalized_effect=norm, normalized_margin_effect=norm_m, ), ) return results def sweep_attn_hook_z_denoise( corrupt_prompt: str, clean_activations, *, target_token: str, distractor_token: str, clean_tgt: float, corrupt_tgt: float, clean_margin: float, corrupt_margin: float, ) -> List[PatchResult]: # Single-head protocol — misses multi-head substrates. results: List[PatchResult] = [] n_heads = model.cfg.n_heads for layer in range(model.cfg.n_layers): hook_name = f"blocks.{layer}.attn.hook_z" for head in range(n_heads): patched_logits = model.run_with_hooks( corrupt_prompt, fwd_hooks=[ ( hook_name, lambda act, *, hook, ca=clean_activations, hd=head: patch_last_pos_attn_z_head_from_cache( act, hook, ca, hd, ), ), ], ) tg, dd, marg, norm, norm_m = patch_metrics( patched_logits, target_token, distractor_token, corrupt_tgt, clean_tgt, corrupt_margin, clean_margin, direction="denoise", ) results.append( PatchResult( site="hook_z_den", layer=layer, head=head, target_logit=tg, distractor_logit=dd, logit_margin=marg, normalized_effect=norm, normalized_margin_effect=norm_m, ), ) return results def top_k_scores(results: List[PatchResult], k: int) -> List[PatchResult]: def score(r: PatchResult) -> float: if not math.isnan(r.normalized_effect): return r.normalized_effect return r.logit_margin return sorted(results, key=score, reverse=True)[:k] def degenerate_best(results: List[PatchResult]) -> PatchResult: """Prefer normalized_effect when finite; otherwise maximize target logit.""" def key(r: PatchResult) -> Tuple[int, float, float]: fin = int(not math.isnan(r.normalized_effect)) ne = r.normalized_effect if fin else float("-inf") return (fin, ne, r.target_logit) return max(results, key=key) def degenerate_best_residual_competitive(rows: List[PatchResult]) -> PatchResult | None: """Best residual-site result excluding trivial final-block readout replay.""" comp = excluding_readout_residual_rows(rows) return degenerate_best(comp) if comp else None def top_k_layer_ids( results: List[PatchResult], k: int, *, exclude_layers: Optional[Set[int]] = None, ) -> List[int]: rows = ( results if exclude_layers is None else [r for r in results if r.layer not in exclude_layers] ) return [r.layer for r in top_k_scores(rows, k)] READOUT_UPPER_BOUND_NOTE = ( "Final block hook_resid_post is an analytical upper bound (donor residual → LM head), " "not a competitive causal site vs earlier layers." ) def jaccard(a: Set[int], b: Set[int]) -> float: if not a and not b: return 1.0 return len(a & b) / len(a | b) def print_explanation_robustness( label_to_results: Dict[str, List[PatchResult]], k: int, surface_targets: Dict[str, Tuple[float, float]], min_abs_gap: float, ) -> None: """Quantify overlap of top-k causal layers across falsification surfaces.""" eligible_labels: List[str] = [] skipped: List[str] = [] for lab, (c_raw, r_raw) in surface_targets.items(): if lab not in label_to_results: continue if robustness_surface_eligible(c_raw, r_raw): eligible_labels.append(lab) else: reason = [] if r_raw >= c_raw - 1e-9: reason.append("corrupt_tgt>=clean_tgt") elif abs(c_raw - r_raw) < min_abs_gap: reason.append(f"|gap|<{min_abs_gap}") skipped.append(f"{lab} ({', '.join(reason)})") if not eligible_labels: print( "\n=== EXPLANATION ROBUSTNESS (skipped — no eligible surfaces after filters) ===" ) print("Skipped:", "; ".join(skipped) if skipped else "(none)") return filtered = {lab: label_to_results[lab] for lab in eligible_labels} labels = list(filtered.keys()) sets: Dict[str, Set[int]] = { lab: set( top_k_layer_ids( filtered[lab], k, exclude_layers={final_residual_layer_idx()}, ), ) for lab in labels } pairs = list(combinations(labels, 2)) mean_j = sum(jaccard(sets[a], sets[b]) for a, b in pairs) / max(len(pairs), 1) consensus: Set[int] = set.intersection(*sets.values()) if sets else set() n = len(labels) majority: Set[int] = { layer for layer in set().union(*sets.values()) if sum(1 for s in sets.values() if layer in s) > n / 2 } print( f"\n=== EXPLANATION ROBUSTNESS (top-{k} layers among blocks 0–{final_residual_layer_idx() - 1} only — " f"omit final residual as readout upper bound; eligible surfaces: " f"|clean_tgt−corrupt_tgt|≥{min_abs_gap}, corrupt50% eligible) top-{k} layers: {sorted(majority)}") def format_patch_line(r: PatchResult) -> str: nt = ( "nan" if math.isnan(r.normalized_effect) else f"{r.normalized_effect:.3f}" ) nmar = ( "nan" if math.isnan(r.normalized_margin_effect) else f"{r.normalized_margin_effect:.3f}" ) return ( f"L{r.layer:02d} tgt={r.target_logit:.3f} margin={r.logit_margin:.3f} " f"nt={nt} nm={nmar}" ) def format_head_line(r: PatchResult) -> str: nt = "nan" if math.isnan(r.normalized_effect) else f"{r.normalized_effect:.3f}" nmar = "nan" if math.isnan(r.normalized_margin_effect) else f"{r.normalized_margin_effect:.3f}" return ( f"L{r.layer:02d} H{r.head:02d} tgt={r.target_logit:.3f} " f"margin={r.logit_margin:.3f} nt={nt} nm={nmar}" ) def format_layer_effect_pair(r: PatchResult, decimals: int = 2) -> str: """Compact L + nt/nm for tables (decimals=2) or summaries.""" w = decimals nt = "nan" if math.isnan(r.normalized_effect) else f"{r.normalized_effect:.{w}f}" nmar = "nan" if math.isnan(r.normalized_margin_effect) else f"{r.normalized_margin_effect:.{w}f}" return f"L{r.layer} nt={nt} nm={nmar}" def print_ranked_residual_mlp( title: str, rows: List[PatchResult], k: int, *, annotate_readout_bound: bool = True, ) -> None: print(f"\n=== {title} ===") competitive = excluding_readout_residual_rows(rows) if annotate_readout_bound else rows if annotate_readout_bound: print(f"(Top-{k} omits L{final_residual_layer_idx():02d} hook_resid_post: {READOUT_UPPER_BOUND_NOTE})") for r in top_k_scores(competitive, k): print(format_patch_line(r)) if annotate_readout_bound: rb = residual_readout_bound_row(rows) if rb is not None: nt_rb = ( "nan" if math.isnan(rb.normalized_effect) else f"{rb.normalized_effect:.3f}" ) nmt_rb = ( "nan" if math.isnan(rb.normalized_margin_effect) else f"{rb.normalized_margin_effect:.3f}" ) print( f"L{rb.layer:02d} tgt={rb.target_logit:.3f} margin={rb.logit_margin:.3f} " f"nt={nt_rb} nm={nmt_rb} [upper bound — not ranked vs earlier blocks]" ) def print_ranked_heads(title: str, rows: List[PatchResult], k: int) -> None: print(f"\n=== {title} (top {k} head sites) ===") print("(nt=logit frac; nm=margin frac — same semantics as residual tables.)") for r in top_k_scores(rows, k): print(format_head_line(r)) @dataclass class BenchmarkQuickOutcome: spec: BenchmarkSpec clean_tgt: float corrupt_tgt: float margin_clean: float margin_corr: float top_den_layer: int top_den_norm: float top_den_norm_margin: float top_noise_layer: int top_noise_norm: float top_noise_norm_margin: float thin_gap_warn: bool inverted_gap: bool denoise_nt_nm_conflict: bool noise_nt_nm_conflict: bool def try_benchmark_quick(spec: BenchmarkSpec) -> Optional[BenchmarkQuickOutcome]: try: model.to_single_token(spec.target_token) model.to_single_token(spec.distractor_token) except Exception as exc: print(f"[SKIP {spec.id}] token check failed ({exc})") return None cl, ca = model.run_with_cache(spec.clean_prompt) crl, cc = model.run_with_cache(spec.corrupt_prompt) c_tgt = _token_logit(cl, spec.target_token) r_tgt = _token_logit(crl, spec.target_token) c_dist = _token_logit(cl, spec.distractor_token) r_dist = _token_logit(crl, spec.distractor_token) m_c = c_tgt - c_dist m_r = r_tgt - r_dist inverted = bool(r_tgt >= c_tgt - 1e-9) thin = bool(abs(c_tgt - r_tgt) < MIN_ABS_TARGET_GAP_FOR_ROBUSTNESS) den = sweep_resid_post_denoise( spec.corrupt_prompt, ca, target_token=spec.target_token, distractor_token=spec.distractor_token, clean_tgt=c_tgt, corrupt_tgt=r_tgt, clean_margin=m_c, corrupt_margin=m_r, ) nz = sweep_resid_post_noise( spec.clean_prompt, cc, target_token=spec.target_token, distractor_token=spec.distractor_token, clean_tgt=c_tgt, corrupt_tgt=r_tgt, clean_margin=m_c, corrupt_margin=m_r, ) best_den = degenerate_best_residual_competitive(den) best_nz = degenerate_best_residual_competitive(nz) if best_den is None: best_den = degenerate_best(den) if best_nz is None: best_nz = degenerate_best(nz) dn = best_den.normalized_effect nn = best_nz.normalized_effect dnm = best_den.normalized_margin_effect nnm = best_nz.normalized_margin_effect return BenchmarkQuickOutcome( spec=spec, clean_tgt=c_tgt, corrupt_tgt=r_tgt, margin_clean=m_c, margin_corr=m_r, top_den_layer=int(best_den.layer), top_den_norm=float(dn if not math.isnan(dn) else float("nan")), top_den_norm_margin=float(dnm if not math.isnan(dnm) else float("nan")), top_noise_layer=int(best_nz.layer), top_noise_norm=float(nn if not math.isnan(nn) else float("nan")), top_noise_norm_margin=float(nnm if not math.isnan(nnm) else float("nan")), thin_gap_warn=thin, inverted_gap=inverted, denoise_nt_nm_conflict=nt_nm_sign_conflict(dn, dnm), noise_nt_nm_conflict=nt_nm_sign_conflict(nn, nnm), ) def run_print_benchmark_suite() -> None: print( "\n=== MULTI-DOMAIN BENCHMARK HARNESS (residual hooks only — cheap cross-task lens) ===" ) print( "Each row scans hook_resid_post denoise/noise like the Primary France block; " "MLP/head sweeps stay on the canonical task to limit compute." ) print( f"Best-layer column uses competitive blocks only (drops L{final_residual_layer_idx():02d} hook_resid_post " f"— readout replay); thin-gap flag uses |clean_tgt−corrupt_tgt|<{MIN_ABS_TARGET_GAP_FOR_ROBUSTNESS}." ) rows_use: List[BenchmarkQuickOutcome] = [] for spec in BENCHMARK_SUITE: o = try_benchmark_quick(spec) if o: rows_use.append(o) hdr = ( f"{'benchmark_id':<22} {'domain':<18} {'dn_L':>5} {'d_nt':>6} {'d_nm':>6} " f"{'nz_L':>5} {'n_nt':>6} {'n_nm':>6} {'gap':>6} {'flg':^9}" ) print(hdr) print("-" * len(hdr)) for o in rows_use: dn_s = ( " nan " if math.isnan(o.top_den_norm) else f"{o.top_den_norm:6.3f}" ) dnm_s = ( " nan " if math.isnan(o.top_den_norm_margin) else f"{o.top_den_norm_margin:6.3f}" ) nz_s = ( " nan " if math.isnan(o.top_noise_norm) else f"{o.top_noise_norm:6.3f}" ) nnm_s = ( " nan " if math.isnan(o.top_noise_norm_margin) else f"{o.top_noise_norm_margin:6.3f}" ) gap = o.clean_tgt - o.corrupt_tgt flags = "" if o.inverted_gap: flags += "I" if o.thin_gap_warn: flags += "T" if o.denoise_nt_nm_conflict: flags += "d" if o.noise_nt_nm_conflict: flags += "m" flags = flags or "-" print( f"{o.spec.id:<22} {o.spec.domain:<18} {o.top_den_layer:5d} {dn_s} {dnm_s} " f"{o.top_noise_layer:5d} {nz_s} {nnm_s} {gap:+6.2f} {flags:^9}" ) print( "(flg: I=inverted tgt gap; T=|clean_tgt−corrupt_tgt| below robust threshold; " "d=denoise winner nt/nm opposite sign; m=noise winner nt/nm opposite sign. " "d_nt/n_nt=logit frac; d_nm/n_nm=margin frac (tgt−dist).)" ) eligible = [ r for r in rows_use if robustness_surface_eligible(r.clean_tgt, r.corrupt_tgt) ] if eligible: c_den = Counter(r.top_den_layer for r in eligible) c_nz = Counter(r.top_noise_layer for r in eligible) den_mode = sorted(c_den.items(), key=lambda x: (-x[1], x[0]))[:6] nz_mode = sorted(c_nz.items(), key=lambda x: (-x[1], x[0]))[:6] print( "\nAcross eligible benchmarks (invert/thin flagged rows excluded here), " "dominant-layer histograms:" ) lr = final_residual_layer_idx() print( f"(Winner layers exclude readout-bound L{lr:02d}; cross-task concentration in penultimate/noise bests " f"often matters more than spread denoise bests — asymmetry ⇒ distinct sufficiency vs necessity geometry.)" ) print(f" denoise bests: {den_mode}") print(f" noise bests: {nz_mode}") # --------------------------------------------------------- # Main experiment (default prompt pair) # --------------------------------------------------------- run_print_benchmark_suite() clean_logits, clean_activations = model.run_with_cache(CLEAN_PROMPT) corrupt_logits, corrupt_activations = model.run_with_cache(CORRUPT_PROMPT) ct_clean = _token_logit(clean_logits, TARGET_TOKEN) cd_clean = _token_logit(clean_logits, DISTRACTOR_TOKEN) ct_corrupt = _token_logit(corrupt_logits, TARGET_TOKEN) cd_corrupt = _token_logit(corrupt_logits, DISTRACTOR_TOKEN) marg_clean = ct_clean - cd_clean marg_corrupt = ct_corrupt - cd_corrupt denom_rest = ct_clean - ct_corrupt print("\n=== BASELINES (default pair) ===") print( f"Clean run: tgt {TARGET_TOKEN!r} logit={ct_clean:.4f}, " f"distractor {DISTRACTOR_TOKEN!r}={cd_clean:.4f}, margin={marg_clean:.4f}" ) print( f"Corrupt run: tgt={ct_corrupt:.4f}, distractor={cd_corrupt:.4f}, margin={marg_corrupt:.4f}" ) print( f"Target-gap (clean_tgt - corrupt_tgt) for normalization: {denom_rest:.4f}" ) print( f"Margin-gap (clean_margin - corrupt_margin) for contrast-normalized nm: " f"{marg_clean - marg_corrupt:.4f}" ) results_resid_den = sweep_resid_post_denoise( CORRUPT_PROMPT, clean_activations, target_token=TARGET_TOKEN, distractor_token=DISTRACTOR_TOKEN, clean_tgt=ct_clean, corrupt_tgt=ct_corrupt, clean_margin=marg_clean, corrupt_margin=marg_corrupt, ) print_ranked_residual_mlp( "DENoise — hook_resid_post (patch corrupt <- clean)", results_resid_den, 10, ) print( "\n(norm>1 is expected when the corrupt frame already pushes \"capital-like\" completions: " "clean activations grafted at readout stack with corrupt priming, so full Paris logit overshoot " "\"clean-alone\" baseline; treat 1.0 as illustrative, not a tight ceiling. " "nt=logit fraction; nm=margin fraction vs clean−corrupt margin.)" ) results_resid_noise = sweep_resid_post_noise( CLEAN_PROMPT, corrupt_activations, target_token=TARGET_TOKEN, distractor_token=DISTRACTOR_TOKEN, clean_tgt=ct_clean, corrupt_tgt=ct_corrupt, clean_margin=marg_clean, corrupt_margin=marg_corrupt, ) print_ranked_residual_mlp( "NOISE — hook_resid_post (patch clean <- corrupt; necessity-style)", results_resid_noise, 10, ) top_resid_den = degenerate_best_residual_competitive(results_resid_den) top_resid_noise = degenerate_best_residual_competitive(results_resid_noise) if top_resid_den is None: top_resid_den = degenerate_best(results_resid_den) if top_resid_noise is None: top_resid_noise = degenerate_best(results_resid_noise) print( "\n=== INTERPRETATION (defaults: where trajectories diverge) ===" ) print( "Maps localize causal **differences** between scripted clean vs corrupt forwards — " 'not "where Paris facts live"; wrong-answer / distractor coupling can dominate late sites.' ) print( "Sufficiency (denoise) vs necessity (noise) routinely split — e.g. best injector ≠ bottleneck.\n" "Best competitive blocker layers (omit final residual = readout replay bound):\n" f"L{top_resid_den.layer} leads denoise (nt={top_resid_den.normalized_effect:.3f}, " f"nm={top_resid_den.normalized_margin_effect:.3f}); " f"L{top_resid_noise.layer} leads noise toward corrupt (nt={top_resid_noise.normalized_effect:.3f}, " f"nm={top_resid_noise.normalized_margin_effect:.3f})." ) results_mlp = sweep_hook_mlp_out_denoise( CORRUPT_PROMPT, clean_activations, target_token=TARGET_TOKEN, distractor_token=DISTRACTOR_TOKEN, clean_tgt=ct_clean, corrupt_tgt=ct_corrupt, clean_margin=marg_clean, corrupt_margin=marg_corrupt, ) print_ranked_residual_mlp( "DENoise — hook_mlp_out", results_mlp, 10, annotate_readout_bound=False, ) results_heads = sweep_attn_hook_z_denoise( CORRUPT_PROMPT, clean_activations, target_token=TARGET_TOKEN, distractor_token=DISTRACTOR_TOKEN, clean_tgt=ct_clean, corrupt_tgt=ct_corrupt, clean_margin=marg_clean, corrupt_margin=marg_corrupt, ) print_ranked_heads( "DENoise — hook_z (single-head sweep limit)", results_heads, 15, ) top_head = max( results_heads, key=lambda r: r.normalized_effect if not math.isnan(r.normalized_effect) else float("-inf"), ) top_mlp = max( results_mlp, key=lambda r: r.normalized_effect if not math.isnan(r.normalized_effect) else float("-inf"), ) print("\n=== DEPTH SPLIT (hypothesis scaffold, not proof) ===") print( f"Early/low-mid peaks in MLP (e.g., L{top_mlp.layer} nt={top_mlp.normalized_effect:.3f}, " f"nm={top_mlp.normalized_margin_effect:.3f}) vs " f"later residual dominance (denoise L{top_resid_den.layer}) is consonant with " "\"retrieve / consolidate\" upstream vs contextual override nearer readout." ) print( f"Largest isolated single head L{top_head.layer} H{top_head.head} " f"nt={top_head.normalized_effect:.3f} nm={top_head.normalized_margin_effect:.3f}; " "multi-head combos are deliberately not searched here." ) comp_resid_den = excluding_readout_residual_rows(results_resid_den) ranked_resid = top_k_scores(comp_resid_den, min(len(comp_resid_den), model.cfg.n_layers)) if len(ranked_resid) >= 2: l1, l2 = ranked_resid[0].layer, ranked_resid[1].layer duo_logits = model.run_with_hooks( CORRUPT_PROMPT, fwd_hooks=[ ( f"blocks.{l1}.hook_resid_post", lambda act, *, hook, ca=clean_activations: patch_last_pos_inject_from_cache(act, hook, ca), ), ( f"blocks.{l2}.hook_resid_post", lambda act, *, hook, ca=clean_activations: patch_last_pos_inject_from_cache(act, hook, ca), ), ], ) dt, dd_, _, duo_norm, duo_nm = patch_metrics( duo_logits, TARGET_TOKEN, DISTRACTOR_TOKEN, ct_corrupt, ct_clean, marg_corrupt, marg_clean, direction="denoise", ) print("\n=== MULTI-SITE RESidual (defaults, joint top two competitive denoise layers) ===") print( f"Layers L{l1}, L{l2} jointly tgt={dt:.4f} margin={dt-dd_:.4f} " f"nt={duo_norm:.3f} nm={duo_nm:.3f}" ) print("(Compare to singles above.)") phl = int(top_head.layer) phh = int(top_head.head) pml = int(top_mlp.layer) pr_late = int(top_resid_den.layer) path_joint_logits = model.run_with_hooks( CORRUPT_PROMPT, fwd_hooks=[ ( f"blocks.{phl}.attn.hook_z", lambda act, *, hook, ca=clean_activations, h=phh: patch_last_pos_attn_z_head_from_cache( act, hook, ca, h, ), ), ( f"blocks.{pml}.hook_mlp_out", lambda act, *, hook, ca=clean_activations: patch_last_pos_inject_from_cache( act, hook, ca, ), ), ], ) pj_tgt, _pj_dd, pj_mr, pn_joint, pn_joint_m = patch_metrics( path_joint_logits, TARGET_TOKEN, DISTRACTOR_TOKEN, ct_corrupt, ct_clean, marg_corrupt, marg_clean, direction="denoise", ) path_washout_logits = model.run_with_hooks( CORRUPT_PROMPT, fwd_hooks=[ ( f"blocks.{phl}.attn.hook_z", lambda act, *, hook, ca=clean_activations, h=phh: patch_last_pos_attn_z_head_from_cache( act, hook, ca, h, ), ), ( f"blocks.{pml}.hook_mlp_out", lambda act, *, hook, ca=clean_activations: patch_last_pos_inject_from_cache( act, hook, ca, ), ), ( f"blocks.{pr_late}.hook_resid_post", lambda act, *, hook, ca=clean_activations: patch_last_pos_inject_from_cache( act, hook, ca, ), ), ], ) _, _, _, pn_wash, pn_wash_m = patch_metrics( path_washout_logits, TARGET_TOKEN, DISTRACTOR_TOKEN, ct_corrupt, ct_clean, marg_corrupt, marg_clean, direction="denoise", ) single_late_den = next(r for r in results_resid_den if r.layer == pr_late) print("\n=== PATH PATCH (donor head_z @ last pos + donor mlp_out — late residual NOT stacked) ===") print( f"L{phl}H{phh} z + L{pml} mlp_out jointly -> tgt={pj_tgt:.4f} margin={pj_mr:.4f} " f"nt={pn_joint:.3f} nm={pn_joint_m:.3f}. " "Later blocks run on this perturbed stream (not Wang frozen-attn path patching)." ) print( "\nContrast — stacking hook_resid_post at " f"L{pr_late} after z/mlp: full residual overwrite **erases** upstream grafts toward the donor run; " f"triple nt≈solo resid " f"(triple nt={pn_wash:.3f} nm={pn_wash_m:.3f} vs solo L{pr_late} " f"nt={single_late_den.normalized_effect:.3f} nm={single_late_den.normalized_margin_effect:.3f})." ) frontier_lab.print_mediation_scaffold( path_joint_nt=pn_joint, path_joint_nm=pn_joint_m, head_layer=phl, head_idx=phh, mlp_layer=pml, resid_denoise_layer=int(top_resid_den.layer), resid_noise_layer=int(top_resid_noise.layer), ko_layer=int(top_resid_den.layer), wash_triple_nt=float(pn_wash), solo_resid_nt=float(single_late_den.normalized_effect), ) KO_LAYER = top_resid_den.layer ablate_hook_name = f"blocks.{KO_LAYER}.hook_resid_post" ko_only_logits = model.run_with_hooks( CORRUPT_PROMPT, fwd_hooks=[(ablate_hook_name, zero_last_pos_2d)], ) ko_tgt = _token_logit(ko_only_logits, TARGET_TOKEN) ko_dist = _token_logit(ko_only_logits, DISTRACTOR_TOKEN) print("\n=== LOCALIZED RESIDUAL KNOCKOUT (default corrupt prompt) ===") print(f"Zero final-pos residual after block {KO_LAYER}.") print( f"Corrupt tgt {ct_corrupt:.4f} -> {ko_tgt:.4f}; margins {marg_corrupt:.4f} -> {(ko_tgt-ko_dist):.4f}" ) print( "If Paris logits fall but Rome falls faster here, you've removed a corridor that was loudly " "coupled to *wrong-city* continuation under this corrupt scaffold — downstream layers can " "still salvage Paris (seen in post-knockout sweeps), so treat as redundancy, not a single bottleneck." ) def resid_falsification_sweep( v: FalsificationPrompt, *, ablate_hook_resid_layer: int | None = None, ) -> Tuple[float, float, float, float, List[PatchResult]]: """Returns clean_tgt, corrupt_tgt, denom, corrupt_margin, results.""" cl, ca = model.run_with_cache(v.clean) crl, _cc = model.run_with_cache(v.corrupt) c_tgt = _token_logit(cl, v.target_token) r_tgt = _token_logit(crl, v.target_token) c_dst = _token_logit(cl, v.distractor_token) r_dst = _token_logit(crl, v.distractor_token) marg_c_clean = c_tgt - c_dst marg_corrupt_fwd = r_tgt - r_dst denom = c_tgt - r_tgt outs: List[PatchResult] = [] base_hooks = [] if ablate_hook_resid_layer is not None: ablate_name = f"blocks.{ablate_hook_resid_layer}.hook_resid_post" base_hooks.append((ablate_name, zero_last_pos_2d)) for layer in range(model.cfg.n_layers): if ablate_hook_resid_layer is not None and layer == ablate_hook_resid_layer: continue hook_name = f"blocks.{layer}.hook_resid_post" fwd_hooks = list(base_hooks) + [ ( hook_name, lambda tens, *, hook, act_cache=ca: patch_last_pos_inject_from_cache( tens, hook, act_cache, ), ), ] patched = model.run_with_hooks(v.corrupt, fwd_hooks=fwd_hooks) tg, dd, marg, norm, norm_m = patch_metrics( patched, v.target_token, v.distractor_token, r_tgt, c_tgt, marg_corrupt_fwd, marg_c_clean, direction="denoise", ) outs.append( PatchResult( site="hook_resid_post_den_falsify", layer=layer, target_logit=tg, distractor_logit=dd, logit_margin=marg, normalized_effect=norm, normalized_margin_effect=norm_m, ), ) corrupt_margin = marg_corrupt_fwd return c_tgt, r_tgt, denom, corrupt_margin, outs results_by_label: Dict[str, List[PatchResult]] = {} surface_targets: Dict[str, Tuple[float, float]] = {} print("\n=== FALSIFICATION — denoise residuals across surfaces ===") print( "Per-row norms get noisy when |clean_tgt−corrupt_tgt| is tiny; pooled Jaccard excludes those surfaces. " "Per-surface bests and top‑3 omit final hook_resid_post (readout replay bound). " "nt=logit frac vs tgt gap; nm=margin frac vs (clean−corr) tgt−dist margin." ) for v in PROMPT_VARIANTS: c_tgt, r_tgt, denom, corr_marg, pr = resid_falsification_sweep(v) results_by_label[v.label] = pr surface_targets[v.label] = (c_tgt, r_tgt) best = degenerate_best_residual_competitive(pr) or degenerate_best(pr) pr_rank = excluding_readout_residual_rows(pr) pool = pr_rank if pr_rank else pr top3 = top_k_scores(pool, min(3, len(pool))) layers = ", ".join(format_layer_effect_pair(r, decimals=2) for r in top3) denoise_norm_degenerate = r_tgt >= c_tgt - 1e-9 tiny_gap = abs(c_tgt - r_tgt) < MIN_ABS_TARGET_GAP_FOR_ROBUSTNESS print( f"\n[{v.label}] distractor={v.distractor_token!r} " f"clean_tgt={c_tgt:.3f} corrupt_tgt={r_tgt:.3f} gap={denom:.3f} corrupt_margin={corr_marg:.3f}" + ( " [corrupt tgt >= clean — denoise norm disabled]" if denoise_norm_degenerate else (" [|gap| small — noisy norm]" if tiny_gap else "") ) + "\n" f" best L{best.layer} nt={best.normalized_effect:.3f} nm={best.normalized_margin_effect:.3f} " f"tgt={best.target_logit:.3f} margin={best.logit_margin:.3f} top3: {layers}" ) print_explanation_robustness( results_by_label, k=5, surface_targets=surface_targets, min_abs_gap=MIN_ABS_TARGET_GAP_FOR_ROBUSTNESS, ) _, _, _, _, pr_after_ko = resid_falsification_sweep( FalsificationPrompt( CLEAN_PROMPT, CORRUPT_PROMPT, TARGET_TOKEN, DISTRACTOR_TOKEN, label="defaults", ), ablate_hook_resid_layer=KO_LAYER, ) best_after = degenerate_best_residual_competitive(pr_after_ko) or degenerate_best(pr_after_ko) rank_pool = excluding_readout_residual_rows(pr_after_ko) or list(pr_after_ko) sorted_after = sorted( rank_pool, key=lambda r: ( r.normalized_effect if not math.isnan(r.normalized_effect) else float("-inf") ), reverse=True, ) def _floored(row: PatchResult, floor_tgt: float) -> bool: return row.target_logit > floor_tgt + 0.05 above_floor_simple = [r for r in sorted_after if _floored(r, ko_tgt)] print("\n=== KNOCKOUT + RE-SWEEP (defaults; knockout still applied) ===") print( f"Knockout L{KO_LAYER}; best survivor L{best_after.layer} " f"nt={best_after.normalized_effect:.3f} nm={best_after.normalized_margin_effect:.3f} " f"tgt={best_after.target_logit:.3f}" ) print( "Above knockout tgt floor (~+0.05): " + ", ".join( f"L{r.layer}(t={r.target_logit:.2f})" for r in above_floor_simple[:8] ) ) print("(Downstream patches reinject signal after earlier hook zeroing.)") print("\n=== AUTOMATED TEMPLATE STRESS (mutants vs default scaffolding) ===") for stress_lab, cp_stress, corp_stress in AUTO_STRESS_PAIRS: v_stress = FalsificationPrompt(cp_stress, corp_stress, TARGET_TOKEN, DISTRACTOR_TOKEN, label=stress_lab) s_ct, s_rt, _, _s_cm, spr = resid_falsification_sweep(v_stress) sb = degenerate_best_residual_competitive(spr) or degenerate_best(spr) tag = "" if s_rt >= s_ct - 1e-9: tag = " [inverted tgt gap]" elif abs(s_ct - s_rt) < MIN_ABS_TARGET_GAP_FOR_ROBUSTNESS: tag = " [|gap| None: print("\n=== WHAT QUESTION THIS STACK ANSWERS ===") print( "Operational question: causal **circuit discovery anchored to one scripted distracting surface**, " "then **stress-tested across paraphrases + light template mutation** — task-local mediation, " "not an invariant Paris module." ) print( "Benchmark harness: declarative prompts live in benchmark_specs.py; extend BENCHMARK_SUITE for " "new domains while keeping residual sweeps bounded—switch on full MLP/head ladders per task if compute allows." ) print( "Frontier kit (PG_FRONTIER=1): subspace Δ-patch + random orthonormal baseline, holdout pairs, " "pairwise |cos Δ| across surfaces, toy SAE latent transfer, fingerprint attack w/ patch metrics, " "resid ranking + reroute when blocking denoise site; tune PG_RANDOM_SUBSPACE_TRIALS, " "PG_SUBSPACE_HOLDOUT_PAIRS, PG_FRONTIER_PAIRWISE_COS, PG_SAE_STEPS / PG_ATTACK_TRIALS." ) print( "Milestone (PG_MILESTONE=1 or PG_FROZEN_ATTN=1 / PG_SAELENS=1): frozen clean→corrupt attn.hook_pattern " "graft; SAELens top-k xfer at pretrained hook (PG_SAE_RELEASE+PG_SAE_ID or PG_SAE_DISK)." ) print_method_note() _SCRAF = os.environ.get("PG_FRONTIER", "").strip().lower() if _SCRAF in ("1", "true", "yes", "on", "all"): _pairs = [(v.clean, v.corrupt) for v in PROMPT_VARIANTS[:8]] frontier_lab.run_frontier_suite( model, clean_prompt=CLEAN_PROMPT, corrupt_prompt=CORRUPT_PROMPT, target_token=TARGET_TOKEN, distractor_token=DISTRACTOR_TOKEN, clean_act=clean_activations, corr_act=corrupt_activations, clean_tgt=ct_clean, corrupt_tgt=ct_corrupt, clean_margin=marg_clean, corrupt_margin=marg_corrupt, exclude_readout_layer=final_residual_layer_idx(), focal_resid_layer=int(top_resid_den.layer), prompt_pairs_for_contrast=_pairs, path_joint_nt=float(pn_joint), path_joint_nm=float(pn_joint_m), wash_triple_nt=float(pn_wash), solo_resid_nt=float(single_late_den.normalized_effect), reroute_exclude_layer=int(top_resid_den.layer), ) if milestone_interp.milestone_any_env(): milestone_interp.run_milestone_bundle( model, corrupt_prompt=CORRUPT_PROMPT, target_token=TARGET_TOKEN, distractor_token=DISTRACTOR_TOKEN, clean_act=clean_activations, corr_act=corrupt_activations, clean_tgt=ct_clean, corrupt_tgt=ct_corrupt, clean_margin=marg_clean, corrupt_margin=marg_corrupt, default_attn_layer_for_frozen_pattern=int(phl), )