Upload final_experiment.py with huggingface_hub

final_experiment.py

@@ -0,0 +1,513 @@
+"""
+Final Gradient Clipping Experiment: Testing Physics-of-AI Predictions
+
+Key insights from previous experiments:
+1. With extreme imbalance (99:1), neither model learns rare class
+2. Gradient clipping's benefit is in STABILITY, not learning rare classes per se
+3. The key effect is on WEIGHT NORM STABILITY and GRADIENT SPIKE HANDLING
+
+This experiment tests:
+1. Prediction 2: Representation Collapse - effective dim variance without clipping
+2. Prediction 4: Rare Sample Learning - using moderate imbalance (80:20)
+3. NEW: Weight norm stability analysis
+4. NEW: Gradient spike analysis at rare sample positions
+"""
+
+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
+
+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."""
+    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]:
+    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 create_dataset_moderate_imbalance(n_samples=1000, rare_ratio=0.2, seed=SEED):
+    """Create dataset with moderate imbalance (e.g., 80:20)."""
+    set_seeds(seed)
+    
+    n_rare = int(n_samples * rare_ratio)
+    n_common = n_samples - n_rare
+    
+    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 create_dataset_extreme_imbalance(n_samples=1000, n_rare=10, seed=SEED):
+    """Create dataset with extreme imbalance (99:1)."""
+    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.1, init_weights=None,
+                        track_every: int = 50) -> Dict:
+    """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: []},
+        'accuracy_steps': [],
+        'weight_norm_changes': [],  # Track sudden changes
+    }
+    
+    step = 0
+    n_samples = len(inputs)
+    prev_weight_norm = None
+    
+    for epoch in range(n_epochs):
+        model.train()
+        
+        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())
+            
+            current_weight_norm = sum(p.data.norm(2).item() ** 2 for p in model.parameters()) ** 0.5
+            metrics['weight_norms'].append(current_weight_norm)
+            
+            # Track weight norm change
+            if prev_weight_norm is not None:
+                metrics['weight_norm_changes'].append(abs(current_weight_norm - prev_weight_norm))
+            else:
+                metrics['weight_norm_changes'].append(0)
+            prev_weight_norm = current_weight_norm
+            
+            # Track 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)
+                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)
+                metrics['accuracy_steps'].append(step)
+            
+            step += 1
+    
+    return metrics
+
+
+def run_experiment_suite():
+    """Run complete experiment suite with both imbalance levels."""
+    print("="*70)
+    print("FINAL GRADIENT CLIPPING EXPERIMENT")
+    print("Testing Physics-of-AI Predictions")
+    print("="*70)
+    
+    # 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()}
+    
+    results = {}
+    
+    # =========================================================================
+    # EXPERIMENT 1: Extreme Imbalance (99:1) - Original Setup
+    # =========================================================================
+    print("\n" + "="*70)
+    print("EXPERIMENT 1: EXTREME IMBALANCE (99:1)")
+    print("="*70)
+    
+    inputs_extreme, targets_extreme, rare_extreme = create_dataset_extreme_imbalance(
+        n_samples=1000, n_rare=10, seed=SEED
+    )
+    print(f"Dataset: {(targets_extreme == 0).sum().item()} common, {(targets_extreme == 1).sum().item()} rare")
+    
+    print("\nTraining WITHOUT clipping...")
+    metrics_extreme_no_clip = train_with_tracking(
+        inputs_extreme, targets_extreme, rare_extreme,
+        clip_grad=False, n_epochs=5, lr=0.1, 
+        init_weights=init_weights, track_every=100
+    )
+    
+    print("Training WITH clipping...")
+    metrics_extreme_with_clip = train_with_tracking(
+        inputs_extreme, targets_extreme, rare_extreme,
+        clip_grad=True, max_norm=1.0, n_epochs=5, lr=0.1,
+        init_weights=init_weights, track_every=100
+    )
+    
+    results['extreme'] = {
+        'no_clip': metrics_extreme_no_clip,
+        'with_clip': metrics_extreme_with_clip,
+        'rare_indices': rare_extreme
+    }
+    
+    # =========================================================================
+    # EXPERIMENT 2: Moderate Imbalance (80:20)
+    # =========================================================================
+    print("\n" + "="*70)
+    print("EXPERIMENT 2: MODERATE IMBALANCE (80:20)")
+    print("="*70)
+    
+    inputs_moderate, targets_moderate, rare_moderate = create_dataset_moderate_imbalance(
+        n_samples=1000, rare_ratio=0.2, seed=SEED
+    )
+    print(f"Dataset: {(targets_moderate == 0).sum().item()} common, {(targets_moderate == 1).sum().item()} rare")
+    
+    print("\nTraining WITHOUT clipping...")
+    metrics_moderate_no_clip = train_with_tracking(
+        inputs_moderate, targets_moderate, rare_moderate,
+        clip_grad=False, n_epochs=10, lr=0.1, 
+        init_weights=init_weights, track_every=100
+    )
+    
+    print("Training WITH clipping...")
+    metrics_moderate_with_clip = train_with_tracking(
+        inputs_moderate, targets_moderate, rare_moderate,
+        clip_grad=True, max_norm=1.0, n_epochs=10, lr=0.1,
+        init_weights=init_weights, track_every=100
+    )
+    
+    results['moderate'] = {
+        'no_clip': metrics_moderate_no_clip,
+        'with_clip': metrics_moderate_with_clip,
+        'rare_indices': rare_moderate
+    }
+    
+    return results
+
+
+def plot_final_comparison(results: Dict, filename: str):
+    """Create final comparison plot."""
+    fig = plt.figure(figsize=(20, 20))
+    gs = fig.add_gridspec(5, 2, hspace=0.35, wspace=0.25)
+    
+    # =========================================================================
+    # Row 1: Weight Norm Stability (Key Physics-of-AI Insight)
+    # =========================================================================
+    ax1 = fig.add_subplot(gs[0, 0])
+    ax2 = fig.add_subplot(gs[0, 1])
+    
+    # Extreme imbalance
+    steps = range(len(results['extreme']['no_clip']['weight_norms']))
+    ax1.plot(steps, results['extreme']['no_clip']['weight_norms'], 'r-', alpha=0.7, linewidth=1, label='Without Clip')
+    ax1.plot(steps, results['extreme']['with_clip']['weight_norms'], 'g-', alpha=0.7, linewidth=1, label='With Clip')
+    ax1.set_ylabel('Weight Norm', fontsize=11)
+    ax1.set_title('EXTREME (99:1) - Weight Norm Evolution', fontsize=12, fontweight='bold')
+    ax1.legend()
+    ax1.grid(True, alpha=0.3)
+    
+    # Moderate imbalance
+    steps = range(len(results['moderate']['no_clip']['weight_norms']))
+    ax2.plot(steps, results['moderate']['no_clip']['weight_norms'], 'r-', alpha=0.7, linewidth=1, label='Without Clip')
+    ax2.plot(steps, results['moderate']['with_clip']['weight_norms'], 'g-', alpha=0.7, linewidth=1, label='With Clip')
+    ax2.set_title('MODERATE (80:20) - Weight Norm Evolution', fontsize=12, fontweight='bold')
+    ax2.legend()
+    ax2.grid(True, alpha=0.3)
+    
+    # =========================================================================
+    # Row 2: Weight Norm Changes (Stability Metric)
+    # =========================================================================
+    ax3 = fig.add_subplot(gs[1, 0])
+    ax4 = fig.add_subplot(gs[1, 1])
+    
+    # Extreme
+    steps = range(len(results['extreme']['no_clip']['weight_norm_changes']))
+    ax3.plot(steps, results['extreme']['no_clip']['weight_norm_changes'], 'r-', alpha=0.5, linewidth=0.5, label='Without Clip')
+    ax3.plot(steps, results['extreme']['with_clip']['weight_norm_changes'], 'g-', alpha=0.5, linewidth=0.5, label='With Clip')
+    ax3.set_ylabel('|Weight Norm Change|', fontsize=11)
+    ax3.set_title('EXTREME - Weight Norm Changes (Stability)', fontsize=12, fontweight='bold')
+    ax3.legend()
+    ax3.grid(True, alpha=0.3)
+    
+    # Moderate
+    steps = range(len(results['moderate']['no_clip']['weight_norm_changes']))
+    ax4.plot(steps, results['moderate']['no_clip']['weight_norm_changes'], 'r-', alpha=0.5, linewidth=0.5, label='Without Clip')
+    ax4.plot(steps, results['moderate']['with_clip']['weight_norm_changes'], 'g-', alpha=0.5, linewidth=0.5, label='With Clip')
+    ax4.set_title('MODERATE - Weight Norm Changes (Stability)', fontsize=12, fontweight='bold')
+    ax4.legend()
+    ax4.grid(True, alpha=0.3)
+    
+    # =========================================================================
+    # Row 3: Gradient Norms
+    # =========================================================================
+    ax5 = fig.add_subplot(gs[2, 0])
+    ax6 = fig.add_subplot(gs[2, 1])
+    
+    # Extreme
+    steps = range(len(results['extreme']['no_clip']['grad_norms']))
+    ax5.plot(steps, results['extreme']['no_clip']['grad_norms'], 'r-', alpha=0.3, linewidth=0.5, label='Without Clip')
+    ax5.plot(steps, results['extreme']['with_clip']['grad_norms'], 'g-', alpha=0.3, linewidth=0.5, label='With Clip')
+    ax5.axhline(y=1.0, color='black', linestyle='--', linewidth=2, label='Clip threshold')
+    ax5.set_ylabel('Gradient Norm', fontsize=11)
+    ax5.set_title('EXTREME - Gradient Norms', fontsize=12, fontweight='bold')
+    ax5.legend()
+    ax5.grid(True, alpha=0.3)
+    
+    # Moderate
+    steps = range(len(results['moderate']['no_clip']['grad_norms']))
+    ax6.plot(steps, results['moderate']['no_clip']['grad_norms'], 'r-', alpha=0.3, linewidth=0.5, label='Without Clip')
+    ax6.plot(steps, results['moderate']['with_clip']['grad_norms'], 'g-', alpha=0.3, linewidth=0.5, label='With Clip')
+    ax6.axhline(y=1.0, color='black', linestyle='--', linewidth=2, label='Clip threshold')
+    ax6.set_title('MODERATE - Gradient Norms', fontsize=12, fontweight='bold')
+    ax6.legend()
+    ax6.grid(True, alpha=0.3)
+    
+    # =========================================================================
+    # Row 4: Effective Dimension
+    # =========================================================================
+    ax7 = fig.add_subplot(gs[3, 0])
+    ax8 = fig.add_subplot(gs[3, 1])
+    
+    # Extreme
+    ax7.plot(results['extreme']['no_clip']['effective_dim_steps'], 
+             results['extreme']['no_clip']['effective_dims'], 
+             'r-o', alpha=0.7, linewidth=2, markersize=4, label='Without Clip')
+    ax7.plot(results['extreme']['with_clip']['effective_dim_steps'], 
+             results['extreme']['with_clip']['effective_dims'], 
+             'g-o', alpha=0.7, linewidth=2, markersize=4, label='With Clip')
+    ax7.set_ylabel('Effective Dimension', fontsize=11)
+    ax7.set_title('EXTREME - Effective Dimensionality', fontsize=12, fontweight='bold')
+    ax7.legend()
+    ax7.grid(True, alpha=0.3)
+    
+    # Moderate
+    ax8.plot(results['moderate']['no_clip']['effective_dim_steps'], 
+             results['moderate']['no_clip']['effective_dims'], 
+             'r-o', alpha=0.7, linewidth=2, markersize=4, label='Without Clip')
+    ax8.plot(results['moderate']['with_clip']['effective_dim_steps'], 
+             results['moderate']['with_clip']['effective_dims'], 
+             'g-o', alpha=0.7, linewidth=2, markersize=4, label='With Clip')
+    ax8.set_title('MODERATE - Effective Dimensionality', fontsize=12, fontweight='bold')
+    ax8.legend()
+    ax8.grid(True, alpha=0.3)
+    
+    # =========================================================================
+    # Row 5: Class Accuracies
+    # =========================================================================
+    ax9 = fig.add_subplot(gs[4, 0])
+    ax10 = fig.add_subplot(gs[4, 1])
+    
+    # Extreme - Rare class B
+    ax9.plot(results['extreme']['no_clip']['accuracy_steps'], 
+             results['extreme']['no_clip']['class_accuracies'][1], 
+             'r-', alpha=0.7, linewidth=2, label='Without Clip')
+    ax9.plot(results['extreme']['with_clip']['accuracy_steps'], 
+             results['extreme']['with_clip']['class_accuracies'][1], 
+             'g-', alpha=0.7, linewidth=2, label='With Clip')
+    ax9.set_ylabel('Rare Class B Accuracy', fontsize=11)
+    ax9.set_xlabel('Training Step', fontsize=11)
+    ax9.set_title('EXTREME - Rare Class Accuracy', fontsize=12, fontweight='bold')
+    ax9.legend()
+    ax9.grid(True, alpha=0.3)
+    ax9.set_ylim([0, 1.05])
+    
+    # Moderate - Rare class B
+    ax10.plot(results['moderate']['no_clip']['accuracy_steps'], 
+              results['moderate']['no_clip']['class_accuracies'][1], 
+              'r-', alpha=0.7, linewidth=2, label='Without Clip')
+    ax10.plot(results['moderate']['with_clip']['accuracy_steps'], 
+              results['moderate']['with_clip']['class_accuracies'][1], 
+              'g-', alpha=0.7, linewidth=2, label='With Clip')
+    ax10.set_xlabel('Training Step', fontsize=11)
+    ax10.set_title('MODERATE - Rare Class Accuracy', fontsize=12, fontweight='bold')
+    ax10.legend()
+    ax10.grid(True, alpha=0.3)
+    ax10.set_ylim([0, 1.05])
+    
+    fig.suptitle('Gradient Clipping Analysis: Physics-of-AI Predictions\n'
+                 'Comparing Extreme (99:1) vs Moderate (80:20) Class Imbalance',
+                 fontsize=14, fontweight='bold', y=1.01)
+    
+    plt.savefig(filename, dpi=150, bbox_inches='tight')
+    plt.close()
+    print(f"Final comparison plot saved to: {filename}")
+
+
+def compute_statistics(results: Dict) -> Dict:
+    """Compute summary statistics for all experiments."""
+    stats = {}
+    
+    for imbalance in ['extreme', 'moderate']:
+        no_clip = results[imbalance]['no_clip']
+        with_clip = results[imbalance]['with_clip']
+        
+        stats[imbalance] = {
+            'weight_norm_std': {
+                'no_clip': np.std(no_clip['weight_norms']),
+                'with_clip': np.std(with_clip['weight_norms']),
+            },
+            'weight_change_mean': {
+                'no_clip': np.mean(no_clip['weight_norm_changes']),
+                'with_clip': np.mean(with_clip['weight_norm_changes']),
+            },
+            'weight_change_max': {
+                'no_clip': np.max(no_clip['weight_norm_changes']),
+                'with_clip': np.max(with_clip['weight_norm_changes']),
+            },
+            'grad_norm_max': {
+                'no_clip': np.max(no_clip['grad_norms']),
+                'with_clip': np.max(with_clip['grad_norms']),
+            },
+            'effective_dim_std': {
+                'no_clip': np.std(no_clip['effective_dims']),
+                'with_clip': np.std(with_clip['effective_dims']),
+            },
+            'final_rare_acc': {
+                'no_clip': no_clip['class_accuracies'][1][-1] if no_clip['class_accuracies'][1] else 0,
+                'with_clip': with_clip['class_accuracies'][1][-1] if with_clip['class_accuracies'][1] else 0,
+            },
+        }
+    
+    return stats
+
+
+def print_summary(stats: Dict):
+    """Print formatted summary."""
+    print("\n" + "="*70)
+    print("EXPERIMENT SUMMARY")
+    print("="*70)
+    
+    for imbalance in ['extreme', 'moderate']:
+        s = stats[imbalance]
+        label = "EXTREME (99:1)" if imbalance == 'extreme' else "MODERATE (80:20)"
+        
+        print(f"\n{label}")
+        print("-" * 50)
+        
+        print(f"\n[PREDICTION 2] Representation Collapse (Effective Dim Variance):")
+        print(f"  WITHOUT Clipping: {s['effective_dim_std']['no_clip']:.6f}")
+        print(f"  WITH Clipping:    {s['effective_dim_std']['with_clip']:.6f}")
+        supported = s['effective_dim_std']['no_clip'] > s['effective_dim_std']['with_clip']
+        print(f"  Verdict: {'SUPPORTED' if supported else 'NOT SUPPORTED'}")
+        
+        print(f"\n[PREDICTION 4] Rare Sample Learning:")
+        print(f"  Final Rare Accuracy (WITHOUT): {s['final_rare_acc']['no_clip']:.1%}")
+        print(f"  Final Rare Accuracy (WITH):    {s['final_rare_acc']['with_clip']:.1%}")
+        supported = s['final_rare_acc']['with_clip'] >= s['final_rare_acc']['no_clip']
+        print(f"  Verdict: {'SUPPORTED' if supported else 'NOT SUPPORTED'}")
+        
+        print(f"\n[STABILITY] Weight Norm Analysis:")
+        print(f"  Weight Norm Std (WITHOUT): {s['weight_norm_std']['no_clip']:.4f}")
+        print(f"  Weight Norm Std (WITH):    {s['weight_norm_std']['with_clip']:.4f}")
+        print(f"  Max Weight Change (WITHOUT): {s['weight_change_max']['no_clip']:.4f}")
+        print(f"  Max Weight Change (WITH):    {s['weight_change_max']['with_clip']:.4f}")
+        
+        print(f"\n[GRADIENT] Analysis:")
+        print(f"  Max Gradient Norm (WITHOUT): {s['grad_norm_max']['no_clip']:.4f}")
+        print(f"  Max Gradient Norm (WITH):    {s['grad_norm_max']['with_clip']:.4f}")
+        print(f"  Clipping Ratio: {s['grad_norm_max']['no_clip'] / 1.0:.1f}x threshold")
+
+
+def main():
+    # Run experiments
+    results = run_experiment_suite()
+    
+    # Generate plots
+    print("\n" + "="*70)
+    print("GENERATING PLOTS")
+    print("="*70)
+    
+    plot_final_comparison(results, "final_comparison.png")
+    
+    # Compute and print statistics
+    stats = compute_statistics(results)
+    print_summary(stats)
+    
+    return results, stats
+
+
+if __name__ == "__main__":
+    results, stats = main()
+    print("\n" + "="*70)
+    print("EXPERIMENT COMPLETE!")
+    print("="*70)