diff options
| -rw-r--r-- | README.md | 6 | ||||
| -rw-r--r-- | docs/blog-jlens-frequency.md | 77 | ||||
| -rw-r--r-- | results.md | 39 | ||||
| -rw-r--r-- | src/wu_row_norm_check.py | 189 |
4 files changed, 290 insertions, 21 deletions
@@ -20,6 +20,11 @@ Workspace in Language Models"* (2026, 4. A causal loss-reweighting test (2x loss weight on 'q' targets vs two controls) tests whether effective frequency causally demotes a token's J-lens norm. See `results.md` for the latest numbers. +5. The geometric half generalizes: GPT-2's unembedding row norms (V = 50,257, + wte == tied lm_head) anti-correlate with token log-frequency + (r ≈ -0.45/-0.49, gpt2/gpt2-medium; monotone across frequency deciles) — + the learned W_U geometry is not a 65-char vocabulary artifact + (`src/wu_row_norm_check.py`, CPU-only). See `docs/blog-jlens-frequency.md` for the write-up and `results.md` for the numbers. The three independent adversarial reviews that shaped the project @@ -44,6 +49,7 @@ src/jlens_v3.py FAITHFUL J-lens: rows of W_U * J_l (canonical) src/synthetic_pair.py frequency-matched synthetic pair experiment src/loss_reweight.py causal loss-reweighting experiment src/gpt2_jlens.py GPT-2 scale test (under-powered; see results.md) +src/wu_row_norm_check.py at-scale W_U row-norm check (CPU-only; results.md §5) tests/ unit tests (see scripts/test.sh) scripts/test.sh canonical test command docs/blog-jlens-frequency.md write-up (Feynman-style) diff --git a/docs/blog-jlens-frequency.md b/docs/blog-jlens-frequency.md index 8ab2e46..f5a459e 100644 --- a/docs/blog-jlens-frequency.md +++ b/docs/blog-jlens-frequency.md @@ -1,4 +1,4 @@ -# What the Jacobian Lens Actually Measures +# What the Jacobian Lens Measures ### A small replication of Anthropic's J-lens, the token-frequency confound we found, and the bug we almost published *This is a story about trying to look inside a language model. We found something @@ -11,18 +11,22 @@ stronger result.* **The short version.** We reimplemented Anthropic's Jacobian lens faithfully (verified against their released code) and found that the J-space ranking is strongly confounded by token frequency: rare tokens score high, common tokens -score low (r ≈ -0.6 to -0.7 at every layer, p ~ 10^-9 or less). Anthropic never +score low (r ≈ -0.6 to -0.7 at every layer, p ~ 10^-9 or less; against +log-frequency — the natural scale for Zipfian data — the unembedding geometry +alone hits r = -0.69). Anthropic never controls for frequency — not in the paper, not in the released code. Digging into *why* gave us the most interesting result: the frequency signal splits into two separable parts. Half lives in the static geometry of the model's word-scoring matrix — baked into the lens by definition, so any user inherits it. A smaller, layer-dependent part lives in what the layers themselves do, and -vanishes at the final layer. A frequency-matched synthetic pair shows the lens -also carries genuine structure signal (with a caveat we're resolving), and a -causal test found the demotion effect is small under loss reweighting. We are -**not** claiming the J-space doesn't exist. We're claiming that any -"privileged subspace" interpretation needs a frequency control first. The full -story — numbers, mistakes, and all — is below. +vanishes at the final layer. The static half survives at real scale: GPT-2's +unembedding rows anti-correlate with token log-frequency too (r ≈ -0.45/-0.49, +V = 50,257, n = 46,887). A frequency-matched synthetic pair shows the lens +also carries genuine structure signal (confirmed by a clean-boundary control at +~1.3x), and a causal test found the demotion effect is small under loss +reweighting. We are **not** claiming the J-space doesn't exist. We're claiming +that any "privileged subspace" interpretation needs a frequency control first. +The full story — numbers, mistakes, and all — is below. --- @@ -193,7 +197,9 @@ think, the actual novel mechanistic contribution of this project: **the frequency confound is not one thing.** 1. **Static geometry — the ruler.** The unembedding row norms ||W_U[k]|| - themselves anti-correlate with frequency (r = -0.61, Spearman -0.81). + themselves anti-correlate with frequency — most strongly against + log-frequency, the natural scale for Zipfian data: r(||W_U[k]||, log10 f) = + -0.69 (raw frequency -0.61, Spearman -0.81). Rare tokens get bigger rows in the word-scoring matrix. Since the lens reads through W_U by definition, any user of the lens — including Anthropic's capacity analysis — inherits this bias automatically. A @@ -365,31 +371,64 @@ Let us be very careful here, because it would be easy to overclaim. subtract mean concept directions. Those are real experimental controls — but none of them is a token-frequency control, which is the specific thing our claim is about. -- We are **not** saying the lens is useless. The synthetic-pair result shows it +- We are **not** claiming the lens is useless. The synthetic-pair result shows it carries real structure signal. +- We are **not** claiming that ranking by lens-vector *norm* is the same as + ranking by *lens output on real activations*. Our numbers rank tokens by the + norm of their faithful J-lens vector — a summary of the readout geometry — + not by how strongly, or how often, those directions actually fire in running + text. Anthropic's capacity claim is about the latter (occupancy). The norm + still matters: it is the size of a token's lever on the readout, and the W_U + geometry that drives it enters any readout-based analysis by definition. But + the gap between "geometry is frequency-confounded" and "the capacity claim + is frequency-confounded" is real, and it is the specific gap an at-scale + occupancy test has to close. - We are **not** saying "it's just linear algebra." Our toy models don't show the compression Anthropic sees in large models; that's a limitation of toy models, not evidence against large ones. What we **are** saying is narrower and, we think, more durable: on the paper's -own measurement, J-lens *rankings* are strongly confounded by token frequency -at every scale we can test, and frequency is a variable any J-lens analysis -should control for. Whether the confound survives at Anthropic's scale is an -empirical question — one we are taking to bigger models next. +own measurement, J-lens *norm-rankings* are strongly confounded by token +frequency at every scale we can test, and frequency is a variable any J-lens +analysis should control for. Whether the confound survives at Anthropic's scale +is an empirical question — one we are taking to bigger models next, and one we +already have a first, partial answer to for the geometric half (Section 10). ## 10. What's next +We did try bigger once already, and we owe you the number, because a reader +who opens the repo will find it either way: an early probe on GPT-2 small +(`src/gpt2_jlens.py`) returned an average correlation of only r ≈ -0.18 across +layers. We do not count it as evidence, for three concrete reasons: it sampled +96 token positions out of a 50,257-token vocabulary; it averaged over only 100 +sampled tokens per batch; and it measured a subtly different quantity +(norm-per-batch rather than norm-of-the-mean). It was a directional probe, and +it pointed weak. It is logged in `results.md`, flagged do-not-cite — but a +post that promises "bigger models next" should not pretend the attempt never +happened. + Toy scale answers the methodological question. Scale answers the real one. We want to run the faithful lens on a real language model (V = 50K, d = 768 — the regime where Anthropic's claims live) with proper statistical power, and to run the occupancy test their capacity claim is actually about. That's the next post. -The learned-geometry finding makes one piece of that cheap and decisive: if -the W_U row-norm anti-correlation is a general property of softmax-output -models trained on Zipfian data, it should appear in GPT-2's unembedding matrix -directly — no Jacobian computation at all. One line of code at scale, and it -either generalizes or it doesn't. +The learned-geometry finding makes one piece of that cheaper than the probe +was: if the W_U row-norm anti-correlation is a general property of +softmax-output models trained on Zipfian data, it should appear in GPT-2's +unembedding matrix directly — no Jacobian computation at all. So we ran it: +GPT-2's unembedding row norms correlate with token log-frequency at V = 50,257 +(r ≈ -0.45 on gpt2-small, -0.49 on gpt2-medium, n = 46,887 tokens seen in +wikitext-103; Spearman -0.46 to -0.50). The decile picture is monotone in both +models — the rarest tenth of tokens have the biggest rows — and the top-norm +tokens are rare ones (SPONSORED, soDeliveryDate, Reviewer). The geometric half +of the confound is not a 65-character vocabulary artifact: it is present in +two real, independently trained language models at Anthropic's vocabulary +scale. One honest caveat: it is weaker than at toy scale (-0.45/-0.49 vs +-0.69), so the char-scale result was not a size artifact, but neither is the +at-scale effect as strong as the toy suggested. The layer-dynamics half — +what the layers themselves add beyond the ruler — still needs the faithful +lens at scale, and that remains the next post. ## 11. How to reproduce everything @@ -182,6 +182,40 @@ Reading: 4. Absolute norms differ across experiments (base L3 q = 0.0109 vs these 0.015-0.018) — only within-experiment comparisons are meaningful. +## 5. At-scale W_U row-norm check (GPT-2, V=50,257) — RESOLVED + +The geometric half of the confound generalizes. Correlating GPT-2 unembedding +row norms (wte == tied lm_head, so this is exactly W_U's rows) against token +unigram log-frequencies from wikitext-103-raw (117.9M tokens; 46,887/50,257 +vocab covered): + +``` + model r(raw) r(log10, all V) r(log10, seen) Spearman(seen) + gpt2 (d=768) -0.092 -0.451 -0.468 -0.462 + gpt2-medium -0.104 -0.493 -0.515 -0.498 +``` + +n = 46,887 seen tokens; 95% CIs [-0.475, -0.461] (gpt2) and +[-0.522, -0.508] (gpt2-medium). Log10-freq decile -> mean row norm is +monotone in both models (gpt2: 4.02 -> 3.42; medium: 3.93 -> 3.16): rare +tokens get bigger unembedding rows. Top-norm tokens are rare or zero-count +(SPONSORED, soDeliveryDate, Reviewer, advertisement). + +Reading: the learned W_U geometry finding from the 65-char model carries to +real LMs at V = 50,257 — present in two independently trained models, monotone +across the whole vocabulary — but it is WEAKER than at toy scale +(r(log10) -0.45/-0.49 vs -0.69; Spearman -0.46/-0.50 vs -0.81). So the +char-scale result was not a vocabulary-size artifact, and neither is the +at-scale effect as strong as the toy suggested. The layer-dynamics component +(partial r after regressing out W_U) at scale remains untested — that needs +the faithful J-lens on GPT-2 (future work). + +Reproduce, CPU-only, no training, no GPU: +``` +.venv/bin/python src/wu_row_norm_check.py --models gpt2,gpt2-medium +``` +Full log: `outputs/wu_row_norm_check.log`. + ## 4. Historical / do-not-copy - Original proxy finding (r = -0.65, `jlens_v2`): superseded by the faithful @@ -190,5 +224,6 @@ Reading: (`gpt2_jlens.py`) gave avg r = -0.18 across layers — UNDER-POWERED (96 token positions, n=100 sampled tokens) and a different quantity (norm-per-batch vs norm-of-mean). Kept here only as a reproducibility - warning: **do not cite the -0.18 as evidence**; the at-scale test with the - faithful lens is future work. + warning: **do not cite the -0.18 as evidence**; the at-scale W_U geometry + check (section 5) is the powered version of the geometric claim, and the + at-scale test with the faithful lens is future work. diff --git a/src/wu_row_norm_check.py b/src/wu_row_norm_check.py new file mode 100644 index 0000000..e7e7473 --- /dev/null +++ b/src/wu_row_norm_check.py @@ -0,0 +1,189 @@ +"""GPT-2 W_U row-norm frequency check. + +At-scale test of the jspace-nanogpt geometric finding: at char scale, +the unembedding row norms anti-correlate with token frequency +(r(||W_U[k]||, log10 f) = -0.693, Spearman -0.808) and this geometry is +LEARNED (r ~ 0 at init). Does it generalize to a real LM with +V = 50,257? + +Cheap by design: load lm_head/wte rows, take row norms, correlate against +unigram token counts from a corpus. No Jacobian, no training, no GPU. + +Usage: + .venv/bin/python src/wu_row_norm_check.py [--models gpt2,gpt2-medium] + +Corpus: wikitext-103-raw-v1 (train) from HF. Counts are token-level +unigram frequencies after GPT-2 byte-level BPE encoding. +""" +import argparse +import time + +import numpy as np +import safetensors.torch +import torch +from huggingface_hub import hf_hub_download +from transformers import GPT2TokenizerFast + +WIKITEXT_REPO = "Salesforce/wikitext" +WIKITEXT_PARQUETS = [ + "wikitext-103-raw-v1/train-00000-of-00002.parquet", + "wikitext-103-raw-v1/train-00001-of-00002.parquet", +] + + +def pearson(x, y): + x = x - x.mean() + y = y - y.mean() + denom = np.sqrt((x * x).sum() * (y * y).sum()) + return float((x * y).sum() / denom) if denom > 0 else float("nan") + + +def spearman(x, y): + def rankdata(a): + order = np.argsort(a, kind="mergesort") + ranks = np.empty_like(order, dtype=float) + ranks[order] = np.arange(1, a.size + 1) + return ranks + + return pearson(rankdata(x), rankdata(y)) + + +def load_wte(model_id: str) -> np.ndarray: + path = hf_hub_download(model_id, "model.safetensors") + st = safetensors.torch.load_file(path) + candidates = ["transformer.wte.weight", "wte.weight", "model.embed_tokens.weight"] + key = next((k for k in candidates if k in st), None) + if key is None: + raise KeyError( + f"no wte/embedding weight found in {model_id}; " + f"have: {sorted(st.keys())[:10]} ..." + ) + wte = st[key] + print(f"[{model_id}] {key} {tuple(wte.shape)} dtype={wte.dtype}") + # GPT-2 ties lm_head to wte, so row norms of wte == row norms of W_U. + return torch.norm(wte.float(), dim=1).numpy() + + +def load_corpus() -> list[str]: + import pyarrow.parquet as pq + + texts: list[str] = [] + for name in WIKITEXT_PARQUETS: + path = hf_hub_download(WIKITEXT_REPO, name, repo_type="dataset") + print(f"loading corpus shard from {path}") + table = pq.read_table(path, columns=["text"]) + texts.extend(table.column("text").to_pylist()) + print(f"corpus: {len(texts)} rows") + return texts + + +def count_tokens(tokenizer, texts: list[str], batch_size: int = 1000) -> np.ndarray: + counts = np.zeros(tokenizer.vocab_size, dtype=np.int64) + n_tokens = 0 + t0 = time.time() + for start in range(0, len(texts), batch_size): + chunk = texts[start : start + batch_size] + enc = tokenizer(chunk, add_special_tokens=False) + for ids in enc["input_ids"]: + np.add.at(counts, ids, 1) + n_tokens += len(ids) + if (start // batch_size) % 20 == 0: + print( + f" rows {start}/{len(texts)} tokens {n_tokens:,} " + f"elapsed {time.time() - t0:.0f}s" + ) + print(f"done: {n_tokens:,} tokens in {time.time() - t0:.0f}s") + return counts + + +def report(model_id: str, norms: np.ndarray, counts: np.ndarray, tokenizer): + freqs = counts.astype(np.float64) + seen = freqs > 0 + logf = np.log10(freqs + 1.0) # smoothed; +1 keeps unseen tokens finite + + def r_raw(mask): + return pearson(freqs[mask], norms[mask]) + + def r_log(mask): + return pearson(logf[mask], norms[mask]) + + def r_sp(mask): + return spearman(freqs[mask], norms[mask]) + + n = norms.size + n_seen = int(seen.sum()) + print(f"\n===== {model_id} ===== V={n} tokens seen={n_seen} " + f"({100 * n_seen / n:.1f}%)") + print(f" r(raw freq, norm) all V: {r_raw(np.ones(n, bool)):+.3f}") + print(f" r(raw freq, norm) seen only: {r_raw(seen):+.3f}") + print(f" r(log10 freq, norm) all V: {r_log(np.ones(n, bool)):+.3f}") + print(f" r(log10 freq, norm) seen only: {r_log(seen):+.3f}") + for thr in (5, 100, 1000): + m = freqs >= thr + print(f" r(log10 freq, norm) freq>={thr:<6} " + f"(n={int(m.sum()):>5}): {r_log(m):+.3f}") + print(f" Spearman(raw, norm) seen only: {r_sp(seen):+.3f}") + + # Decile table: quantiles of log-frequency vs mean norm. + if n_seen > 10: + qs = np.quantile(logf[seen], np.linspace(0, 1, 11)) + idx = np.digitize(logf, qs[1:-1]) + print(" log10-freq decile -> mean norm:") + for d in range(10): + m = (idx == d) & seen + if m.sum() == 0: + continue + lo, hi = qs[d], qs[d + 1] + print(f" [{lo:5.2f},{hi:5.2f}] n={int(m.sum()):>5} " + f"mean norm={norms[m].mean():.3f}") + + # Top tokens by norm (narrative color). + order = np.argsort(norms)[::-1][:15] + print(" top-15 tokens by row norm (token | norm | count | rank-by-freq):") + for i in order: + tok = tokenizer.decode([int(i)]) + rank = int((freqs > freqs[i]).sum()) + 1 + print(f" {tok!r:>14} {norms[i]:.3f} {int(freqs[i]):>8,} #{rank:,}") + return { + "model": model_id, + "r_raw_all": r_raw(np.ones(n, bool)), + "r_log_all": r_log(np.ones(n, bool)), + "r_log_seen": r_log(seen), + "spearman_seen": r_sp(seen), + "n_seen": n_seen, + } + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--models", default="gpt2", help="comma-separated HF ids") + ap.add_argument("--no-corpus", action="store_true", + help="skip tokenization (use saved counts if present)") + args = ap.parse_args() + + tokenizer = GPT2TokenizerFast.from_pretrained("openai-community/gpt2") + if args.no_corpus: + counts = np.load("/tmp/wu_counts.npy") + else: + texts = load_corpus() + counts = count_tokens(tokenizer, texts) + np.save("/tmp/wu_counts.npy", counts) + + results = [] + for mid in args.models.split(","): + mid = mid.strip() + norms = load_wte(mid) + results.append(report(mid, norms, counts, tokenizer)) + + print("\n===== SUMMARY =====") + for r_ in results: + print( + f"{r_['model']}: r(raw)={r_['r_raw_all']:+.3f} " + f"r(log10, all V)={r_['r_log_all']:+.3f} " + f"r(log10, seen)={r_['r_log_seen']:+.3f} " + f"Spearman(seen)={r_['spearman_seen']:+.3f}" + ) + + +if __name__ == "__main__": + main() |