Read one GPT-2 neuron's top activations with TransformerLens
Pull the contexts where one GPT-2 small neuron fires hardest, then test your label with activation bands and your own sentences. Covers TransformerLens 4.0.
Before you name a neuron, check the top contexts, the lower activation bands, and sentences you write yourself. Everything runs on CPU.
[1] Install CPU-only torch and TransformerLens in a virtual environment. Pull torch from the CPU index first so pip doesn't fetch the large GPU build.
$ python3 -m venv .venv && . .venv/bin/activate$ pip install torch --index-url https://download.pytorch.org/whl/cpu$ pip install transformer_lens$ pip list | grep -iE "^(torch|transformer-lens|transformers|datasets) "datasets 5.1.0torch 2.14.1+cputransformer-lens 4.2.0transformers 5.20.2[2] HookedTransformer, which older tutorials use, was removed in 4.0. Importing it fails like this:
$ python -c "from transformer_lens import HookedTransformer"ImportError: 'HookedTransformer' was removed in TransformerLens 4.0. Use TransformerBridge.boot_transformers(name), then enable_compatibility_mode() for HookedTransformer-equivalent numerics. See docs/source/content/migrating_to_v4.md.[3] Create the script. It runs layer 5, MLP neuron 1512 over the first 300 documents of the public dataset NeelNanda/pile-10k (128 tokens each) and keeps only that neuron's values. Keeping all 3,072 neurons uses much more memory.
import sys, torchfrom datasets import load_datasetfrom transformer_lens.model_bridge import TransformerBridge LAYER, NEURON = 5, 1512N_DOCS, CTX = 300, 128HOOK = f"blocks.{LAYER}.mlp.hook_post" torch.set_grad_enabled(False)model = TransformerBridge.boot_transformers("gpt2", device="cpu")model.enable_compatibility_mode() def ctx(toks, i, left=8, right=2): s = model.to_string(toks[max(0, i - left):i]) w = model.to_string(toks[i:i + 1]) r = model.to_string(toks[i + 1:i + 1 + right]) return (s + "[[" + w + "]]" + r).replace("\n", " ") def collect(): ds = load_dataset("NeelNanda/pile-10k", split="train") acts, toks = [], [] for d in range(N_DOCS): t = model.to_tokens(ds[d]["text"])[:, :CTX] _, cache = model.run_with_cache(t, names_filter=HOOK) acts.append(cache[HOOK][0, :, NEURON].clone()) toks.append(t[0]) return torch.cat(acts), torch.cat(toks) cmd = sys.argv[1]if cmd in ("top", "bands"): A, T = collect() print(f"tokens={len(A)} max={A.max():.2f} mean={A.mean():.3f} >1.0={(A > 1.0).sum().item()}") if cmd == "top": v, i = A.topk(20) for a, j in zip(v, i): print(f"{a:5.2f} {ctx(T, j)}") else: top = A.max().item() g = torch.Generator().manual_seed(0) for lo, hi in [(0.75, 1.0), (0.5, 0.75), (0.25, 0.5)]: idx = ((A >= lo * top) & (A < hi * top)).nonzero().flatten() print(f"--- {lo:.0%}~{hi:.0%} of max: {len(idx)} tokens, 6 random") for j in idx[torch.randperm(len(idx), generator=g)[:6]]: print(f"{A[j]:5.2f} {ctx(T, j)}")elif cmd == "probe": for s in sys.argv[2:]: t = model.to_tokens(s) _, cache = model.run_with_cache(t, names_filter=HOOK) a = cache[HOOK][0, :, NEURON] print(" ".join(f"{model.to_string(t[0, k:k+1]).strip() or '·'}:{a[k]:.2f}" for k in range(1, t.shape[1])))[4] Print the top 20 contexts. The token inside [[ ]] is where the neuron fired.
$ python neuron.py toptokens=36744 max=3.04 mean=-0.068 >1.0=122 3.04 show appears to be paddling in the[[ wrong]] direction as 2.68 cloud<|endoftext|>Q: Query returns[[ wrong]] result 2.51 rewarded in an environment of growing corruption and[[ mis]]management. 2.49 hypothalamic lesions alter intake of an im[[balanced]] amino acid 2.48 in an environment of growing corruption and mis[[management]]. 2.43 that, it rubbed The Associated Press the[[ wrong]] way simply 2.37 over internet (maybe i'm using the[[ wrong]] key word 2.33 owing be rewarded in an environment of growing[[ corruption]] and mis 2.25 misinterpret things, miss things, draw the[[ wrong]] conclusions, 2.23 price. Through cries of policy paralysis and[[ corruption]], Man 2.16 SCA), noted that international cases of[[ misconduct]] in wholesale 2.15 E results inaccurate/obsolete/inc[[orrect]]/out 2.08 of that money after learning of the federal[[ corruption]] charges lodged 2.03 ’ (16-22) recent[[ poor]] play minus 2.00 s emergency-management agency sent out a[[ false]] alert warning 1.95 stand to benefit from lower tax brackets,[[ higher]] estate tax 1.95 on-facilitated technique, which[[ greatly]] simplified sequencing 1.91 a little bit of current - this avoids[[ excessive]] cross over 1.81 /obsolete/incorrect/out[[dated]]? 1.81 hadex G<|endoftext|>Because of the[[ extreme]] temperatures of[5] Split the same data into bands relative to the maximum and sample at random. The top 20 is only part of the highest band.
$ python neuron.py bandstokens=36744 max=3.04 mean=-0.068 >1.0=122--- 75%~100% of max: 7 tokens, 6 random 2.37 over internet (maybe i'm using the[[ wrong]] key word 2.33 owing be rewarded in an environment of growing[[ corruption]] and mis 2.43 that, it rubbed The Associated Press the[[ wrong]] way simply 2.49 hypothalamic lesions alter intake of an im[[balanced]] amino acid 2.48 in an environment of growing corruption and mis[[management]]. 2.68 cloud<|endoftext|>Q: Query returns[[ wrong]] result --- 50%~75% of max: 32 tokens, 6 random 2.23 price. Through cries of policy paralysis and[[ corruption]], Man 1.60 Please enter an email address Email address is[[ invalid]] Fill out 1.72 industry is facing the immense challenge of processing[[ huge]] amounts of 1.95 stand to benefit from lower tax brackets,[[ higher]] estate tax 1.68 However, there is considerable phenotypic[[ variability]], extending 1.58 "off" as it has a very[[ low]] number of--- 25%~50% of max: 164 tokens, 6 random 0.90 of the 5 experiments, they show that[[ irrelevant]] sound ( 0.87 RowView, and datagridview[[ only]] displays values 1.10 thought it meant this character tends to do[[ more]] than he 0.91 time. Penalty relief will[[ only]] apply to 0.96 year deal. There is still[[ little]] to no 0.95 immense challenge of processing huge amounts of documents[[ efficiently]]. PredThe top band (7 tokens) is wrongness and misconduct. The 50–75% band shows size and degree words such as huge, higher and low. The 25–50% band (164 tokens) shows only, more and little. A "wrongness neuron" label fits only the top band, and the lower bands hold far more tokens.
[6] Test the label with sentences you write. Add an antonym and a different use of the same word as controls.
$ python neuron.py probe "You took the wrong exit." "You took the right exit." "The audit found fraud and mismanagement." "The answer was incorrect." "The water was extremely cold." "Management approved the plan."You:-0.03 took:-0.13 the:-0.11 wrong:2.83 exit:-0.12 .:-0.17You:-0.03 took:-0.13 the:-0.11 right:0.61 exit:-0.14 .:-0.15The:-0.09 audit:-0.13 found:-0.17 fraud:1.65 and:-0.16 mis:2.55 management:2.15 .:-0.16The:-0.09 answer:-0.13 was:-0.10 incorrect:2.13 .:-0.17The:-0.09 water:-0.10 was:-0.08 extremely:0.30 cold:0.39 .:-0.06Management:-0.13 approved:-0.16 the:-0.08 plan:-0.10 .:-0.06[7] Record the result in one line: layer and neuron, data and token count, representative contexts per band, and the control values. This run: "L5N1512 · pile-10k 300 docs, 36,744 tokens · top = wrongness/misconduct, middle = size/degree, lower = quantity limiters · right 0.61".
References: TransformerLens · Migrating to TransformerLens 4.0 · NeelNanda/pile-10k · PyTorch install