Buckets:
tahamajs/Sysmem2_in_AI / ComputerAssignments /CA6_systematic_generalization /src /models /grokking_analyzer.py
| """ | |
| Grokking Phenomenon Analyzer | |
| Implementation of grokking detection and analysis for systematic generalization. | |
| Grokking refers to delayed generalization where models suddenly transition from | |
| memorization to true understanding after extended training. | |
| Based on: "Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets" | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| from typing import List, Dict, Tuple, Any, Optional | |
| from dataclasses import dataclass, field | |
| from collections import deque | |
| class GrokkingMetrics: | |
| """Stores metrics for grokking analysis.""" | |
| epochs: List[int] = field(default_factory=list) | |
| train_losses: List[float] = field(default_factory=list) | |
| train_accuracies: List[float] = field(default_factory=list) | |
| test_losses: List[float] = field(default_factory=list) | |
| test_accuracies: List[float] = field(default_factory=list) | |
| generalization_gaps: List[float] = field(default_factory=list) | |
| weight_norms: List[float] = field(default_factory=list) | |
| gradient_norms: List[float] = field(default_factory=list) | |
| def add_metrics( | |
| self, | |
| epoch: int, | |
| train_loss: float, | |
| train_acc: float, | |
| test_loss: float, | |
| test_acc: float, | |
| weight_norm: Optional[float] = None, | |
| gradient_norm: Optional[float] = None, | |
| ): | |
| """Add metrics for a single epoch.""" | |
| self.epochs.append(epoch) | |
| self.train_losses.append(train_loss) | |
| self.train_accuracies.append(train_acc) | |
| self.test_losses.append(test_loss) | |
| self.test_accuracies.append(test_acc) | |
| self.generalization_gaps.append(train_acc - test_acc) | |
| if weight_norm is not None: | |
| self.weight_norms.append(weight_norm) | |
| if gradient_norm is not None: | |
| self.gradient_norms.append(gradient_norm) | |
| class GrokkingDetector: | |
| """ | |
| Detects grokking phenomenon in training dynamics. | |
| Grokking characteristics: | |
| 1. Initial overfitting (high train acc, low test acc) | |
| 2. Sudden improvement in test accuracy | |
| 3. Convergence to high generalization | |
| """ | |
| def __init__( | |
| self, | |
| window_size: int = 20, | |
| improvement_threshold: float = 0.3, | |
| accuracy_threshold: float = 0.8, | |
| ): | |
| self.window_size = window_size | |
| self.improvement_threshold = improvement_threshold | |
| self.accuracy_threshold = accuracy_threshold | |
| self.grokking_detected = False | |
| self.grokking_epoch = None | |
| def detect(self, metrics: GrokkingMetrics) -> Dict[str, Any]: | |
| """ | |
| Detect if grokking has occurred. | |
| Args: | |
| metrics: Training metrics | |
| Returns: | |
| Dictionary with detection results | |
| """ | |
| if len(metrics.epochs) < self.window_size * 2: | |
| return { | |
| "grokking_detected": False, | |
| "reason": "Insufficient data", | |
| "epochs_needed": self.window_size * 2 - len(metrics.epochs), | |
| } | |
| test_accs = metrics.test_accuracies | |
| # Check for sudden improvement | |
| for i in range(self.window_size, len(test_accs) - self.window_size): | |
| before_window = test_accs[i - self.window_size : i] | |
| after_window = test_accs[i : i + self.window_size] | |
| before_avg = np.mean(before_window) | |
| after_avg = np.mean(after_window) | |
| improvement = after_avg - before_avg | |
| # Check grokking criteria | |
| if ( | |
| improvement > self.improvement_threshold | |
| and after_avg > self.accuracy_threshold | |
| and before_avg < self.accuracy_threshold | |
| ): | |
| self.grokking_detected = True | |
| self.grokking_epoch = metrics.epochs[i] | |
| return { | |
| "grokking_detected": True, | |
| "grokking_epoch": self.grokking_epoch, | |
| "improvement": improvement, | |
| "before_accuracy": before_avg, | |
| "after_accuracy": after_avg, | |
| "generalization_gap_before": np.mean( | |
| metrics.generalization_gaps[i - self.window_size : i] | |
| ), | |
| "generalization_gap_after": np.mean( | |
| metrics.generalization_gaps[i : i + self.window_size] | |
| ), | |
| } | |
| return { | |
| "grokking_detected": False, | |
| "reason": "No sudden improvement detected", | |
| "max_test_accuracy": max(test_accs), | |
| "final_test_accuracy": test_accs[-1], | |
| } | |
| def analyze_phases(self, metrics: GrokkingMetrics) -> Dict[str, Any]: | |
| """ | |
| Analyze different phases of training. | |
| Phases: | |
| 1. Memorization: High train acc, low test acc | |
| 2. Transition: Improving test acc | |
| 3. Generalization: High train and test acc | |
| """ | |
| if len(metrics.epochs) < 30: | |
| return {"error": "Insufficient data for phase analysis"} | |
| train_accs = np.array(metrics.train_accuracies) | |
| test_accs = np.array(metrics.test_accuracies) | |
| gaps = np.array(metrics.generalization_gaps) | |
| # Identify phases | |
| memorization_phase = [] | |
| transition_phase = [] | |
| generalization_phase = [] | |
| for i, (train_acc, test_acc, gap) in enumerate( | |
| zip(train_accs, test_accs, gaps) | |
| ): | |
| if train_acc > 0.9 and test_acc < 0.5: | |
| memorization_phase.append(i) | |
| elif gap > 0.1 and test_acc > 0.5: | |
| transition_phase.append(i) | |
| elif train_acc > 0.9 and test_acc > 0.8: | |
| generalization_phase.append(i) | |
| return { | |
| "memorization_phase": { | |
| "epochs": [metrics.epochs[i] for i in memorization_phase], | |
| "duration": len(memorization_phase), | |
| "avg_gap": ( | |
| np.mean([gaps[i] for i in memorization_phase]) | |
| if memorization_phase | |
| else 0 | |
| ), | |
| }, | |
| "transition_phase": { | |
| "epochs": [metrics.epochs[i] for i in transition_phase], | |
| "duration": len(transition_phase), | |
| "avg_gap": ( | |
| np.mean([gaps[i] for i in transition_phase]) | |
| if transition_phase | |
| else 0 | |
| ), | |
| }, | |
| "generalization_phase": { | |
| "epochs": [metrics.epochs[i] for i in generalization_phase], | |
| "duration": len(generalization_phase), | |
| "avg_gap": ( | |
| np.mean([gaps[i] for i in generalization_phase]) | |
| if generalization_phase | |
| else 0 | |
| ), | |
| }, | |
| } | |
| class GrokkingMonitor: | |
| """ | |
| Real-time monitoring of grokking during training. | |
| """ | |
| def __init__(self, patience: int = 100, check_interval: int = 10): | |
| self.patience = patience | |
| self.check_interval = check_interval | |
| self.metrics = GrokkingMetrics() | |
| self.detector = GrokkingDetector() | |
| self.best_test_acc = 0.0 | |
| self.epochs_without_improvement = 0 | |
| def update( | |
| self, | |
| epoch: int, | |
| train_loss: float, | |
| train_acc: float, | |
| test_loss: float, | |
| test_acc: float, | |
| model: Optional[nn.Module] = None, | |
| ) -> Dict[str, Any]: | |
| """ | |
| Update monitoring with new metrics. | |
| Args: | |
| epoch: Current epoch | |
| train_loss: Training loss | |
| train_acc: Training accuracy | |
| test_loss: Test loss | |
| test_acc: Test accuracy | |
| model: Optional model for weight/gradient norms | |
| Returns: | |
| Status dictionary | |
| """ | |
| # Compute norms if model provided | |
| weight_norm = None | |
| gradient_norm = None | |
| if model is not None: | |
| weight_norm = sum(p.norm().item() for p in model.parameters()) | |
| gradient_norm = sum( | |
| p.grad.norm().item() for p in model.parameters() if p.grad is not None | |
| ) | |
| # Add metrics | |
| self.metrics.add_metrics( | |
| epoch, | |
| train_loss, | |
| train_acc, | |
| test_loss, | |
| test_acc, | |
| weight_norm, | |
| gradient_norm, | |
| ) | |
| # Check for improvement | |
| if test_acc > self.best_test_acc: | |
| self.best_test_acc = test_acc | |
| self.epochs_without_improvement = 0 | |
| else: | |
| self.epochs_without_improvement += 1 | |
| # Periodic grokking detection | |
| status = {"epoch": epoch} | |
| if epoch % self.check_interval == 0: | |
| detection_result = self.detector.detect(self.metrics) | |
| status.update(detection_result) | |
| if detection_result.get("grokking_detected"): | |
| status["message"] = ( | |
| f"🎉 Grokking detected at epoch {detection_result['grokking_epoch']}!" | |
| ) | |
| # Early stopping check | |
| if self.epochs_without_improvement >= self.patience: | |
| status["should_stop"] = True | |
| status["reason"] = "No improvement in test accuracy" | |
| else: | |
| status["should_stop"] = False | |
| return status | |
| def plot_training_dynamics(self, save_path: Optional[str] = None): | |
| """Plot training dynamics showing grokking.""" | |
| fig, axes = plt.subplots(2, 2, figsize=(15, 10)) | |
| fig.suptitle( | |
| "Grokking Analysis: Training Dynamics", fontsize=16, fontweight="bold" | |
| ) | |
| epochs = self.metrics.epochs | |
| # Plot 1: Accuracies | |
| axes[0, 0].plot( | |
| epochs, self.metrics.train_accuracies, label="Train Accuracy", linewidth=2 | |
| ) | |
| axes[0, 0].plot( | |
| epochs, self.metrics.test_accuracies, label="Test Accuracy", linewidth=2 | |
| ) | |
| axes[0, 0].set_xlabel("Epoch") | |
| axes[0, 0].set_ylabel("Accuracy") | |
| axes[0, 0].set_title("Training vs Test Accuracy") | |
| axes[0, 0].legend() | |
| axes[0, 0].grid(True, alpha=0.3) | |
| # Mark grokking point if detected | |
| if self.detector.grokking_detected: | |
| grok_idx = epochs.index(self.detector.grokking_epoch) | |
| axes[0, 0].axvline( | |
| x=self.detector.grokking_epoch, | |
| color="red", | |
| linestyle="--", | |
| label="Grokking Point", | |
| ) | |
| axes[0, 0].scatter( | |
| [self.detector.grokking_epoch], | |
| [self.metrics.test_accuracies[grok_idx]], | |
| color="red", | |
| s=100, | |
| zorder=5, | |
| ) | |
| # Plot 2: Losses | |
| axes[0, 1].plot( | |
| epochs, self.metrics.train_losses, label="Train Loss", linewidth=2 | |
| ) | |
| axes[0, 1].plot( | |
| epochs, self.metrics.test_losses, label="Test Loss", linewidth=2 | |
| ) | |
| axes[0, 1].set_xlabel("Epoch") | |
| axes[0, 1].set_ylabel("Loss") | |
| axes[0, 1].set_title("Training vs Test Loss") | |
| axes[0, 1].legend() | |
| axes[0, 1].grid(True, alpha=0.3) | |
| axes[0, 1].set_yscale("log") | |
| # Plot 3: Generalization Gap | |
| axes[1, 0].plot( | |
| epochs, self.metrics.generalization_gaps, linewidth=2, color="purple" | |
| ) | |
| axes[1, 0].axhline(y=0, color="black", linestyle="--", alpha=0.5) | |
| axes[1, 0].set_xlabel("Epoch") | |
| axes[1, 0].set_ylabel("Generalization Gap") | |
| axes[1, 0].set_title("Generalization Gap (Train - Test Accuracy)") | |
| axes[1, 0].grid(True, alpha=0.3) | |
| # Plot 4: Weight Norms | |
| if self.metrics.weight_norms: | |
| axes[1, 1].plot( | |
| epochs[: len(self.metrics.weight_norms)], | |
| self.metrics.weight_norms, | |
| linewidth=2, | |
| color="green", | |
| ) | |
| axes[1, 1].set_xlabel("Epoch") | |
| axes[1, 1].set_ylabel("Weight Norm") | |
| axes[1, 1].set_title("Model Weight Norm Evolution") | |
| axes[1, 1].grid(True, alpha=0.3) | |
| else: | |
| axes[1, 1].text( | |
| 0.5, | |
| 0.5, | |
| "Weight norms not tracked", | |
| ha="center", | |
| va="center", | |
| transform=axes[1, 1].transAxes, | |
| ) | |
| plt.tight_layout() | |
| if save_path: | |
| plt.savefig(save_path, dpi=300, bbox_inches="tight") | |
| print(f"Plot saved to {save_path}") | |
| return fig | |
| class GrokkingExperiment: | |
| """ | |
| Complete experiment setup for studying grokking. | |
| """ | |
| def __init__( | |
| self, | |
| model: nn.Module, | |
| train_loader: Any, | |
| test_loader: Any, | |
| optimizer: torch.optim.Optimizer, | |
| criterion: nn.Module, | |
| device: str = "cpu", | |
| ): | |
| self.model = model | |
| self.train_loader = train_loader | |
| self.test_loader = test_loader | |
| self.optimizer = optimizer | |
| self.criterion = criterion | |
| self.device = device | |
| self.monitor = GrokkingMonitor(patience=200, check_interval=10) | |
| def train_epoch(self) -> Tuple[float, float]: | |
| """Train for one epoch.""" | |
| self.model.train() | |
| total_loss = 0.0 | |
| correct = 0 | |
| total = 0 | |
| for inputs, targets in self.train_loader: | |
| inputs, targets = inputs.to(self.device), targets.to(self.device) | |
| self.optimizer.zero_grad() | |
| outputs = self.model(inputs) | |
| loss = self.criterion(outputs, targets) | |
| loss.backward() | |
| self.optimizer.step() | |
| total_loss += loss.item() | |
| _, predicted = outputs.max(1) | |
| total += targets.size(0) | |
| correct += predicted.eq(targets).sum().item() | |
| avg_loss = total_loss / len(self.train_loader) | |
| accuracy = correct / total | |
| return avg_loss, accuracy | |
| def evaluate(self) -> Tuple[float, float]: | |
| """Evaluate on test set.""" | |
| self.model.eval() | |
| total_loss = 0.0 | |
| correct = 0 | |
| total = 0 | |
| with torch.no_grad(): | |
| for inputs, targets in self.test_loader: | |
| inputs, targets = inputs.to(self.device), targets.to(self.device) | |
| outputs = self.model(inputs) | |
| loss = self.criterion(outputs, targets) | |
| total_loss += loss.item() | |
| _, predicted = outputs.max(1) | |
| total += targets.size(0) | |
| correct += predicted.eq(targets).sum().item() | |
| avg_loss = total_loss / len(self.test_loader) | |
| accuracy = correct / total | |
| return avg_loss, accuracy | |
| def run(self, num_epochs: int = 1000, verbose: bool = True) -> GrokkingMetrics: | |
| """ | |
| Run complete grokking experiment. | |
| Args: | |
| num_epochs: Number of epochs to train | |
| verbose: Whether to print progress | |
| Returns: | |
| Complete metrics | |
| """ | |
| for epoch in range(num_epochs): | |
| train_loss, train_acc = self.train_epoch() | |
| test_loss, test_acc = self.evaluate() | |
| status = self.monitor.update( | |
| epoch, train_loss, train_acc, test_loss, test_acc, self.model | |
| ) | |
| if verbose and epoch % 10 == 0: | |
| print( | |
| f"Epoch {epoch:4d} | " | |
| f"Train: {train_acc:.3f} | " | |
| f"Test: {test_acc:.3f} | " | |
| f"Gap: {train_acc - test_acc:.3f}" | |
| ) | |
| if status.get("grokking_detected"): | |
| print(f" {status['message']}") | |
| if status.get("should_stop"): | |
| print(f"Early stopping at epoch {epoch}: {status['reason']}") | |
| break | |
| # Final analysis | |
| detection_result = self.monitor.detector.detect(self.monitor.metrics) | |
| phase_analysis = self.monitor.detector.analyze_phases(self.monitor.metrics) | |
| if verbose: | |
| print("\n" + "=" * 60) | |
| print("Grokking Analysis Results:") | |
| print("=" * 60) | |
| print(f"Grokking Detected: {detection_result['grokking_detected']}") | |
| if detection_result["grokking_detected"]: | |
| print(f" Grokking Epoch: {detection_result['grokking_epoch']}") | |
| print(f" Improvement: {detection_result['improvement']:.3f}") | |
| print(f" Before Accuracy: {detection_result['before_accuracy']:.3f}") | |
| print(f" After Accuracy: {detection_result['after_accuracy']:.3f}") | |
| print("\nPhase Analysis:") | |
| for phase_name, phase_info in phase_analysis.items(): | |
| if "error" not in phase_info: | |
| print(f" {phase_name}: {phase_info['duration']} epochs") | |
| return self.monitor.metrics | |
| # Example usage | |
| if __name__ == "__main__": | |
| print("Grokking Phenomenon Analyzer") | |
| print("=" * 60) | |
| print("This module provides tools for detecting and analyzing grokking") | |
| print("in neural network training for systematic generalization tasks.") | |
| print("\n✓ Module loaded successfully!") | |
Xet Storage Details
- Size:
- 17.3 kB
- Xet hash:
- 1a76cee45e2826e65d05bafda783ff4abe34d0d430ca171868fafc5121b5e03e
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.