diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/jlens_v3.py | 238 |
1 files changed, 238 insertions, 0 deletions
diff --git a/src/jlens_v3.py b/src/jlens_v3.py new file mode 100644 index 0000000..7112cb1 --- /dev/null +++ b/src/jlens_v3.py @@ -0,0 +1,238 @@ +""" +J-lens v3: FAITHFUL implementation of Anthropic's Jacobian lens. + +Paper: "Verbalizable Representations Form a Global Workspace in Language Models" +https://transformer-circuits.pub/2026/workspace/index.html + +Definition (from the paper's Methods): + J_l = E_{t, t' >= t, prompt}[ d h_final,t' / d h_l,t ] (d_model x d_model per layer) + J-lens vectors at layer l = rows of W_U * J_l + lens(h_l) = softmax(W_U * norm(J_l * h_l)) + +The average is over source positions t, all future positions t' >= t, and a corpus +of prompts. The result is a single d x d matrix per layer that maps the residual +stream at layer l to the FINAL residual stream, read out through the model's own +unembedding (W_U = lm_head.weight in nanoGPT). + +This is different from the earlier proxy (jlens_v2) which differentiated +log_softmax(logits) directly, folding the token-dependent softmax calibration +into the averaged object. + +Implementation: basis-chunked vector-Jacobian products. For a chunk C of basis +vectors {e_c}, compute + S_c = sum_{t'} (h_final,t' . e_c) +and backpropagate S_c to h_l (the captured residual stream at layer l). +Because h_final,t' does not depend on h_l,t for t' < t (causality), summing over +all t' yields exactly sum over valid pairs. One batched autograd call per chunk +(is_grads_batched=True) gives all rows in the chunk. + +Also implements the BOTH-WAYS comparison: old proxy ||E[grad log p(k|x)]|| vs +faithful ||(W_U J_l)[k]||, each correlated against token frequency, per layer. + +Usage (inside Docker on meru): + python3 src/jlens_v3.py --checkpoint out-shakespeare-char/ckpt.pt \ + --data_dir data/shakespeare_char --n_prompts 20 --layers 2,3,4 +""" + +import sys, os +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +import argparse +import torch +import numpy as np +import pickle + +import jlens_v2 # for load_model + + +def make_batches(train_data, block_size, batch_size, n_prompts, device, seed=1337): + """Random contiguous prompts from train.bin (n_prompts prompts of batch_size).""" + g = torch.Generator().manual_seed(seed) + batches = [] + for _ in range(n_prompts): + ix = torch.randint(len(train_data) - block_size, (batch_size,), generator=g) + x = torch.stack([torch.from_numpy( + train_data[i:i + block_size].astype(np.int64)) for i in ix]) + y = torch.stack([torch.from_numpy( + train_data[i + 1:i + 1 + block_size].astype(np.int64)) for i in ix]) + batches.append((x.to(device), y.to(device))) + return batches + + +def compute_faithful_jlens(model, layer_idx, batches, device, chunk=32): + """ + Compute the faithful J-lens VECTORS (rows of W_U * J_l) directly. + + (W_U J_l)[k] = E_{t, t'>=t, prompt}[ d(W_U[k] . h_final,t') / d h_l,t ] + = E[ grad of the k-th RAW LOGIT (pre-softmax) w.r.t. h_l ] + + Instead of materializing the d x d residual Jacobian (384 basis VJPs), we + probe with the unembedding rows W_U[k] (V=65 probes) — mathematically + identical, ~6x cheaper. One batched autograd call per prompt per layer + (is_grads_batched=True) computes all V token vectors at once. + """ + d = model.config.n_embd + V = model.config.vocab_size + n_layer = model.config.n_layer + last_block = model.transformer.h[n_layer - 1] + W_U = model.lm_head.weight.detach().float() # (V, d), nanoGPT weight tying + + accum = torch.zeros(V, d, device='cpu') + n_pairs = 0 + + for (x, y) in batches: + B, T = x.shape + resid_captured = {} + + def hook(module, inp, out): + resid_captured['h_l'] = out + + def hook_final(module, inp, out): + resid_captured['h_final'] = out + + h1 = model.transformer.h[layer_idx].register_forward_hook(hook) + h2 = last_block.register_forward_hook(hook_final) + logits, loss = model(x, y) # graph connected + h1.remove() + h2.remove() + + h_l = resid_captured['h_l'] # (B, T, d) + h_final = resid_captured['h_final'] # (B, T, d) + # sum over all future positions t' (causal: t' < t contributes zero grad) + h_final_sum = h_final.sum(dim=1) # (B, d) + + # grad_outputs[v, b, :] = W_U[v] -> batched VJPs for all V tokens + grad_outputs = W_U.to(device).unsqueeze(1).expand(V, B, d).contiguous() + grads = torch.autograd.grad( + h_final_sum, h_l, grad_outputs=grad_outputs, + is_grads_batched=True)[0] # (V, B, T, d) + accum += grads.detach().cpu().reshape(V, -1, d).sum(dim=1) + + n_pairs += B * T * (T + 1) // 2 # valid (t, t' >= t) pairs + del logits, loss, h_l, h_final, h_final_sum, grads, grad_outputs + torch.cuda.empty_cache() + + accum /= n_pairs + return accum # (V, d): faithful J-lens vectors, rows of W_U * J_l + + +def compute_proxy_norms(model, layer_idx, batches, device, vocab_size): + """ + Old proxy (jlens_v2 quantity): || E_x[ grad log_softmax(logits)[:, :, k] ] || + averaged over positions — norm of the mean gradient, per token. + """ + d = model.config.n_embd + accum = torch.zeros(vocab_size, d, device='cpu') + counts = torch.zeros(vocab_size, device='cpu') + for (x, y) in batches: + B, T = x.shape + resid_captured = {} + + def hook(module, inp, out): + resid_captured['val'] = out + + handle = model.transformer.h[layer_idx].register_forward_hook(hook) + logits, loss = model(x, y) + handle.remove() + log_probs = torch.nn.functional.log_softmax(logits, dim=-1) + # sum over positions -> (V,); batched VJP over all vocab at once + lp_sum = log_probs.sum(dim=(0, 1)) + grad_outputs = torch.eye(vocab_size, device=device) + grad = torch.autograd.grad( + lp_sum, resid_captured['val'], grad_outputs=grad_outputs, + is_grads_batched=True)[0] # (V, B, T, d) + accum += grad.detach().cpu().reshape(vocab_size, -1, d).sum(dim=1) + counts += B * T + del logits, loss, log_probs, lp_sum, grad + torch.cuda.empty_cache() + norms = {} + for token_id in range(vocab_size): + norms[token_id] = (accum[token_id] / counts[token_id]).norm().item() + return norms + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument('--checkpoint', required=True) + ap.add_argument('--data_dir', default='data/shakespeare_char') + ap.add_argument('--n_prompts', type=int, default=20) + ap.add_argument('--batch_size', type=int, default=16) + ap.add_argument('--layers', default='2,3,4', help='comma-separated layer indices') + ap.add_argument('--device', default='cuda') + ap.add_argument('--chunk', type=int, default=32) + args = ap.parse_args() + + device = args.device + model, config = jlens_v2.load_model(args.checkpoint, device) + model.eval() + d = config.n_embd + V = config.vocab_size + n_layer = config.n_layer + print(f"Model: {n_layer} layers, d={d}, V={V}, V/d={V/d:.2f}x") + + train_data = np.memmap(f'{args.data_dir}/train.bin', dtype=np.uint16, mode='r') + with open(f'{args.data_dir}/meta.pkl', 'rb') as f: + meta = pickle.load(f) + itos = meta['itos'] + + # token frequencies + counts = np.bincount(train_data, minlength=V).astype(float) + freq = counts / counts.sum() * 100 # percent + + batches = make_batches(train_data, config.block_size, args.batch_size, + args.n_prompts, device) + + layers = [int(l) for l in args.layers.split(',')] + W_U = model.lm_head.weight.detach().float() # (V, d) — nanoGPT weight tying + + print(f"\n{'='*78}") + print("BOTH-WAYS COMPARISON: old proxy vs faithful J-lens (rows of W_U * J_l)") + print(f"{'='*78}") + + for layer_idx in layers: + print(f"\n--- Layer {layer_idx} ---") + print(" Computing faithful J-lens vectors (W_U-probed VJPs)...") + faithful_vecs = compute_faithful_jlens(model, layer_idx, batches, device, + args.chunk) # (V, d) + faithful_norms = {k: faithful_vecs[k].norm().item() for k in range(V)} + + if layer_idx == n_layer - 1: + # Validation: J_{L-1} should be identity, so faithful vectors == W_U rows + sims = torch.nn.functional.cosine_similarity( + faithful_vecs.float(), W_U.float(), dim=1) + print(f" [validation] last layer: mean cos-sim(faithful, W_U rows) = " + f"{sims.mean().item():.4f} (expect ~1.0 if J=identity)") + + print(" Computing old proxy norms...") + proxy_norms = compute_proxy_norms(model, layer_idx, batches, device, V) + + # correlations with frequency + f_arr = np.array([freq[k] for k in range(V)]) + p_arr = np.array([proxy_norms[k] for k in range(V)]) + fa_arr = np.array([faithful_norms[k] for k in range(V)]) + r_proxy = np.corrcoef(p_arr, f_arr)[0, 1] + r_faith = np.corrcoef(fa_arr, f_arr)[0, 1] + + print(f" freq corr | old proxy: r = {r_proxy:+.3f}") + print(f" freq corr | faithful: r = {r_faith:+.3f}") + + # top/bottom by faithful norm + srt = sorted(faithful_norms.items(), key=lambda kv: kv[1], reverse=True) + def esc(s): + return s.replace(chr(10), '\\n') + print(" Top 5 by faithful norm:") + for tid, n in srt[:5]: + print(f" '{esc(itos[tid])}' freq={freq[tid]:.3f}% norm={n:.4f}") + print(" Bottom 5 by faithful norm:") + for tid, n in srt[-5:]: + print(f" '{esc(itos[tid])}' freq={freq[tid]:.3f}% norm={n:.4f}") + + torch.save({'faithful_vecs': faithful_vecs, 'faithful_norms': faithful_norms, + 'proxy_norms': proxy_norms}, + f'outputs/jlens_v3_layer{layer_idx}.pt') + print(f" Saved outputs/jlens_v3_layer{layer_idx}.pt") + + print("\nDONE.") + + +if __name__ == '__main__': + main() |