"""Claim 4: Spectral bias analysis. Compute graph Laplacian eigenvectors (Fiedler vectors) and compare alignment with learned embeddings. Verify that geometry stems from spectral bias. """ import sys sys.path.insert(0, '/tmp/geometric_memory') import os import numpy as np import torch from pathlib import Path from geometric_memory.in_weights.config import get_args from geometric_memory.tokenizing import get_tokenizer from geometric_memory.models import get_model from scipy.sparse import csr_matrix from scipy.sparse.linalg import eigsh OUTPUT_DIR = '/tmp/claim4_analysis' os.makedirs(OUTPUT_DIR, exist_ok=True) def build_args(): cli_args = [ '--training_recipe', 'staged_full_path', '--model_family', 'gpt', '--graph_type', 'star', '--star_degree', '10', '--star_subtree_degree', '10', '--path_length', '4', '--total_nodes', '-1', '--add_forward_edges', '--add_backward_edges', '--edge_memorization_epochs', '20', '--path_finetuning_epochs', '20', '--edge_memorization_batch_size', '64', '--path_finetuning_batch_size', '64', '--edge_memorization_learning_rate', '0.01', '--path_finetuning_learning_rate', '0.0005', '--optimizer_weight_decay', '0.0', '--edge_memorization_warmup_steps', '0', '--path_finetuning_warmup_steps', '0', '--disable_edge_memorization_lr_decay', '--disable_path_finetuning_lr_decay', '--edge_memorization_eval_interval_epochs', '5', '--path_finetuning_eval_interval_epochs', '1', '--embedding_dimension', '384', '--attention_head_count', '8', '--transformer_layer_count', '12', '--dropout_rate', '0.0', '--experiment_log_root', '/tmp/claim3_run1', '--dataset_root', '/tmp/geometric_memory/data/datasets/in_weights_graphs', '--dataset_name', 'in_weights', '--train_split_ratio', '0.75', '--path_prefix_pause_token_count', '1', '--exclude_task_token_in_prefix', '--no-enable_wandb', ] return get_args(cli_args) def extract_embeddings_with_edges(checkpoint_path, tokenizer, device): """Extract embeddings AND track which edge (u,v) each embedding belongs to.""" args = build_args() tokenizer = get_tokenizer(args) args.vocab_size = tokenizer.vocab_size args.block_size = max(64, args.path_length * 3) args.teacherless_token = tokenizer.encode("$")[0] if args.use_teacherless_inputs else None args.use_flash = False model = get_model(args) state_dict = torch.load(checkpoint_path, map_location=device, weights_only=False) model.load_state_dict(state_dict) model = model.to(device) model.eval() from geometric_memory.in_weights.data_loader import EdgeMemorizationDataset pretrain_path = str(Path(args.dataset_directory) / 'star_deg_10_deg_tree_10_path_4_nodes_1111_sd_00_fb_11_selfedge_0_pretrain.txt') dataset = EdgeMemorizationDataset( tokenizer=tokenizer, data_path=pretrain_path, device=device, teacherless_token_id=args.teacherless_token, drop_pause_token=True, include_task_token_in_prefix=False, eval_mode=True, ) all_embeddings = [] all_edges = [] # (u, v) for each edge with torch.no_grad(): for i in range(len(dataset)): seq = dataset[i].to(device) # (seq_len,) if hasattr(model, 'transformer') and hasattr(model.transformer, 'wte'): embeddings = model.transformer.wte(seq.unsqueeze(0)) elif hasattr(model, 'wte'): embeddings = model.wte(seq.unsqueeze(0)) else: for name, module in model.named_modules(): if 'embed' in name.lower() and hasattr(module, 'weight'): embeddings = module(seq.unsqueeze(0)) break else: continue last_emb = embeddings[0, -1, :] # (dim,) all_embeddings.append(last_emb.cpu().numpy()) # Extract edge from sequence: for edge memorization, seq = [task_token, u, v] # or [task_token, u] depending on format seq_tokens = seq.cpu().numpy() if len(seq_tokens) >= 2: u = int(seq_tokens[-2]) # source node v = int(seq_tokens[-1]) # target node all_edges.append((u, v)) else: all_edges.append((0, 0)) # fallback return np.array(all_embeddings), np.array(all_edges) def aggregate_to_node_embeddings(embeddings, edges, n_nodes): """Aggregate edge embeddings to node-level embeddings by averaging. For each node, average all embeddings where that node appears as source or target. """ node_embeddings = {i: [] for i in range(n_nodes)} for emb, (u, v) in zip(embeddings, edges): node_embeddings[u].append(emb) node_embeddings[v].append(emb) # Average node_emb_matrix = np.zeros((n_nodes, embeddings.shape[1])) for i in range(n_nodes): if len(node_embeddings[i]) > 0: node_emb_matrix[i] = np.mean(node_embeddings[i], axis=0) else: node_emb_matrix[i] = np.zeros(embeddings.shape[1]) return node_emb_matrix def compute_laplacian(adj): """Compute normalized graph Laplacian.""" n = adj.shape[0] deg = adj.sum(axis=1) D = np.diag(deg) L = D - adj # Normalized Laplacian: I - D^(-1/2) A D^(-1/2) D_inv_sqrt = np.diag(1.0 / np.sqrt(np.where(deg > 0, deg, 1))) L_norm = np.eye(n) - D_inv_sqrt @ adj @ D_inv_sqrt return L_norm, L def compute_fiedler_vectors(L_norm, k=10): """Compute top k eigenvectors of the normalized Laplacian.""" n = L_norm.shape[0] try: eigenvalues, eigenvectors = eigsh(L_norm, k=k, which='SM', tol=1e-6) # Sort by eigenvalue idx = np.argsort(eigenvalues) eigenvalues = eigenvalues[idx] eigenvectors = eigenvectors[:, idx] return eigenvalues, eigenvectors except Exception as e: print(f"Warning: eigsh failed: {e}") return None, None def compute_alignment(node_emb, spectral_vec): """Compute alignment between node embeddings and spectral vectors. For each spectral dimension, compute the correlation between: - The spectral coordinate of each node - The projection of that node's embedding onto the first PCA component """ from sklearn.decomposition import PCA n_nodes = node_emb.shape[0] # PCA to reduce embedding dimensionality pca = PCA(n_components=min(10, node_emb.shape[1])) node_emb_pca = pca.fit_transform(node_emb) correlations = [] for i in range(min(spectral_vec.shape[1], node_emb_pca.shape[1])): vec = spectral_vec[:, i] emb_dim = node_emb_pca[:, i] if np.std(vec) > 1e-10 and np.std(emb_dim) > 1e-10: corr = np.corrcoef(vec, emb_dim)[0, 1] correlations.append(corr) else: correlations.append(0.0) return np.array(correlations) def procrustes_alignment(A, B): """Compute Procrustes similarity between two sets of vectors.""" # Normalize A = A / (np.linalg.norm(A, axis=1, keepdims=True) + 1e-10) B = B / (np.linalg.norm(B, axis=1, keepdims=True) + 1e-10) # Find optimal rotation from scipy.linalg import svd U, S, Vt = svd(B.T @ A) R = U @ Vt # Compute similarity projected = A @ R.T similarity = np.sum(projected * B) / (np.linalg.norm(A) * np.linalg.norm(B) + 1e-10) return similarity def main(): print("=" * 60) print("Claim 4: Spectral Bias Analysis") print("=" * 60) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"Device: {device}") # Load embeddings from Claim 1 (unfrozen) c1_path = '/tmp/claim3_run1/in_weights/in_weights_star-d10-dt10-p4-n1111_gpt-L12-D384-H8_staged_full_path_bs64x64-lr0p01x0p0005_tl0-rev0-sd00-fb11-selfedge0-task0-split0_20260716-081327/checkpoints/in_weights_star-d10-dt10-p4-n1111_gpt-L12-D384-H8_staged_full_path_bs64x64-lr0p01x0p0005_tl0-rev0-sd00-fb11-selfedge0-task0-split0_20260716-081327_final_model.pt' print("\n--- Extracting Claim 1 embeddings ---") emb_c1, edges_c1 = extract_embeddings_with_edges(c1_path, None, device) print(f"Claim 1: {len(emb_c1)} embeddings, shape {emb_c1.shape}") unique_edges = set() for e in edges_c1: unique_edges.add(tuple(e)) print(f"Unique edges: {len(unique_edges)}") # Determine number of nodes from edges unique_nodes = set() for u, v in edges_c1: unique_nodes.add(u) unique_nodes.add(v) n_nodes = max(unique_nodes) + 1 print(f"Unique nodes: {n_nodes}") # Aggregate to node-level embeddings print("\n--- Aggregating to node-level embeddings ---") node_emb = aggregate_to_node_embeddings(emb_c1, edges_c1, n_nodes) print(f"Node embeddings: {node_emb.shape}") # Build adjacency matrix from the graph structure print("\n--- Building adjacency matrix ---") adj = np.zeros((n_nodes, n_nodes), dtype=np.float64) # Count edges from training data edge_counts = {} for u, v in edges_c1: edge_counts[(u, v)] = edge_counts.get((u, v), 0) + 1 for (u, v), count in edge_counts.items(): adj[u, v] = count adj[v, u] = count # Make undirected for Laplacian # For star graph with degree=10, subtree_degree=10: # Node 0 is center, nodes 1-10 are leaves # Add explicit star structure for i in range(1, n_nodes): if adj[0, i] == 0: adj[0, i] = 1 adj[i, 0] = 1 print(f"Adjacency matrix shape: {adj.shape}") print(f"Adjacency matrix:\n{adj}") # Compute Laplacian L_norm, L = compute_laplacian(adj) print(f"\nNormalized Laplacian computed") # Compute Fiedler vectors (top 10 eigenvectors) eigenvalues, eigenvectors = compute_fiedler_vectors(L_norm, k=min(10, n_nodes)) if eigenvalues is not None: print(f"\nEigenvalues (sorted): {eigenvalues}") print(f"Fiedler vector (2nd smallest eigenvalue): {eigenvalues[1]:.6f}") # Save Fiedler vectors np.save(f'{OUTPUT_DIR}/fiedler_vectors.npy', eigenvectors) np.save(f'{OUTPUT_DIR}/fiedler_eigenvalues.npy', eigenvalues) # Compute alignment between node embeddings and Fiedler vectors print("\n--- Computing alignment ---") try: alignment = compute_alignment(node_emb, eigenvectors) print(f"Node-level alignment (correlations):") for i, corr in enumerate(alignment): print(f" Spectral dim {i+1}: {corr:.4f}") np.save(f'{OUTPUT_DIR}/alignment_correlations.npy', alignment) except ImportError: print("sklearn not available, using manual correlation") alignment = [] for i in range(min(eigenvectors.shape[1], node_emb.shape[1])): vec = eigenvectors[:, i] emb_dim = node_emb[:, i] if i < node_emb.shape[1] else np.zeros(n_nodes) if np.std(vec) > 1e-10 and np.std(emb_dim) > 1e-10: corr = np.corrcoef(vec, emb_dim)[0, 1] alignment.append(corr) else: alignment.append(0.0) alignment = np.array(alignment) print(f"Manual alignment: {alignment}") # Procrustes alignment print("\n--- Procrustes alignment ---") spectral_reduced = eigenvectors[:, :min(10, node_emb.shape[1])] # Normalize node_emb_norm = node_emb / (np.linalg.norm(node_emb, axis=1, keepdims=True) + 1e-10) spectral_norm = spectral_reduced / (np.linalg.norm(spectral_reduced, axis=1, keepdims=True) + 1e-10) # For Procrustes alignment, project both to same dimensionality first from sklearn.decomposition import PCA pca = PCA(n_components=10) node_emb_pca = pca.fit_transform(node_emb) node_emb_pca_norm = node_emb_pca / (np.linalg.norm(node_emb_pca, axis=1, keepdims=True) + 1e-10) # Now both are 10D - do orthogonal Procrustes # M = node_emb_pca_norm.T @ spectral_norm # (10, 10) M = node_emb_pca_norm.T @ spectral_norm from scipy.linalg import svd U, S_vals, Vt = svd(M) R = Vt.T @ U.T # (10, 10) - orthogonal rotation # Rotate and compute similarity rotated = node_emb_pca_norm @ R similarity = np.sum(rotated * spectral_norm) / (np.linalg.norm(rotated) * np.linalg.norm(spectral_norm) + 1e-10) print(f"Procrustes similarity: {similarity:.4f}") print(f"(Higher = better alignment between spectral and learned structure)") # Save results np.save(f'{OUTPUT_DIR}/procrustes_similarity.npy', similarity) # Visualize: plot Fiedler vectors colored by node embedding norm import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt fig, axes = plt.subplots(2, 5, figsize=(20, 8)) axes = axes.flatten() norms = np.linalg.norm(node_emb, axis=1) for i in range(min(10, eigenvectors.shape[1])): vec = eigenvectors[:, i] scatter = axes[i].scatter(vec, norms, c=range(len(vec)), cmap='viridis', s=100, edgecolors='black', linewidth=0.5) axes[i].set_title(f'Fiedler #{i+1} (λ={eigenvalues[i]:.4f})') axes[i].set_xlabel('Spectral coordinate') axes[i].set_ylabel('Embedding norm') axes[i].grid(True, alpha=0.3) plt.tight_layout() plt.savefig(f'{OUTPUT_DIR}/fiedler_vs_embeddings.png', dpi=150, bbox_inches='tight') plt.close() print(f"\nSaved Fiedler vs embeddings plot to {OUTPUT_DIR}/fiedler_vs_embeddings.png") # Visualize: UMAP of node embeddings colored by Fiedler coordinates try: import umap reducer = umap.UMAP(n_neighbors=min(5, n_nodes-1), min_dist=0.1, metric='cosine') node_emb_2d = reducer.fit_transform(node_emb) fig, axes = plt.subplots(2, 5, figsize=(25, 10)) axes = axes.flatten() for i in range(min(10, eigenvectors.shape[1])): vec = eigenvectors[:, i] scatter = axes[i].scatter(node_emb_2d[:, 0], node_emb_2d[:, 1], c=vec, cmap='viridis', s=100, alpha=0.7, edgecolors='black', linewidth=0.5) axes[i].set_title(f'Fiedler #{i+1} (λ={eigenvalues[i]:.4f})') axes[i].set_xlabel('UMAP 1') axes[i].set_ylabel('UMAP 2') plt.colorbar(scatter, ax=axes[i], fraction=0.046) plt.tight_layout() plt.savefig(f'{OUTPUT_DIR}/umap_colored_by_fiedler.png', dpi=150, bbox_inches='tight') plt.close() print(f"Saved UMAP colored by Fiedler vectors to {OUTPUT_DIR}/umap_colored_by_fiedler.png") except ImportError: print("UMAP not available, skipping UMAP visualization") # Visualize: scatter plot of node embeddings (PCA) colored by spectral coordinate try: from sklearn.decomposition import PCA pca = PCA(n_components=2) node_emb_pca = pca.fit_transform(node_emb) fig, ax = plt.subplots(figsize=(10, 8)) scatter = ax.scatter(node_emb_pca[:, 0], node_emb_pca[:, 1], c=eigenvectors[:, 1], cmap='viridis', s=150, edgecolors='black', linewidth=1) ax.set_title('Node Embeddings (PCA) colored by Fiedler vector') ax.set_xlabel('PCA 1') ax.set_ylabel('PCA 2') plt.colorbar(scatter, ax=ax) # Annotate nodes for i in range(n_nodes): ax.annotate(str(i), (node_emb_pca[i, 0], node_emb_pca[i, 1]), fontsize=12, ha='center', va='center') plt.tight_layout() plt.savefig(f'{OUTPUT_DIR}/node_embeddings_pca.png', dpi=150, bbox_inches='tight') plt.close() print(f"Saved PCA plot to {OUTPUT_DIR}/node_embeddings_pca.png") except ImportError: print("sklearn not available, skipping PCA visualization") print("\n" + "=" * 60) print("Claim 4: Spectral Bias Analysis Complete") print("=" * 60) print(f"Output directory: {OUTPUT_DIR}") print("Files generated:") for f in sorted(os.listdir(OUTPUT_DIR)): size = os.path.getsize(f'{OUTPUT_DIR}/{f}') print(f" - {OUTPUT_DIR}/{f} ({size:,} bytes)") else: print("Failed to compute Fiedler vectors") if __name__ == '__main__': main()