summaryrefslogtreecommitdiff
path: root/src/jlens.py
diff options
context:
space:
mode:
Diffstat (limited to 'src/jlens.py')
-rw-r--r--src/jlens.py441
1 files changed, 441 insertions, 0 deletions
diff --git a/src/jlens.py b/src/jlens.py
new file mode 100644
index 0000000..515dc08
--- /dev/null
+++ b/src/jlens.py
@@ -0,0 +1,441 @@
+"""
+J-lens: Jacobian Lens for Transformer Models
+
+Replicates Anthropic's technique from:
+"Verbalizable Representations Form a Global Workspace in Language Models"
+https://transformer-circuits.pub/2026/workspace/index.html
+
+Core idea: For each token in the vocabulary, compute the average gradient
+of log p(token) with respect to the residual stream at each layer,
+averaged over many contexts. This reveals which concepts are "verbalizable"
+— readily available for the model to report on.
+
+Usage:
+ python jlens.py --model checkpoints/ckpt.pt --data data/shakespeare_char
+"""
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+import numpy as np
+import argparse
+import json
+import os
+import pickle
+from pathlib import Path
+from collections import defaultdict
+
+
+def load_model(checkpoint_path, model_class, device='cuda'):
+ """Load a trained nanoGPT model from checkpoint."""
+ checkpoint = torch.load(checkpoint_path, map_location=device)
+ # nanoGPT stores model args, state_dict + optimizer in checkpoint
+ model_args = checkpoint['model_args']
+
+ # Create model with saved config
+ model = model_class(model_args)
+
+ # Fix state dict keys (nanoGPT wraps in DataParallel)
+ state_dict = checkpoint['model']
+ unwanted_prefix = '_orig_mod.'
+ for k in list(state_dict.keys()):
+ if k.startswith(unwanted_prefix):
+ state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
+
+ model.load_state_dict(state_dict)
+ model.to(device)
+ model.eval()
+ return model, model_args
+
+
+def compute_jlens_single_token(model, token_id, dataloader, layer_idx, device='cuda'):
+ """
+ Compute J-lens vector for a single token at a specific layer.
+
+ J_l(token, layer) = E_x [ ∇_{resid[layer]} log p(token | x) ]
+
+ Where the expectation is taken over all positions in the corpus.
+ """
+ vectors = []
+
+ with torch.no_grad():
+ for batch_idx, (x, y) in enumerate(dataloader):
+ x, y = x.to(device), y.to(device)
+ B, T = x.shape
+
+ # We need gradients, so we'll do forward passes with hooks
+ # Strategy: use torch.autograd.grad on a forward pass
+ # where we capture residual stream activations
+
+ # Register hook to capture residual stream at target layer
+ activations = {}
+
+ def make_hook():
+ def hook(module, input, output):
+ # output is (B, T, n_embd)
+ # Detach then require grad so we can compute gradient through it
+ activations['resid'] = output.detach().requires_grad_(True)
+ return activations['resid']
+ return hook
+
+ # Find the target layer
+ target_block = model.transformer.h[layer_idx]
+ # nanoGPT architecture: h = x + attn(ln1(x)), then x = h + mlp(ln2(h))
+ # We want the residual stream AFTER the attention + MLP of this layer
+ # which is the output of the block
+
+ handle = target_block.register_forward_hook(make_hook())
+
+ # Forward pass
+ logits, loss = model(x, y)
+
+ handle.remove()
+
+ # Now compute gradient of log p(token_id) w.r.t. residual stream
+ # log p(token_id) at each position = log_softmax(logits)[:, :, token_id]
+ log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
+ token_log_probs = log_probs[:, :, token_id].sum() # sum over B, T
+
+ # Gradient of this sum w.r.t. the captured activations
+ grad = torch.autograd.grad(
+ token_log_probs,
+ activations['resid'],
+ retain_graph=False
+ )[0] # Shape: (B, T, n_embd)
+
+ vectors.append(grad.detach().cpu())
+
+ # Cleanup
+ del logits, loss, log_probs, activations, grad
+ torch.cuda.empty_cache()
+
+ # Average over all positions in the corpus
+ all_vectors = torch.cat([v.reshape(-1, v.shape[-1]) for v in vectors], dim=0)
+ jlens_vector = all_vectors.mean(dim=0) # Shape: (n_embd,)
+
+ return jlens_vector
+
+
+def compute_jlens_all_tokens(model, dataloader, layer_idx, vocab_size, device='cuda'):
+ """
+ Compute J-lens vectors for all tokens at a specific layer.
+
+ Returns: dict mapping token_id -> jlens_vector (n_embd,)
+ """
+ jlens_vectors = {}
+
+ for token_id in range(vocab_size):
+ vec = compute_jlens_single_token(model, token_id, dataloader, layer_idx, device)
+ jlens_vectors[token_id] = vec
+
+ if (token_id + 1) % 10 == 0:
+ print(f" Token {token_id + 1}/{vocab_size} done")
+
+ return jlens_vectors
+
+
+def compute_jlens_all_layers(model, dataloader, n_layers, vocab_size, device='cuda',
+ use_batched=True):
+ """
+ Compute J-lens vectors for all layers and all tokens.
+
+ Uses batched approach: for each context, compute gradients for ALL tokens
+ at once using vector-Jacobian products. Much faster than per-token.
+
+ Returns: dict mapping layer_idx -> {token_id: jlens_vector}
+ """
+ all_layer_vectors = defaultdict(dict)
+
+ if use_batched:
+ # Optimized: compute all token J-lens vectors simultaneously
+ # For each context position, the gradient of log p(token) w.r.t. resid
+ # for all tokens is just the Jacobian of the unembedding layer
+ # which equals W_U^T * (one_hot(token) - softmax(logits))
+ # Wait, let me think about this more carefully...
+
+ print("Using batched J-lens computation...")
+
+ for layer_idx in range(n_layers):
+ print(f"\nLayer {layer_idx}/{n_layers}...")
+
+ layer_accum = torch.zeros(vocab_size, model.config.n_embd, device='cpu')
+ token_count = torch.zeros(vocab_size, device='cpu')
+
+ with torch.no_grad():
+ for batch_idx, (x, y) in enumerate(dataloader):
+ x, y = x.to(device), y.to(device)
+ B, T = x.shape
+
+ # Capture residual stream at target layer
+ resid_captured = {}
+
+ def make_hook(resid_dict):
+ def hook(module, input, output):
+ resid_dict['val'] = output.detach().requires_grad_(True)
+ return resid_dict['val']
+ return hook
+
+ target_block = model.transformer.h[layer_idx]
+ handle = target_block.register_forward_hook(make_hook(resid_captured))
+
+ logits, loss = model(x, y)
+ handle.remove()
+
+ # Now: for each token in vocab, we want d(logit_t)/d(resid)
+ # This is the Jacobian of unembedding w.r.t. residual stream
+ # Chain rule: d(logit_t)/d(resid) = W_U[t, :] * d(layer_out)/d(resid)
+ # where layer_out is the final layer output after all remaining layers
+ # plus the direct path through the residual stream.
+
+ # Actually, since we captured resid at layer L, and the model
+ # applies remaining layers resid_L -> ... -> resid_final -> logits,
+ # the gradient d(logits)/d(resid_L) = d(logits)/d(resid_final) * d(resid_final)/d(resid_L)
+ #
+ # We can compute this by:
+ # 1. Get logits
+ # 2. For EACH position, compute gradient of logit for EACH token
+ # w.r.t. the captured residual stream
+ # 3. Average across positions
+
+ # Vectorized approach: compute gradients for ALL tokens simultaneously
+ # using torch.autograd.grad with list of outputs
+
+ # For efficiency, compute per position, then aggregate
+ log_probs = torch.nn.functional.log_softmax(logits, dim=-1) # (B, T, vocab)
+
+ # For each position (b, t), we need jacobian of log_probs[b,t,:] w.r.t. resid[b,t,:]
+ # This is (vocab, n_embd) per position
+ # We can batch by computing gradient of sum_{tokens} a_i * log_p(token_i)
+ # where a_i cycles through standard basis vectors
+
+ # Practical approach for small vocab (nanoGPT: 65 tokens):
+ # Just loop over tokens, compute gradient, and accumulate
+
+ resid = resid_captured['val'] # (B, T, n_embd)
+
+ for token_id in range(vocab_size):
+ # Gradient of log_p(token_id) summed over all positions
+ token_log_prob = log_probs[:, :, token_id].sum()
+
+ grad = torch.autograd.grad(
+ token_log_prob, resid, retain_graph=(token_id < vocab_size - 1)
+ )[0] # (B, T, n_embd)
+
+ # Accumulate: sum of gradients across all positions
+ layer_accum[token_id] += grad.detach().cpu().reshape(-1, model.config.n_embd).sum(dim=0)
+ token_count[token_id] += B * T
+
+ del logits, loss, log_probs, resid
+ del grad # pyright: ignore[reportPossiblyUnboundVariable]
+ torch.cuda.empty_cache()
+
+ if (batch_idx + 1) % 10 == 0:
+ print(f" Batch {batch_idx + 1}/{len(dataloader)}")
+
+ # Average: divide sum by count
+ for token_id in range(vocab_size):
+ if token_count[token_id] > 0:
+ all_layer_vectors[layer_idx][token_id] = layer_accum[token_id] / token_count[token_id]
+ else:
+ all_layer_vectors[layer_idx][token_id] = torch.zeros(model.config.n_embd)
+
+ print(f" Layer {layer_idx} complete. Saved {vocab_size} token vectors.")
+
+ return dict(all_layer_vectors)
+
+
+def save_jlens(jlens_data, output_path, metadata=None):
+ """Save J-lens vectors to disk."""
+ output = {
+ 'metadata': metadata or {},
+ 'vectors': {
+ str(layer): {
+ str(token_id): vec.numpy() for token_id, vec in tokens.items()
+ }
+ for layer, tokens in jlens_data.items()
+ }
+ }
+
+ os.makedirs(os.path.dirname(output_path), exist_ok=True)
+ with open(output_path, 'wb') as f:
+ pickle.dump(output, f)
+
+ print(f"Saved J-lens data to {output_path}")
+
+
+def load_jlens(path):
+ """Load saved J-lens vectors."""
+ with open(path, 'rb') as f:
+ data = pickle.load(f)
+
+ # Convert back to tensors
+ jlens = {}
+ for layer_str, tokens in data['vectors'].items():
+ layer = int(layer_str)
+ jlens[layer] = {}
+ for token_id_str, vec in tokens.items():
+ jlens[layer][int(token_id_str)] = torch.from_numpy(vec)
+
+ return jlens, data['metadata']
+
+
+def analyze_jlens(jlens_data, itos, n_layers, output_dir='outputs'):
+ """Analyze and visualize J-lens vectors."""
+ os.makedirs(output_dir, exist_ok=True)
+ vocab_size = len(itos)
+
+ print(f"\n{'='*60}")
+ print("J-LENS ANALYSIS")
+ print(f"{'='*60}")
+
+ for layer_idx in range(n_layers):
+ if layer_idx not in jlens_data:
+ continue
+
+ layer_vectors = jlens_data[layer_idx]
+
+ # Compute norm of each token's J-lens vector
+ norms = {}
+ for token_id, vec in layer_vectors.items():
+ norms[token_id] = vec.norm().item()
+
+ # Sort by norm (most "verbalizable" tokens first)
+ sorted_tokens = sorted(norms.items(), key=lambda x: x[1], reverse=True)
+
+ print(f"\n--- Layer {layer_idx} ---")
+ print(f"Top 10 most verbalizable tokens:")
+ for token_id, norm in sorted_tokens[:10]:
+ token_str = itos[token_id].replace('\n', '\\n')
+ print(f" '{token_str}': norm={norm:.4f}")
+
+ print(f"Bottom 5 least verbalizable tokens:")
+ for token_id, norm in sorted_tokens[-5:]:
+ token_str = itos[token_id].replace('\n', '\\n')
+ print(f" '{token_str}': norm={norm:.4f}")
+
+ # Compute J-space "capacity" — how many tokens have significant norm?
+ print(f"\n--- J-space Capacity ---")
+ for layer_idx in range(n_layers):
+ if layer_idx not in jlens_data:
+ continue
+ layer_vectors = jlens_data[layer_idx]
+ norms = torch.tensor([v.norm().item() for v in layer_vectors.values()])
+
+ # Count "active" tokens (norm > median * 2)
+ threshold = norms.median() * 2
+ active = (norms > threshold).sum().item()
+ print(f" Layer {layer_idx}: {active}/{vocab_size} tokens active (threshold={threshold:.4f})")
+
+
+if __name__ == '__main__':
+ parser = argparse.ArgumentParser(description='J-lens: Jacobian Lens for nanoGPT')
+ parser.add_argument('--checkpoint', type=str, required=True,
+ help='Path to model checkpoint')
+ parser.add_argument('--data_dir', type=str, default='data/shakespeare_char',
+ help='Path to data directory')
+ parser.add_argument('--output_dir', type=str, default='outputs/jlens',
+ help='Directory for saving outputs')
+ parser.add_argument('--batch_size', type=int, default=32,
+ help='Batch size for processing')
+ parser.add_argument('--max_batches', type=int, default=100,
+ help='Max batches to process (limit for speed)')
+ parser.add_argument('--layers', type=str, default=None,
+ help='Comma-separated layer indices (default: all)')
+ parser.add_argument('--device', type=str, default='cuda',
+ help='Device to use')
+
+ args = parser.parse_args()
+
+ # Import model from local src
+ import sys
+ sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))
+ from model import GPT, GPTConfig
+
+ # Load model
+ print(f"Loading model from {args.checkpoint}")
+ model, model_args = load_model(args.checkpoint, GPT, args.device)
+ print(f"Model: {model_args.n_layer} layers, {model_args.n_embd} dim, "
+ f"{model_args.n_head} heads, {model_args.vocab_size} vocab")
+
+ # Load data
+ data_dir = Path(args.data_dir)
+ train_data = np.memmap(data_dir / 'train.bin', dtype=np.uint16, mode='r')
+ val_data = np.memmap(data_dir / 'val.bin', dtype=np.uint16, mode='r')
+
+ # Load vocab mappings
+ meta_path = data_dir / 'meta.pkl'
+ if meta_path.exists():
+ with open(meta_path, 'rb') as f:
+ meta = pickle.load(f)
+ itos = meta['itos']
+ stoi = meta['stoi']
+ else:
+ # Default char-level vocab
+ chars = sorted(list(set(open(data_dir / 'input.txt').read())))
+ stoi = {ch: i for i, ch in enumerate(chars)}
+ itos = {i: ch for i, ch in enumerate(chars)}
+
+ print(f"Vocabulary size: {len(itos)}")
+ print(f"Train data: {len(train_data):,} tokens")
+
+ # Create dataloader
+ def get_batch(split):
+ data = train_data if split == 'train' else val_data
+ block_size = model_args.block_size
+ ix = torch.randint(len(data) - block_size, (args.batch_size,))
+ x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64))
+ for i in ix])
+ y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64))
+ for i in ix])
+ return x, y
+
+ class SimpleDataset(torch.utils.data.IterableDataset):
+ def __iter__(self):
+ while True:
+ yield get_batch('train')
+
+ dataset = SimpleDataset()
+ dataloader = DataLoader(dataset, batch_size=None, num_workers=0)
+
+ # Limit to max_batches
+ limited_dataloader = []
+ for i, batch in enumerate(dataloader):
+ if i >= args.max_batches:
+ break
+ limited_dataloader.append(batch)
+
+ print(f"Processing {len(limited_dataloader)} batches of size {args.batch_size}")
+
+ # Determine layers to process
+ if args.layers:
+ layers_to_process = [int(l) for l in args.layers.split(',')]
+ else:
+ layers_to_process = list(range(model_args.n_layer))
+
+ print(f"Computing J-lens for layers: {layers_to_process}")
+
+ # Compute J-lens for selected layers
+ jlens_data = {}
+ for layer_idx in layers_to_process:
+ print(f"\nComputing J-lens for layer {layer_idx}...")
+ layer_vectors = compute_jlens_all_tokens(
+ model, limited_dataloader, layer_idx,
+ model_args.vocab_size, args.device
+ )
+ jlens_data[layer_idx] = layer_vectors
+
+ # Save results
+ save_path = os.path.join(args.output_dir, 'jlens_vectors.pkl')
+ metadata = {
+ 'model_args': vars(model_args),
+ 'num_batches': len(limited_dataloader),
+ 'batch_size': args.batch_size,
+ 'layers_processed': layers_to_process,
+ 'vocab_size': model_args.vocab_size,
+ }
+ save_jlens(jlens_data, save_path, metadata)
+
+ # Analyze
+ analyze_jlens(jlens_data, itos, model_args.n_layer, args.output_dir)
+
+ print(f"\nDone! Results saved to {args.output_dir}")
generated by cgit v1.2.3 (git 2.46.0) at 2026-08-07 22:47:42 +0000