File size: 6,678 Bytes
253d988 | 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 | """
Random-feature ablation control for the v9c SAE induction-feature ablation result.
The headline finding is: zeroing the top-50 induction-scoring features drops top-1
ICL accuracy from 57.75% -> 47.65% (-10.1pp). Reviewers will fairly ask: what if
50 *random* features did the same thing?
This script:
1. Loads v9c SAE + Gemma-2-2B.
2. Samples 50 *non-induction-cluster* features uniformly at random, for each of
5 random seeds.
3. Runs the same feature ablation as ablations.py on each random sample.
4. Reports mean accuracy drop across the 5 random samples vs the 10.1pp top-50
induction drop.
Output: results/random_feature_ablation.json.
"""
import argparse
import json
import random
import time
from functools import partial
from pathlib import Path
import numpy as np
import pandas as pd
import torch
from sae_gemma.model_utils import load_model
from sae_gemma.paths import (
HOOK_NAME,
INDUCTION_PROBES_PATH,
RESULTS_DIR,
SAE_MAIN_DIR,
REPO_ROOT,
)
CANDIDATE_IDS_PATH = RESULTS_DIR / "induction_candidate_ids.json"
ABLATION_RESULTS_PATH = RESULTS_DIR / "ablation_results.json"
def load_sae_local(sae_path: Path, device: str):
from sae_lens.saes.sae import SAE
sae = SAE.load_from_disk(str(sae_path), device=device)
sae.eval()
return sae
@torch.no_grad()
def measure_icl_accuracy(
model,
sae,
probes_df: pd.DataFrame,
feature_ids_to_ablate: list[int],
device: str,
batch_size: int = 16,
) -> float:
"""Run model with SAE-patch + ablation; return top-1 accuracy on induction probes."""
feature_mask = torch.zeros(sae.cfg.d_sae, dtype=torch.bool, device=device)
if feature_ids_to_ablate:
feature_mask[feature_ids_to_ablate] = True
correct = 0
total = 0
probe_tokens = [torch.tensor(t, dtype=torch.long, device=device) for t in probes_df["tokens"]]
answer_tokens = probes_df["B"].tolist()
for i in range(0, len(probe_tokens), batch_size):
batch = probe_tokens[i: i + batch_size]
max_len = max(len(seq) for seq in batch)
padded = torch.zeros(len(batch), max_len, dtype=torch.long, device=device)
seq_lens = []
for j, seq in enumerate(batch):
padded[j, : len(seq)] = seq
seq_lens.append(len(seq))
def patch_hook(value, hook):
# value: [batch, seq, d_model]
b, s, d = value.shape
flat = value.reshape(b * s, d).float()
z = sae.encode(flat) # [b*s, d_sae]
z_ablated = z * (~feature_mask).float() # zero target features
recon = sae.decode(z_ablated) # [b*s, d_model]
recon_orig = sae.decode(z)
delta = (recon - recon_orig).reshape(b, s, d).to(value.dtype)
return value + delta
logits = model.run_with_hooks(padded, fwd_hooks=[(HOOK_NAME, patch_hook)])
for j in range(len(batch)):
final_pos = seq_lens[j] - 1
pred = logits[j, final_pos].argmax().item()
if pred == answer_tokens[i + j]:
correct += 1
total += 1
return correct / total
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--n-ablate", type=int, default=50)
parser.add_argument("--n-seeds", type=int, default=5)
parser.add_argument("--seed-base", type=int, default=10000)
parser.add_argument("--batch-size", type=int, default=16)
parser.add_argument("--device", default="cuda")
args = parser.parse_args()
try:
from dotenv import load_dotenv
load_dotenv(REPO_ROOT / ".env")
except ImportError:
pass
print("[random_ab] Loading model + SAE ...", flush=True)
model = load_model(device=args.device)
sae = load_sae_local(SAE_MAIN_DIR, args.device)
d_sae = sae.cfg.d_sae
print(f"[random_ab] Loading induction probes + candidate IDs ...", flush=True)
probes_df = pd.read_parquet(INDUCTION_PROBES_PATH)
candidate_ids = set(json.loads(CANDIDATE_IDS_PATH.read_text(encoding="utf-8")))
non_induction = [i for i in range(d_sae) if i not in candidate_ids]
print(f"[random_ab] {len(non_induction)} non-induction features available", flush=True)
# Reference: baseline accuracy + top-50 induction ablation from ablation_results.json
ablation_results = json.loads(ABLATION_RESULTS_PATH.read_text(encoding="utf-8"))
baseline_acc = ablation_results["baseline_accuracy"]
induction_ns = ablation_results["feature_ablation"]["N"]
induction_accs = ablation_results["feature_ablation"]["accuracy"]
print(f"[random_ab] Reference baseline: {baseline_acc:.4f}", flush=True)
print(f"[random_ab] Reference induction ablation: N={induction_ns} acc={induction_accs}", flush=True)
# Run random ablation
results = []
t0 = time.monotonic()
for s in range(args.n_seeds):
seed = args.seed_base + s
rng = random.Random(seed)
sample = rng.sample(non_induction, args.n_ablate)
print(f"[random_ab] Seed {seed}: ablating {args.n_ablate} random non-induction features ...", flush=True)
acc = measure_icl_accuracy(model, sae, probes_df, sample, args.device, args.batch_size)
drop = baseline_acc - acc
elapsed = time.monotonic() - t0
print(f"[random_ab] acc={acc:.4f} drop={drop:+.4f} ({drop*100:+.2f}pp) elapsed={elapsed/60:.1f}m", flush=True)
results.append({"seed": seed, "n_ablated": args.n_ablate, "feature_ids": sample[:10], "accuracy": acc, "drop": drop})
out = {
"baseline_accuracy": baseline_acc,
"induction_top50_accuracy": induction_accs[-1] if induction_accs else None,
"induction_top50_drop": baseline_acc - (induction_accs[-1] if induction_accs else 0.0),
"random_n_ablated": args.n_ablate,
"random_n_seeds": args.n_seeds,
"random_runs": results,
"random_mean_acc": float(np.mean([r["accuracy"] for r in results])),
"random_mean_drop": float(np.mean([r["drop"] for r in results])),
"random_std_drop": float(np.std([r["drop"] for r in results])),
}
out_path = RESULTS_DIR / "random_feature_ablation.json"
out_path.write_text(json.dumps(out, indent=2), encoding="utf-8")
print(f"\n[random_ab] === RESULTS ===", flush=True)
print(f" Top-50 induction features ablated: drop = {out['induction_top50_drop']*100:.2f}pp", flush=True)
print(f" 50 RANDOM features (mean of {args.n_seeds} seeds): drop = {out['random_mean_drop']*100:.2f}pp ± {out['random_std_drop']*100:.2f}pp", flush=True)
print(f" Saved to {out_path}", flush=True)
if __name__ == "__main__":
main()
|