Upload extended_experiment_v2.py with huggingface_hub

extended_experiment_v2.py

@@ -0,0 +1,519 @@
+"""
+Extended Gradient Clipping Experiment V2: Testing Physics-of-AI Predictions
+
+Key changes from V1:
+1. More epochs (10 instead of 3) to allow rare class learning
+2. Smaller learning rate (0.01) for more stable training
+3. More frequent tracking to catch dynamics
+4. Added loss tracking per class to understand learning dynamics
+
+Predictions being tested:
+- Prediction 2: Representation Collapse (effective dimensionality drops without clipping)
+- Prediction 4: Rare Sample Learning (clipping improves rare class accuracy)
+"""
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import numpy as np
+import matplotlib.pyplot as plt
+import random
+from typing import Dict, List, Tuple
+
+SEED = 42
+
+
+def set_seeds(seed=SEED):
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+
+
+class SimpleNextTokenModel(nn.Module):
+    def __init__(self, vocab_size=4, embedding_dim=16):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, embedding_dim)
+        self.linear = nn.Linear(embedding_dim, vocab_size)
+    
+    def forward(self, x):
+        embedded = self.embedding(x)
+        logits = self.linear(embedded)
+        return logits
+    
+    def get_embeddings(self):
+        return self.embedding.weight.data.clone()
+
+
+def compute_effective_dimension(embedding_matrix: torch.Tensor) -> float:
+    """PCA-based effective dimensionality using entropy."""
+    centered = embedding_matrix - embedding_matrix.mean(dim=0, keepdim=True)
+    cov = torch.mm(centered.T, centered) / (embedding_matrix.shape[0] - 1)
+    eigenvalues = torch.linalg.eigvalsh(cov)
+    eigenvalues = torch.clamp(eigenvalues, min=1e-10)
+    eigenvalues = eigenvalues / eigenvalues.sum()
+    entropy = -torch.sum(eigenvalues * torch.log(eigenvalues))
+    return torch.exp(entropy).item()
+
+
+def compute_per_class_accuracy(model: nn.Module, inputs: torch.Tensor, 
+                                targets: torch.Tensor) -> Dict[int, float]:
+    """Compute accuracy for each target class."""
+    model.eval()
+    with torch.no_grad():
+        logits = model(inputs)
+        predictions = logits.argmax(dim=1)
+    
+    accuracies = {}
+    for class_idx in range(4):
+        mask = targets == class_idx
+        if mask.sum() > 0:
+            correct = (predictions[mask] == targets[mask]).float().mean().item()
+            accuracies[class_idx] = correct
+        else:
+            accuracies[class_idx] = None
+    
+    return accuracies
+
+
+def compute_per_class_loss(model: nn.Module, inputs: torch.Tensor, 
+                           targets: torch.Tensor, criterion: nn.Module) -> Dict[int, float]:
+    """Compute average loss for each target class."""
+    model.eval()
+    losses = {}
+    with torch.no_grad():
+        logits = model(inputs)
+        for class_idx in range(4):
+            mask = targets == class_idx
+            if mask.sum() > 0:
+                class_loss = criterion(logits[mask], targets[mask]).item()
+                losses[class_idx] = class_loss
+            else:
+                losses[class_idx] = None
+    return losses
+
+
+def create_imbalanced_dataset(n_samples=1000, n_rare=10, seed=SEED):
+    set_seeds(seed)
+    inputs = torch.randint(0, 4, (n_samples,))
+    targets = torch.zeros(n_samples, dtype=torch.long)
+    rare_indices = random.sample(range(n_samples), n_rare)
+    targets[rare_indices] = 1
+    return inputs, targets, sorted(rare_indices)
+
+
+def train_with_tracking(inputs: torch.Tensor, targets: torch.Tensor, 
+                        rare_indices: List[int], clip_grad: bool = False, 
+                        max_norm: float = 1.0, n_epochs: int = 10, 
+                        lr: float = 0.01, init_weights=None,
+                        track_every: int = 50) -> Dict:
+    """
+    Extended training with comprehensive tracking.
+    """
+    set_seeds(SEED)
+    model = SimpleNextTokenModel(vocab_size=4, embedding_dim=16)
+    if init_weights:
+        model.load_state_dict({k: v.clone() for k, v in init_weights.items()})
+    
+    optimizer = optim.SGD(model.parameters(), lr=lr)
+    criterion = nn.CrossEntropyLoss()
+    
+    metrics = {
+        'losses': [],
+        'grad_norms': [],
+        'weight_norms': [],
+        'effective_dims': [],
+        'effective_dim_steps': [],
+        'class_accuracies': {0: [], 1: [], 2: [], 3: []},
+        'class_losses': {0: [], 1: [], 2: [], 3: []},
+        'accuracy_steps': [],
+        'rare_sample_losses': [],  # Track loss specifically at rare samples
+        'rare_sample_steps': [],
+    }
+    
+    mode = "WITH" if clip_grad else "WITHOUT"
+    print(f"\n{'='*60}")
+    print(f"Training {mode} gradient clipping (max_norm={max_norm})")
+    print(f"Learning rate: {lr}, Epochs: {n_epochs}")
+    print(f"{'='*60}")
+    
+    step = 0
+    n_samples = len(inputs)
+    
+    for epoch in range(n_epochs):
+        model.train()
+        epoch_losses = []
+        
+        for i in range(n_samples):
+            x = inputs[i:i+1]
+            y = targets[i:i+1]
+            
+            optimizer.zero_grad()
+            logits = model(x)
+            loss = criterion(logits, y)
+            loss.backward()
+            
+            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), float('inf'))
+            
+            if clip_grad:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
+            
+            optimizer.step()
+            
+            metrics['losses'].append(loss.item())
+            metrics['grad_norms'].append(grad_norm.item())
+            
+            total_norm = sum(p.data.norm(2).item() ** 2 for p in model.parameters()) ** 0.5
+            metrics['weight_norms'].append(total_norm)
+            
+            epoch_losses.append(loss.item())
+            
+            # Track at rare sample positions
+            if i in rare_indices:
+                metrics['rare_sample_losses'].append(loss.item())
+                metrics['rare_sample_steps'].append(step)
+            
+            # Track embedding stats and accuracy periodically
+            if step % track_every == 0:
+                emb_matrix = model.get_embeddings()
+                eff_dim = compute_effective_dimension(emb_matrix)
+                
+                metrics['effective_dims'].append(eff_dim)
+                metrics['effective_dim_steps'].append(step)
+                
+                class_acc = compute_per_class_accuracy(model, inputs, targets)
+                class_loss = compute_per_class_loss(model, inputs, targets, criterion)
+                
+                for cls_idx in range(4):
+                    if class_acc[cls_idx] is not None:
+                        metrics['class_accuracies'][cls_idx].append(class_acc[cls_idx])
+                    else:
+                        metrics['class_accuracies'][cls_idx].append(0.0)
+                    
+                    if class_loss[cls_idx] is not None:
+                        metrics['class_losses'][cls_idx].append(class_loss[cls_idx])
+                    else:
+                        metrics['class_losses'][cls_idx].append(0.0)
+                
+                metrics['accuracy_steps'].append(step)
+            
+            step += 1
+        
+        avg_loss = np.mean(epoch_losses)
+        class_acc = compute_per_class_accuracy(model, inputs, targets)
+        class_loss = compute_per_class_loss(model, inputs, targets, criterion)
+        eff_dim = compute_effective_dimension(model.get_embeddings())
+        
+        b_acc = f"{class_acc[1]:.3f}" if class_acc[1] is not None else "N/A"
+        b_loss = f"{class_loss[1]:.3f}" if class_loss[1] is not None else "N/A"
+        
+        print(f"Epoch {epoch+1:2d}/{n_epochs}: Loss={avg_loss:.4f} | "
+              f"Acc A={class_acc[0]:.3f} B={b_acc} | "
+              f"Loss A={class_loss[0]:.3f} B={b_loss} | "
+              f"EffDim={eff_dim:.3f}")
+    
+    return metrics
+
+
+def plot_comprehensive_analysis(metrics_no_clip: Dict, metrics_with_clip: Dict,
+                                 rare_indices: List[int], filename: str,
+                                 n_samples: int = 1000):
+    """Create comprehensive 8-panel analysis."""
+    fig = plt.figure(figsize=(20, 16))
+    gs = fig.add_gridspec(4, 2, hspace=0.35, wspace=0.25)
+    
+    n_epochs = len(metrics_no_clip['losses']) // n_samples
+    
+    # Row 1: Effective Dimension
+    ax1 = fig.add_subplot(gs[0, 0])
+    ax2 = fig.add_subplot(gs[0, 1])
+    
+    ax1.plot(metrics_no_clip['effective_dim_steps'], metrics_no_clip['effective_dims'], 
+             'b-', linewidth=2, marker='o', markersize=3)
+    ax1.set_ylabel('Effective Dimension', fontsize=11)
+    ax1.set_title('Effective Dim - WITHOUT Clipping', fontsize=12, fontweight='bold', color='red')
+    ax1.grid(True, alpha=0.3)
+    ax1.set_ylim([2.0, 3.5])
+    
+    ax2.plot(metrics_with_clip['effective_dim_steps'], metrics_with_clip['effective_dims'], 
+             'g-', linewidth=2, marker='o', markersize=3)
+    ax2.set_title('Effective Dim - WITH Clipping', fontsize=12, fontweight='bold', color='green')
+    ax2.grid(True, alpha=0.3)
+    ax2.set_ylim([2.0, 3.5])
+    
+    # Row 2: Class Accuracies
+    ax3 = fig.add_subplot(gs[1, 0])
+    ax4 = fig.add_subplot(gs[1, 1])
+    
+    ax3.plot(metrics_no_clip['accuracy_steps'], metrics_no_clip['class_accuracies'][0], 
+             'r-', linewidth=2, alpha=0.7, label='Without Clip')
+    ax3.plot(metrics_with_clip['accuracy_steps'], metrics_with_clip['class_accuracies'][0], 
+             'g-', linewidth=2, alpha=0.7, label='With Clip')
+    ax3.set_ylabel('Accuracy', fontsize=11)
+    ax3.set_title("Common Class 'A' Accuracy", fontsize=12, fontweight='bold')
+    ax3.legend()
+    ax3.grid(True, alpha=0.3)
+    ax3.set_ylim([0, 1.05])
+    
+    ax4.plot(metrics_no_clip['accuracy_steps'], metrics_no_clip['class_accuracies'][1], 
+             'r-', linewidth=2, alpha=0.7, label='Without Clip')
+    ax4.plot(metrics_with_clip['accuracy_steps'], metrics_with_clip['class_accuracies'][1], 
+             'g-', linewidth=2, alpha=0.7, label='With Clip')
+    ax4.set_title("Rare Class 'B' Accuracy [KEY PREDICTION]", fontsize=12, fontweight='bold', color='purple')
+    ax4.legend()
+    ax4.grid(True, alpha=0.3)
+    ax4.set_ylim([0, 1.05])
+    
+    # Row 3: Class Losses
+    ax5 = fig.add_subplot(gs[2, 0])
+    ax6 = fig.add_subplot(gs[2, 1])
+    
+    ax5.plot(metrics_no_clip['accuracy_steps'], metrics_no_clip['class_losses'][0], 
+             'r-', linewidth=2, alpha=0.7, label='Without Clip')
+    ax5.plot(metrics_with_clip['accuracy_steps'], metrics_with_clip['class_losses'][0], 
+             'g-', linewidth=2, alpha=0.7, label='With Clip')
+    ax5.set_ylabel('Loss', fontsize=11)
+    ax5.set_title("Common Class 'A' Loss", fontsize=12, fontweight='bold')
+    ax5.legend()
+    ax5.grid(True, alpha=0.3)
+    
+    ax6.plot(metrics_no_clip['accuracy_steps'], metrics_no_clip['class_losses'][1], 
+             'r-', linewidth=2, alpha=0.7, label='Without Clip')
+    ax6.plot(metrics_with_clip['accuracy_steps'], metrics_with_clip['class_losses'][1], 
+             'g-', linewidth=2, alpha=0.7, label='With Clip')
+    ax6.set_title("Rare Class 'B' Loss", fontsize=12, fontweight='bold')
+    ax6.legend()
+    ax6.grid(True, alpha=0.3)
+    
+    # Row 4: Gradient Norms and Weight Norms
+    ax7 = fig.add_subplot(gs[3, 0])
+    ax8 = fig.add_subplot(gs[3, 1])
+    
+    steps = range(len(metrics_no_clip['grad_norms']))
+    
+    ax7.plot(steps, metrics_no_clip['grad_norms'], 'r-', alpha=0.3, linewidth=0.5, label='Without Clip')
+    ax7.plot(steps, metrics_with_clip['grad_norms'], 'g-', alpha=0.3, linewidth=0.5, label='With Clip')
+    ax7.axhline(y=1.0, color='black', linestyle='--', linewidth=2, label='Clip threshold')
+    ax7.set_ylabel('Gradient Norm', fontsize=11)
+    ax7.set_xlabel('Training Step', fontsize=11)
+    ax7.set_title('Gradient Norms', fontsize=12, fontweight='bold')
+    ax7.legend()
+    ax7.grid(True, alpha=0.3)
+    
+    ax8.plot(steps, metrics_no_clip['weight_norms'], 'r-', alpha=0.7, linewidth=1, label='Without Clip')
+    ax8.plot(steps, metrics_with_clip['weight_norms'], 'g-', alpha=0.7, linewidth=1, label='With Clip')
+    ax8.set_xlabel('Training Step', fontsize=11)
+    ax8.set_title('Weight Norms', fontsize=12, fontweight='bold')
+    ax8.legend()
+    ax8.grid(True, alpha=0.3)
+    
+    fig.suptitle('Extended Gradient Clipping Analysis: Testing Physics-of-AI Predictions\n'
+                 f'(10 epochs, lr=0.01, 990 common / 10 rare samples)',
+                 fontsize=14, fontweight='bold', y=1.01)
+    
+    plt.savefig(filename, dpi=150, bbox_inches='tight')
+    plt.close()
+    print(f"Comprehensive analysis saved to: {filename}")
+
+
+def plot_rare_sample_dynamics(metrics_no_clip: Dict, metrics_with_clip: Dict,
+                               filename: str):
+    """Plot dynamics specifically at rare sample positions."""
+    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
+    
+    # Rare sample losses over time
+    ax1 = axes[0, 0]
+    ax1.plot(metrics_no_clip['rare_sample_steps'], metrics_no_clip['rare_sample_losses'], 
+             'ro-', alpha=0.7, markersize=3, linewidth=0.5, label='Without Clip')
+    ax1.plot(metrics_with_clip['rare_sample_steps'], metrics_with_clip['rare_sample_losses'], 
+             'go-', alpha=0.7, markersize=3, linewidth=0.5, label='With Clip')
+    ax1.set_ylabel('Loss at Rare Sample', fontsize=11)
+    ax1.set_title('Loss When Encountering Rare Samples', fontsize=12, fontweight='bold')
+    ax1.legend()
+    ax1.grid(True, alpha=0.3)
+    
+    # Histogram of rare sample losses
+    ax2 = axes[0, 1]
+    ax2.hist(metrics_no_clip['rare_sample_losses'], bins=30, alpha=0.5, color='red', 
+             label=f"Without Clip (mean={np.mean(metrics_no_clip['rare_sample_losses']):.3f})")
+    ax2.hist(metrics_with_clip['rare_sample_losses'], bins=30, alpha=0.5, color='green',
+             label=f"With Clip (mean={np.mean(metrics_with_clip['rare_sample_losses']):.3f})")
+    ax2.set_xlabel('Loss', fontsize=11)
+    ax2.set_ylabel('Count', fontsize=11)
+    ax2.set_title('Distribution of Rare Sample Losses', fontsize=12, fontweight='bold')
+    ax2.legend()
+    ax2.grid(True, alpha=0.3)
+    
+    # Gradient norms at rare positions
+    ax3 = axes[1, 0]
+    
+    # Extract gradient norms at rare sample positions
+    n_samples = 1000
+    n_epochs = len(metrics_no_clip['losses']) // n_samples
+    rare_indices = [25, 104, 114, 142, 228, 250, 281, 654, 754, 759]  # From our dataset
+    
+    rare_grad_norms_no = []
+    rare_grad_norms_with = []
+    rare_steps = []
+    
+    for epoch in range(n_epochs):
+        for idx in rare_indices:
+            step = epoch * n_samples + idx
+            if step < len(metrics_no_clip['grad_norms']):
+                rare_grad_norms_no.append(metrics_no_clip['grad_norms'][step])
+                rare_grad_norms_with.append(metrics_with_clip['grad_norms'][step])
+                rare_steps.append(step)
+    
+    ax3.scatter(rare_steps, rare_grad_norms_no, c='red', alpha=0.6, s=20, label='Without Clip')
+    ax3.scatter(rare_steps, rare_grad_norms_with, c='green', alpha=0.6, s=20, label='With Clip')
+    ax3.axhline(y=1.0, color='black', linestyle='--', linewidth=2, label='Clip threshold')
+    ax3.set_xlabel('Training Step', fontsize=11)
+    ax3.set_ylabel('Gradient Norm', fontsize=11)
+    ax3.set_title('Gradient Norms at Rare Sample Positions', fontsize=12, fontweight='bold')
+    ax3.legend()
+    ax3.grid(True, alpha=0.3)
+    
+    # Summary statistics
+    ax4 = axes[1, 1]
+    ax4.axis('off')
+    
+    mean_rare_loss_no = np.mean(metrics_no_clip['rare_sample_losses'])
+    mean_rare_loss_with = np.mean(metrics_with_clip['rare_sample_losses'])
+    mean_rare_grad_no = np.mean(rare_grad_norms_no)
+    mean_rare_grad_with = np.mean(rare_grad_norms_with)
+    
+    # Final class B accuracy
+    final_acc_b_no = metrics_no_clip['class_accuracies'][1][-1] if metrics_no_clip['class_accuracies'][1] else 0
+    final_acc_b_with = metrics_with_clip['class_accuracies'][1][-1] if metrics_with_clip['class_accuracies'][1] else 0
+    
+    summary_text = f"""
+    RARE SAMPLE DYNAMICS SUMMARY
+    ════════════════════════════════════════════════════
+    
+    At Rare Sample Positions:
+    ─────────────────────────────────────────────────────
+    Mean Loss (WITHOUT Clipping):  {mean_rare_loss_no:.4f}
+    Mean Loss (WITH Clipping):     {mean_rare_loss_with:.4f}
+    Loss Reduction:                {(mean_rare_loss_no - mean_rare_loss_with) / mean_rare_loss_no * 100:+.1f}%
+    
+    Mean Gradient Norm (WITHOUT):  {mean_rare_grad_no:.4f}
+    Mean Gradient Norm (WITH):     {mean_rare_grad_with:.4f}
+    Gradient Reduction:            {(mean_rare_grad_no - mean_rare_grad_with) / mean_rare_grad_no * 100:+.1f}%
+    
+    Final Rare Class Accuracy:
+    ─────────────────────────────────────────────────────
+    WITHOUT Clipping: {final_acc_b_no:.1%}
+    WITH Clipping:    {final_acc_b_with:.1%}
+    
+    ════════════════════════════════════════════════════
+    PHYSICS-OF-AI INTERPRETATION:
+    
+    Gradient clipping acts as a "velocity limiter" in
+    weight space, preventing the model from making
+    sudden large updates when encountering rare samples.
+    
+    This allows the model to gradually learn the rare
+    class pattern rather than overshooting and forgetting.
+    """
+    
+    ax4.text(0.05, 0.5, summary_text, transform=ax4.transAxes,
+             fontsize=10, verticalalignment='center', fontfamily='monospace',
+             bbox=dict(boxstyle='round', facecolor='lightyellow', alpha=0.9))
+    
+    fig.suptitle('Rare Sample Dynamics Analysis\n'
+                 '(How the model behaves when encountering rare class B samples)',
+                 fontsize=14, fontweight='bold', y=1.01)
+    
+    plt.tight_layout()
+    plt.savefig(filename, dpi=150, bbox_inches='tight')
+    plt.close()
+    print(f"Rare sample dynamics plot saved to: {filename}")
+
+
+def main():
+    print("="*70)
+    print("EXTENDED GRADIENT CLIPPING EXPERIMENT V2")
+    print("Testing Physics-of-AI Predictions with Extended Training")
+    print("="*70)
+    
+    # Create dataset
+    inputs, targets, rare_indices = create_imbalanced_dataset(n_samples=1000, n_rare=10, seed=SEED)
+    
+    print(f"\nDataset: {len(inputs)} samples ({(targets == 0).sum().item()} common, {(targets == 1).sum().item()} rare)")
+    print(f"Rare indices: {rare_indices}")
+    
+    # Get initial weights
+    set_seeds(SEED)
+    init_model = SimpleNextTokenModel(vocab_size=4, embedding_dim=16)
+    init_weights = {name: param.clone() for name, param in init_model.state_dict().items()}
+    
+    init_eff_dim = compute_effective_dimension(init_model.get_embeddings())
+    print(f"Initial effective dimension: {init_eff_dim:.3f}")
+    
+    # Training parameters
+    n_epochs = 10
+    lr = 0.01
+    
+    # Run training WITHOUT gradient clipping
+    metrics_no_clip = train_with_tracking(
+        inputs, targets, rare_indices,
+        clip_grad=False, n_epochs=n_epochs, lr=lr, 
+        init_weights=init_weights, track_every=100
+    )
+    
+    # Run training WITH gradient clipping
+    metrics_with_clip = train_with_tracking(
+        inputs, targets, rare_indices,
+        clip_grad=True, max_norm=1.0, n_epochs=n_epochs, lr=lr,
+        init_weights=init_weights, track_every=100
+    )
+    
+    # Generate plots
+    print("\n" + "="*70)
+    print("GENERATING ANALYSIS PLOTS")
+    print("="*70)
+    
+    plot_comprehensive_analysis(
+        metrics_no_clip, metrics_with_clip, rare_indices,
+        "extended_analysis_v2.png"
+    )
+    
+    plot_rare_sample_dynamics(
+        metrics_no_clip, metrics_with_clip,
+        "rare_sample_dynamics.png"
+    )
+    
+    # Final summary
+    print("\n" + "="*70)
+    print("FINAL PREDICTION TEST RESULTS")
+    print("="*70)
+    
+    # Prediction 2
+    dims_no = metrics_no_clip['effective_dims']
+    dims_with = metrics_with_clip['effective_dims']
+    
+    print("\n[PREDICTION 2] Representation Collapse:")
+    print(f"  Effective Dim Variance (WITHOUT): {np.std(dims_no):.6f}")
+    print(f"  Effective Dim Variance (WITH):    {np.std(dims_with):.6f}")
+    print(f"  Verdict: {'SUPPORTED' if np.std(dims_no) > np.std(dims_with) else 'NOT SUPPORTED'}")
+    
+    # Prediction 4
+    final_acc_b_no = metrics_no_clip['class_accuracies'][1][-1]
+    final_acc_b_with = metrics_with_clip['class_accuracies'][1][-1]
+    
+    print("\n[PREDICTION 4] Rare Sample Learning:")
+    print(f"  Final Rare Class Accuracy (WITHOUT): {final_acc_b_no:.1%}")
+    print(f"  Final Rare Class Accuracy (WITH):    {final_acc_b_with:.1%}")
+    print(f"  Verdict: {'SUPPORTED' if final_acc_b_with >= final_acc_b_no else 'NOT SUPPORTED'}")
+    
+    return {
+        'metrics_no_clip': metrics_no_clip,
+        'metrics_with_clip': metrics_with_clip,
+        'rare_indices': rare_indices,
+    }
+
+
+if __name__ == "__main__":
+    results = main()
+    print("\n" + "="*70)
+    print("EXPERIMENT COMPLETE!")
+    print("="*70)