Upload experiment.py with huggingface_hub

experiment.py

@@ -0,0 +1,525 @@
+"""
+Representation Learning Dynamics Experiment
+============================================
+How does a model's internal representation respond to continued training
+on the same task vs. learning a new one?
+
+Experiment design:
+  Phase 1: Train model on Task A (modular addition) until convergence
+  Phase 2: Fork into two branches:
+    Branch A→A: Continue training on Task A (same task, more data)
+    Branch A→B: Switch to Task B (modular subtraction)
+  Track: CKA, subspace angles, gradient alignment, attention entropy,
+         representation variance explained, probing accuracy — all per layer,
+         at every checkpoint.
+
+The key contrast reveals what "learning" looks like at the representation
+level vs. what "forgetting" looks like — and the precise moment they diverge.
+"""
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import numpy as np
+import json
+import os
+import copy
+import time
+from pathlib import Path
+from collections import defaultdict
+from typing import Dict, List, Optional
+
+from model import SmallTransformer, TransformerConfig
+from tasks import (
+    ModularArithmeticDataset, get_probe_data, get_dataloaders,
+    DEFAULT_P, NUM_SPECIAL
+)
+from representation_tracker import (
+    linear_CKA, svcca, subspace_angles, mean_subspace_angle_degrees,
+    gradient_alignment, attention_entropy, task_variance_explained,
+    parameter_delta_cosine, weight_change_magnitude_per_layer,
+    cka_heatmap
+)
+
+
+def evaluate(model, dataloader, device) -> Dict[str, float]:
+    """Evaluate accuracy and loss on a dataset."""
+    model.eval()
+    total_loss = 0
+    total_correct = 0
+    total_count = 0
+
+    with torch.no_grad():
+        for batch in dataloader:
+            input_ids = batch['input_ids'].to(device)
+            labels = batch['labels'].to(device)
+            out = model(input_ids, labels=labels)
+
+            total_loss += out['loss'].item() * input_ids.shape[0]
+            # Accuracy: check last position prediction
+            preds = out['logits'][:, -1, :].argmax(dim=-1)
+            targets = labels[:, -1]
+            total_correct += (preds == targets).sum().item()
+            total_count += input_ids.shape[0]
+
+    return {
+        'loss': total_loss / total_count,
+        'accuracy': total_correct / total_count,
+    }
+
+
+def collect_representations(model, probe_input_ids, device,
+                            token_position: int = -1) -> Dict:
+    """
+    Collect all representation data from a single forward pass on probe data.
+    Returns hidden states, attention patterns, MLP activations.
+    """
+    model.eval()
+    with torch.no_grad():
+        out = model(probe_input_ids.to(device), return_internals=True)
+
+    # Extract at the answer position (last token)
+    hidden_states = [hs[:, token_position, :].cpu()
+                     for hs in out['hidden_states']]
+    attn_weights = [aw.cpu() for aw in out['attn_weights']]
+    mlp_hidden = [mh[:, token_position, :].cpu()
+                  for mh in out['mlp_hidden']]
+
+    return {
+        'hidden_states': hidden_states,
+        'attn_weights': attn_weights,
+        'mlp_hidden': mlp_hidden,
+    }
+
+
+def compute_all_metrics(
+    model, model_init_state, model_phase1_state,
+    reps_current, reps_at_init, reps_at_phase1_end,
+    probe_input_ids_a, probe_labels_a,
+    probe_input_ids_b, probe_labels_b,
+    device, config
+) -> Dict:
+    """
+    Compute the full suite of representation metrics at a single checkpoint.
+    """
+    metrics = {}
+    n_layers = config.n_layers + 1  # +1 for embedding layer
+
+    # === Per-layer metrics ===
+    for layer_idx in range(n_layers):
+        prefix = f'layer_{layer_idx}'
+        curr = reps_current['hidden_states'][layer_idx]
+        init = reps_at_init['hidden_states'][layer_idx]
+        p1 = reps_at_phase1_end['hidden_states'][layer_idx]
+
+        # CKA vs initialization (how far has the representation drifted?)
+        metrics[f'{prefix}/cka_vs_init'] = linear_CKA(curr, init)
+
+        # CKA vs end of Phase 1 (how much has Phase 2 changed things?)
+        metrics[f'{prefix}/cka_vs_phase1'] = linear_CKA(curr, p1)
+
+        # SVCCA for secondary comparison
+        metrics[f'{prefix}/svcca_vs_phase1'] = svcca(curr, p1)
+
+        # Subspace angle vs Phase 1 end
+        k = min(10, curr.shape[0] // 2, curr.shape[1])
+        if k > 0:
+            metrics[f'{prefix}/subspace_angle_vs_phase1'] = \
+                mean_subspace_angle_degrees(curr, p1, k=k)
+        else:
+            metrics[f'{prefix}/subspace_angle_vs_phase1'] = 0.0
+
+    # === Attention entropy per layer ===
+    for layer_idx, aw in enumerate(reps_current['attn_weights']):
+        ent = attention_entropy(aw)
+        metrics[f'layer_{layer_idx+1}/attn_entropy_mean'] = ent['mean_entropy']
+        for h, he in enumerate(ent['per_head_entropy']):
+            metrics[f'layer_{layer_idx+1}/head_{h}_entropy'] = he
+
+    # === Task A representation variance explained ===
+    for layer_idx in range(n_layers):
+        curr = reps_current['hidden_states'][layer_idx]
+        if len(set(probe_labels_a.tolist())) > 1:
+            tve = task_variance_explained(
+                curr, torch.tensor(probe_labels_a, dtype=torch.float), n_components=10
+            )
+            metrics[f'layer_{layer_idx}/task_a_var_explained'] = tve['weighted_r2']
+
+    # === Parameter-space: weight change magnitude ===
+    current_state = {k: v.cpu() for k, v in model.state_dict().items()}
+    wc_from_init = weight_change_magnitude_per_layer(model_init_state, current_state)
+    wc_from_p1 = weight_change_magnitude_per_layer(model_phase1_state, current_state)
+
+    # Aggregate per block
+    for block_idx in range(config.n_layers):
+        init_total = sum(v for k, v in wc_from_init.items()
+                         if f'blocks.{block_idx}' in k)
+        p1_total = sum(v for k, v in wc_from_p1.items()
+                       if f'blocks.{block_idx}' in k)
+        metrics[f'block_{block_idx}/weight_change_from_init'] = init_total
+        metrics[f'block_{block_idx}/weight_change_from_phase1'] = p1_total
+
+    return metrics
+
+
+def train_phase(
+    model, optimizer, dataloader, n_epochs: int,
+    device, phase_name: str,
+    # For metric collection
+    model_init_state, model_phase1_state,
+    reps_at_init, reps_at_phase1_end,
+    probe_input_ids_a, probe_labels_a,
+    probe_input_ids_b, probe_labels_b,
+    eval_loaders: Dict,
+    config: TransformerConfig,
+    checkpoint_every: int = 50,  # steps between metric collection
+    output_dir: str = 'results',
+) -> List[Dict]:
+    """
+    Train for n_epochs, collecting representation metrics periodically.
+    """
+    history = []
+    global_step = 0
+    os.makedirs(output_dir, exist_ok=True)
+
+    for epoch in range(n_epochs):
+        model.train()
+        epoch_loss = 0
+        n_batches = 0
+
+        for batch in dataloader:
+            input_ids = batch['input_ids'].to(device)
+            labels = batch['labels'].to(device)
+
+            out = model(input_ids, labels=labels)
+            loss = out['loss']
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            epoch_loss += loss.item()
+            n_batches += 1
+            global_step += 1
+
+            # Collect metrics at checkpoint intervals
+            if global_step % checkpoint_every == 0:
+                model.eval()
+
+                # Get current representations on probe data
+                reps_current = collect_representations(
+                    model, probe_input_ids_a, device
+                )
+
+                # Compute all metrics
+                step_metrics = compute_all_metrics(
+                    model, model_init_state, model_phase1_state,
+                    reps_current, reps_at_init, reps_at_phase1_end,
+                    probe_input_ids_a, probe_labels_a,
+                    probe_input_ids_b, probe_labels_b,
+                    device, config
+                )
+
+                # Evaluate on all datasets
+                for name, loader in eval_loaders.items():
+                    eval_res = evaluate(model, loader, device)
+                    step_metrics[f'eval/{name}_loss'] = eval_res['loss']
+                    step_metrics[f'eval/{name}_acc'] = eval_res['accuracy']
+
+                # Gradient alignment between tasks
+                # Get one batch from each task
+                batch_a = next(iter(eval_loaders['add_test']))
+                batch_b = next(iter(eval_loaders['subtract_test']))
+
+                def loss_fn(m, b):
+                    return m(b['input_ids'].to(device),
+                             labels=b['labels'].to(device))['loss']
+
+                try:
+                    ga = gradient_alignment(model, batch_a, batch_b, loss_fn)
+                    step_metrics['gradient_alignment_a_vs_b'] = ga
+                except Exception:
+                    step_metrics['gradient_alignment_a_vs_b'] = 0.0
+
+                step_metrics['phase'] = phase_name
+                step_metrics['epoch'] = epoch
+                step_metrics['step'] = global_step
+                step_metrics['train_loss'] = epoch_loss / n_batches
+
+                history.append(step_metrics)
+
+                print(f"[{phase_name}] Step {global_step} | "
+                      f"Loss: {epoch_loss/n_batches:.4f} | "
+                      f"Add acc: {step_metrics.get('eval/add_test_acc', 0):.3f} | "
+                      f"Sub acc: {step_metrics.get('eval/subtract_test_acc', 0):.3f} | "
+                      f"CKA(L1 vs P1): {step_metrics.get('layer_1/cka_vs_phase1', 0):.3f} | "
+                      f"Grad align: {step_metrics.get('gradient_alignment_a_vs_b', 0):.3f}")
+
+                model.train()
+
+        # End of epoch eval
+        print(f"[{phase_name}] Epoch {epoch+1}/{n_epochs} complete, "
+              f"avg loss: {epoch_loss/n_batches:.4f}")
+
+    return history
+
+
+def run_experiment(
+    p: int = DEFAULT_P,
+    n_layers: int = 2,
+    d_model: int = 128,
+    n_heads: int = 4,
+    d_mlp: int = 512,
+    phase1_epochs: int = 100,
+    phase2_epochs: int = 100,
+    lr: float = 1e-3,
+    weight_decay: float = 1.0,
+    batch_size: int = 512,
+    train_frac: float = 0.5,
+    checkpoint_every: int = 50,
+    output_dir: str = 'results',
+    seed: int = 42,
+):
+    """
+    Run the full two-phase experiment.
+    """
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Using device: {device}")
+
+    # === Setup ===
+    config = TransformerConfig(
+        vocab_size=p + NUM_SPECIAL,
+        n_layers=n_layers,
+        d_model=d_model,
+        n_heads=n_heads,
+        d_mlp=d_mlp,
+        max_seq_len=5,
+    )
+
+    model = SmallTransformer(config).to(device)
+    print(f"Model parameters: {model.count_parameters():,}")
+
+    # Save initial state
+    model_init_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
+
+    # Dataloaders
+    loaders = get_dataloaders(p=p, batch_size=batch_size,
+                               train_frac=train_frac, seed=seed)
+
+    # Probe datasets (fixed subsets for consistent metric computation)
+    ds_a = ModularArithmeticDataset('add', p=p, split='test', train_frac=train_frac, seed=seed)
+    ds_b = ModularArithmeticDataset('subtract', p=p, split='test', train_frac=train_frac, seed=seed)
+    probe_ids_a, probe_labels_a = get_probe_data(ds_a, n_samples=min(500, len(ds_a)))
+    probe_ids_b, probe_labels_b = get_probe_data(ds_b, n_samples=min(500, len(ds_b)))
+
+    # Initial representations
+    reps_at_init = collect_representations(model, probe_ids_a, device)
+
+    # ===========================
+    # PHASE 1: Train on Task A
+    # ===========================
+    print("\n" + "=" * 60)
+    print("PHASE 1: Training on Task A (Modular Addition)")
+    print("=" * 60)
+
+    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
+
+    # Dummy phase1 state for Phase 1 tracking (use init)
+    phase1_history = train_phase(
+        model=model,
+        optimizer=optimizer,
+        dataloader=loaders['add_train'],
+        n_epochs=phase1_epochs,
+        device=device,
+        phase_name='phase1_add',
+        model_init_state=model_init_state,
+        model_phase1_state=model_init_state,  # placeholder
+        reps_at_init=reps_at_init,
+        reps_at_phase1_end=reps_at_init,  # placeholder
+        probe_input_ids_a=probe_ids_a,
+        probe_labels_a=probe_labels_a,
+        probe_input_ids_b=probe_ids_b,
+        probe_labels_b=probe_labels_b,
+        eval_loaders=loaders,
+        config=config,
+        checkpoint_every=checkpoint_every,
+        output_dir=output_dir,
+    )
+
+    # Save Phase 1 endpoint
+    model_phase1_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
+    reps_at_phase1_end = collect_representations(model, probe_ids_a, device)
+    phase1_final_eval = evaluate(model, loaders['add_test'], device)
+    print(f"\nPhase 1 final — Add accuracy: {phase1_final_eval['accuracy']:.3f}")
+
+    # Save Phase 1 checkpoint
+    torch.save(model.state_dict(), os.path.join(output_dir, 'phase1_checkpoint.pt'))
+
+    # ===========================
+    # PHASE 2: Fork into A→A and A→B
+    # ===========================
+
+    # Branch A→A: Continue on same task
+    print("\n" + "=" * 60)
+    print("PHASE 2a: Branch A→A (Continue training on Addition)")
+    print("=" * 60)
+
+    model_aa = SmallTransformer(config).to(device)
+    model_aa.load_state_dict(torch.load(os.path.join(output_dir, 'phase1_checkpoint.pt'),
+                                         weights_only=True))
+    optimizer_aa = optim.AdamW(model_aa.parameters(), lr=lr, weight_decay=weight_decay)
+
+    history_aa = train_phase(
+        model=model_aa,
+        optimizer=optimizer_aa,
+        dataloader=loaders['add_train'],
+        n_epochs=phase2_epochs,
+        device=device,
+        phase_name='phase2_aa',
+        model_init_state=model_init_state,
+        model_phase1_state=model_phase1_state,
+        reps_at_init=reps_at_init,
+        reps_at_phase1_end=reps_at_phase1_end,
+        probe_input_ids_a=probe_ids_a,
+        probe_labels_a=probe_labels_a,
+        probe_input_ids_b=probe_ids_b,
+        probe_labels_b=probe_labels_b,
+        eval_loaders=loaders,
+        config=config,
+        checkpoint_every=checkpoint_every,
+        output_dir=output_dir,
+    )
+
+    # Branch A→B: Switch to new task
+    print("\n" + "=" * 60)
+    print("PHASE 2b: Branch A→B (Switch to Subtraction)")
+    print("=" * 60)
+
+    model_ab = SmallTransformer(config).to(device)
+    model_ab.load_state_dict(torch.load(os.path.join(output_dir, 'phase1_checkpoint.pt'),
+                                         weights_only=True))
+    optimizer_ab = optim.AdamW(model_ab.parameters(), lr=lr, weight_decay=weight_decay)
+
+    history_ab = train_phase(
+        model=model_ab,
+        optimizer=optimizer_ab,
+        dataloader=loaders['subtract_train'],
+        n_epochs=phase2_epochs,
+        device=device,
+        phase_name='phase2_ab',
+        model_init_state=model_init_state,
+        model_phase1_state=model_phase1_state,
+        reps_at_init=reps_at_init,
+        reps_at_phase1_end=reps_at_phase1_end,
+        probe_input_ids_a=probe_ids_a,
+        probe_labels_a=probe_labels_a,
+        probe_input_ids_b=probe_ids_b,
+        probe_labels_b=probe_labels_b,
+        eval_loaders=loaders,
+        config=config,
+        checkpoint_every=checkpoint_every,
+        output_dir=output_dir,
+    )
+
+    # ===========================
+    # PHASE 3: Cross-model comparison
+    # ===========================
+    print("\n" + "=" * 60)
+    print("PHASE 3: Cross-model representation comparison")
+    print("=" * 60)
+
+    reps_aa = collect_representations(model_aa, probe_ids_a, device)
+    reps_ab = collect_representations(model_ab, probe_ids_a, device)
+
+    cross_metrics = {}
+    for layer_idx in range(config.n_layers + 1):
+        ha = reps_aa['hidden_states'][layer_idx]
+        hb = reps_ab['hidden_states'][layer_idx]
+        hp1 = reps_at_phase1_end['hidden_states'][layer_idx]
+
+        cross_metrics[f'layer_{layer_idx}/cka_aa_vs_ab'] = linear_CKA(ha, hb)
+        cross_metrics[f'layer_{layer_idx}/cka_aa_vs_p1'] = linear_CKA(ha, hp1)
+        cross_metrics[f'layer_{layer_idx}/cka_ab_vs_p1'] = linear_CKA(hb, hp1)
+        cross_metrics[f'layer_{layer_idx}/subspace_angle_aa_vs_ab'] = \
+            mean_subspace_angle_degrees(ha, hb, k=min(10, ha.shape[0] // 2, ha.shape[1]))
+
+    # CKA heatmaps
+    heatmap_aa_vs_ab = cka_heatmap(reps_aa['hidden_states'], reps_ab['hidden_states'])
+    heatmap_aa_vs_p1 = cka_heatmap(reps_aa['hidden_states'],
+                                    reps_at_phase1_end['hidden_states'])
+    heatmap_ab_vs_p1 = cka_heatmap(reps_ab['hidden_states'],
+                                    reps_at_phase1_end['hidden_states'])
+
+    # Parameter delta cosine
+    params_init = [v for v in model_init_state.values()]
+    params_aa = [v.cpu() for v in model_aa.state_dict().values()]
+    params_ab = [v.cpu() for v in model_ab.state_dict().values()]
+    params_p1 = [v for v in model_phase1_state.values()]
+
+    cross_metrics['param_delta_cosine_aa_vs_ab'] = \
+        parameter_delta_cosine(params_p1, params_aa, params_ab)
+    cross_metrics['param_delta_cosine_aa_vs_p1_from_init'] = \
+        parameter_delta_cosine(params_init, params_p1, params_aa)
+
+    print("\n=== Cross-model metrics ===")
+    for k, v in sorted(cross_metrics.items()):
+        print(f"  {k}: {v:.4f}")
+
+    # ===========================
+    # Save all results
+    # ===========================
+    results = {
+        'config': {
+            'p': p, 'n_layers': n_layers, 'd_model': d_model,
+            'n_heads': n_heads, 'd_mlp': d_mlp,
+            'phase1_epochs': phase1_epochs, 'phase2_epochs': phase2_epochs,
+            'lr': lr, 'weight_decay': weight_decay, 'batch_size': batch_size,
+            'train_frac': train_frac, 'seed': seed,
+            'n_parameters': model.count_parameters(),
+        },
+        'phase1_history': phase1_history,
+        'phase2_aa_history': history_aa,
+        'phase2_ab_history': history_ab,
+        'cross_metrics': cross_metrics,
+        'cka_heatmaps': {
+            'aa_vs_ab': heatmap_aa_vs_ab.tolist(),
+            'aa_vs_p1': heatmap_aa_vs_p1.tolist(),
+            'ab_vs_p1': heatmap_ab_vs_p1.tolist(),
+        },
+    }
+
+    results_path = os.path.join(output_dir, 'experiment_results.json')
+    with open(results_path, 'w') as f:
+        json.dump(results, f, indent=2, default=str)
+    print(f"\nResults saved to {results_path}")
+
+    # Save final models
+    torch.save(model_aa.state_dict(), os.path.join(output_dir, 'model_aa_final.pt'))
+    torch.save(model_ab.state_dict(), os.path.join(output_dir, 'model_ab_final.pt'))
+
+    return results
+
+
+if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--p', type=int, default=DEFAULT_P)
+    parser.add_argument('--n-layers', type=int, default=2)
+    parser.add_argument('--d-model', type=int, default=128)
+    parser.add_argument('--n-heads', type=int, default=4)
+    parser.add_argument('--d-mlp', type=int, default=512)
+    parser.add_argument('--phase1-epochs', type=int, default=100)
+    parser.add_argument('--phase2-epochs', type=int, default=100)
+    parser.add_argument('--lr', type=float, default=1e-3)
+    parser.add_argument('--weight-decay', type=float, default=1.0)
+    parser.add_argument('--batch-size', type=int, default=512)
+    parser.add_argument('--train-frac', type=float, default=0.5)
+    parser.add_argument('--checkpoint-every', type=int, default=50)
+    parser.add_argument('--output-dir', type=str, default='results')
+    parser.add_argument('--seed', type=int, default=42)
+    args = parser.parse_args()
+
+    run_experiment(**vars(args))