This document describes the implementation of the Neural Plasticity Dashboard, a comprehensive visualization system for tracking and analyzing neural plasticity experiments in real-time.
Version: v0.0.1 (2025-04-20)
-
Real-time Experiment Monitoring:
- Track experiment progress while it's running
- View key metrics, phases, and decisions in real-time
- Monitor console output with formatted information
-
Comprehensive Visualizations:
- Warmup phase: Training loss curves and stabilization detection
- Analysis phase: Entropy and gradient heatmaps for attention heads
- Pruning phase: Before/after architecture visualization
- Evaluation phase: Direct comparison of baseline vs. pruned models
-
Multiple Dashboard Options:
- Weights & Biases integration for online tracking
- Standalone HTML dashboard for offline viewing
- Real-time console output for immediate feedback
-
Collaboration Features:
- Shareable dashboard links for team analysis
- Browser integration for automatic dashboard opening
- Support for both local and Colab environments
The WandbDashboard class in utils/neural_plasticity/dashboard/wandb_integration.py provides:
- Integration with Weights & Biases for experiment tracking
- Support for both online and offline modes
- Comprehensive logging of metrics, visualizations, and text samples
- Real-time console output for all experiment phases
- Loss curve visualization with smoothed trend line
- Training metric tracking (loss, perplexity)
- Text samples during training
- Entropy heatmap for attention heads
- Gradient heatmap for attention heads
- Combined importance metrics visualization
- Pruning decision visualization
- Before/after model architecture comparison
- Pruned heads highlighted by layer
- Side-by-side text generation comparison
- Perplexity comparison bar charts
- Attention pattern visualization
The dashboard_demo.py script provides:
- Complete mock experiment simulation for testing
- Generation of all visualizations without running a full experiment
- Exportable HTML dashboard with embedded visualizations
- Browser integration for easy viewing
The colab_integration.py module provides:
- Specialized support for Google Colab environments
- Integration with Colab's display system
- Sharing options for collaborative analysis
- Support for tunneling for local servers
# Activate virtual environment
source .venv/bin/activate
# Run offline demo (creates local dashboard)
python -m utils.neural_plasticity.dashboard.dashboard_demo
# Run online demo (publishes to wandb.ai)
WANDB_MODE=online python -m utils.neural_plasticity.dashboard.dashboard_demo --online# Run experiment with dashboard enabled
python scripts/run_neural_plasticity.py --model_name distilgpt2 --pruning_strategy entropy --pruning_level 0.2 --use_dashboard
# For quick testing
python scripts/run_neural_plasticity.py --model_name distilgpt2 --quick_test --use_dashboard-
Online Mode: Dashboard URL will be provided in console output, or:
open https://wandb.ai/[username]/neural-plasticity/runs/[run_id] -
Offline Mode: Open the local HTML file:
open output/neural_plasticity_[timestamp]/dashboard.html -
Colab Environment: Click the dashboard link displayed in the notebook cell.
The dashboard tracks key metrics:
- Training and evaluation loss
- Perplexity
- Sparsity (percentage of pruned heads)
- Inference speedup
- Model size reduction
- Heatmaps: For entropy, gradients, and attention patterns
- Line Charts: For loss curves and performance tracking
- Bar Charts: For perplexity and other comparative metrics
- Grid Visualizations: For model architecture and pruning decisions
- Text Samples: For qualitative evaluation of model outputs
The dashboard provides callback functions for the experiment:
metrics_callback: For logging all numeric metricssample_callback: For logging text generation samples
- Multi-Cycle Tracking: Enhance the dashboard to track multiple pruning cycles
- Fine-Tuning Visualization: Add visualizations for the fine-tuning phase
- Attention Evolution: Track how attention patterns evolve across pruning cycles
- Model Growth: Visualize neural plasticity with model growth
The Neural Plasticity Dashboard provides a comprehensive visualization system for tracking, analyzing, and sharing neural plasticity experiments. It supports real-time monitoring, detailed phase visualization, and collaborative analysis through multiple viewing options.
🤖 Generated with Claude Code
Co-Authored-By: Claude noreply@anthropic.com