Skip to content

Latest commit

 

History

History
executable file
·
157 lines (112 loc) · 5.17 KB

File metadata and controls

executable file
·
157 lines (112 loc) · 5.17 KB

Neural Plasticity Dashboard Implementation

Overview

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)

Key Features

  1. 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
  2. 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
  3. Multiple Dashboard Options:

    • Weights & Biases integration for online tracking
    • Standalone HTML dashboard for offline viewing
    • Real-time console output for immediate feedback
  4. Collaboration Features:

    • Shareable dashboard links for team analysis
    • Browser integration for automatic dashboard opening
    • Support for both local and Colab environments

Implementation Components

1. Core Dashboard Integration

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

2. Experiment Phase Visualizations

Warmup Phase

  • Loss curve visualization with smoothed trend line
  • Training metric tracking (loss, perplexity)
  • Text samples during training

Analysis Phase

  • Entropy heatmap for attention heads
  • Gradient heatmap for attention heads
  • Combined importance metrics visualization

Pruning Phase

  • Pruning decision visualization
  • Before/after model architecture comparison
  • Pruned heads highlighted by layer

Evaluation Phase

  • Side-by-side text generation comparison
  • Perplexity comparison bar charts
  • Attention pattern visualization

3. Standalone Dashboard Demo

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

4. Colab Integration

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

Usage Instructions

Running the Dashboard Demo

# 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

Running a Full Experiment with Dashboard

# 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

Viewing the Dashboard

  1. Online Mode: Dashboard URL will be provided in console output, or:

    open https://wandb.ai/[username]/neural-plasticity/runs/[run_id]
    
  2. Offline Mode: Open the local HTML file:

    open output/neural_plasticity_[timestamp]/dashboard.html
    
  3. Colab Environment: Click the dashboard link displayed in the notebook cell.

Implementation Details

Metrics Tracking

The dashboard tracks key metrics:

  • Training and evaluation loss
  • Perplexity
  • Sparsity (percentage of pruned heads)
  • Inference speedup
  • Model size reduction

Visualization Methods

  1. Heatmaps: For entropy, gradients, and attention patterns
  2. Line Charts: For loss curves and performance tracking
  3. Bar Charts: For perplexity and other comparative metrics
  4. Grid Visualizations: For model architecture and pruning decisions
  5. Text Samples: For qualitative evaluation of model outputs

Callback System

The dashboard provides callback functions for the experiment:

  1. metrics_callback: For logging all numeric metrics
  2. sample_callback: For logging text generation samples

Next Steps

  1. Multi-Cycle Tracking: Enhance the dashboard to track multiple pruning cycles
  2. Fine-Tuning Visualization: Add visualizations for the fine-tuning phase
  3. Attention Evolution: Track how attention patterns evolve across pruning cycles
  4. Model Growth: Visualize neural plasticity with model growth

Conclusion

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