VAE are one of the most vasinating concepts in machine learning and computer vision. How do we teach computers to understand the meaning of images and shapes ? a very hard question that takes a very clever solution. It took me months of watching youtube videos and reading textbooks by building the math and code to write this.

Before you try the code I would highly suggest looking into my [blog][https://www.alinawaf.com/blog] and exploring my writing about some of the mathematical concepts that make the VAE work.

A comprehensive interactive tool for exploring Variational Autoencoders (VAEs) with real-time visualization, clustering analysis, and latent space manipulation. This project provides an intuitive interface for understanding how VAEs learn representations and generate new data.

• Interactive VAE Training: Train VAE models with adjustable hyperparameters
• Real-time Visualization: Monitor training progress and loss curves
• Latent Space Exploration: Visualize high-dimensional latent representations using PCA and t-SNE
• EM Clustering: Perform Gaussian Mixture Model clustering on learned representations
• Sample Generation: Generate new images by sampling from the latent space
• Latent Space Interpolation: Interactively walk through the latent space using sliders
• Reconstruction Analysis: Compare original images with their reconstructions

• Training loss monitoring
• 2D latent space projections (PCA and t-SNE)
• Clustering results with performance metrics
• Generated samples grid
• Original vs reconstructed images
• Interactive latent space navigation

• Installation
• Quick Start
• Architecture Overview
• Interface Guide
• Technical Details
• Customization
• Troubleshooting
• Contributing

• Python 3.7+
• CUDA-compatible GPU (optional but recommended)

pip install torch torchvision matplotlib numpy scikit-learn seaborn

# If you have the custom EM implementation
pip install scipy

git clone <repository-url>
cd VAE
mkdir data  # MNIST data will be downloaded here automatically
python interactive_vae_explorer.py

1. Run the Application:

   python interactive_vae_explorer.py

2. Basic Workflow:

   • Click "Train VAE" to start training with default parameters
   • Monitor the loss curve in real-time
   • After training, click "Run Clustering" to analyze latent representations
   • Use "Generate Samples" to create new images
   • Explore the latent space using the Z1/Z2 sliders

3. Experiment with Parameters:

   • Adjust latent dimension (Z Dim) before training
   • Modify learning rate and epochs
   • Observe how changes affect the learned representations

The Variational Autoencoder consists of three main components:

Input (28×28 MNIST image) 
    ↓
Linear(784 → 400) + ReLU
    ↓
Linear(400 → z_dim) [μ branch]
Linear(400 → z_dim) [log σ² branch]

z = μ + σ * ε, where ε ~ N(0,1)

Latent vector z (z_dim dimensional)
    ↓
Linear(z_dim → 400) + ReLU
    ↓
Linear(400 → 784) + Sigmoid
    ↓
Reshape to (28×28) image

The VAE optimizes the Evidence Lower BOund (ELBO):

L = E[log p(x|z)] - KL[q(z|x) || p(z)]
  = Reconstruction Loss + KL Divergence

Where:

• Reconstruction Loss: Binary Cross-Entropy between input and output
• KL Divergence: Regularization term ensuring latent distribution approximates N(0,I)

The interface is divided into a 3×4 grid of subplots:

• Controls Panel: Interactive buttons and sliders
• Loss Plot: Real-time training loss monitoring
• Latent PCA: 2D PCA projection of latent space colored by digit class
• Latent t-SNE: 2D t-SNE embedding showing local structure

• Original Images: Sample of input MNIST digits
• Reconstructed Images: VAE reconstructions of the same samples
• Generated Samples: New images generated from random latent vectors
• Interpolation: Real-time latent space walk visualization

• Cluster PCA: Latent space colored by cluster assignments
• Confusion Matrix: Comparison between true digits and cluster assignments
• Clustering Metrics: Adjusted Rand Index (ARI) and Normalized Mutual Information (NMI)
• Latent Walk: Interactive exploration using Z1/Z2 sliders

• Train VAE: Start training with current parameter settings
• Stop Training: Halt ongoing training process
• Run Clustering: Perform EM clustering on latent representations
• Generate Samples: Create new images from random latent codes

• Z Dim (2-20): Dimensionality of latent space
• Learning Rate (0.0001-0.01): Optimizer learning rate
• Epochs (1-50): Number of training epochs
• Z1/Z2 (-3 to 3): Interactive latent space coordinates for real-time generation

The application uses Python threading to enable:

• Concurrent Training: Model training runs in a separate thread
• Real-time Visualization: Matplotlib animation updates plots during training
• Responsive Interface: UI remains interactive during computation

• Efficient Data Loading: Uses PyTorch DataLoader with batching
• Limited History: Training loss history kept in a bounded deque
• GPU Utilization: Automatic CUDA detection and usage

• Purpose: Linear dimensionality reduction for global structure
• Advantages: Fast, deterministic, preserves global variance
• Use Case: Understanding primary directions of variation

• Purpose: Non-linear reduction preserving local neighborhoods
• Advantages: Reveals clusters and local structure
• Limitation: Computationally expensive, limited to samples for speed

The tool attempts to use a custom EM algorithm from EM/EM_ForVae.py:

try:
    from EM.EM_ForVae import em_gmm
    means, covs, weights, responsibilities = em_gmm(X, 10, max_iter=100)
except:
    # Fallback to sklearn
    from sklearn.mixture import GaussianMixture

• Adjusted Rand Index (ARI): Measures clustering quality vs. true labels
• Normalized Mutual Information (NMI): Information-theoretic clustering metric
• Confusion Matrix: Visual comparison of clusters vs. digit classes

To use different datasets, modify the setup_data() method:

def setup_data(self):
    # Replace MNIST with your dataset
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Resize((28, 28))  # Adjust size as needed
    ])
    
    train_dataset = YourDataset(transform=transform)
    # ... rest of setup

Modify the InteractiveVAE class for different architectures:

class InteractiveVAE(nn.Module):
    def __init__(self, z_dim=10, hidden=400, input_size=784):
        super().__init__()
        # Customize architecture here
        self.fc1 = nn.Linear(input_size, hidden)
        # Add more layers, different activations, etc.

Extend the interface by adding new subplot axes:

def setup_interface(self):
    # Add new subplot
    self.ax_new_plot = self.fig.add_subplot(gs[row, col])
    
def update_new_plot(self):
    # Implementation for new visualization
    pass

Error: RuntimeError: CUDA out of memory Solution:

• Reduce batch size in DataLoader
• Decrease model hidden size
• Use CPU instead: device = torch.device("cpu")

Warning: EM_ForVae not found, using sklearn fallback Solution: This is expected behavior. The tool automatically falls back to sklearn's GaussianMixture.

Error: Perplexity issues or convergence problems Solution: The code handles this with try-except blocks. Reduce sample size or adjust perplexity.

Issue: Interface becomes unresponsive Solutions:

• Reduce dataset size for testing
• Increase animation interval
• Use smaller latent dimensions
• Train on GPU if available

Issue: Visualizations remain static Solution:

• Ensure model is trained before clustering
• Check that matplotlib backend supports animation
• Try plt.ion() for interactive mode

# Reduce sample size for visualization
sample_size = min(1000, len(dataset))
indices = np.random.choice(len(dataset), sample_size, replace=False)

# Increase batch size (if memory allows)
DataLoader(dataset, batch_size=512, num_workers=4)

# Reduce animation frequency
FuncAnimation(fig, update_func, interval=2000)  # Update every 2 seconds

• Cluster Separation: Different digits form distinct clusters
• Smooth Transitions: Similar digits are close in latent space
• Meaningful Directions: Moving in latent space changes semantic properties

• PCA: Better for understanding global structure and primary variations
• t-SNE: Better for identifying local clusters and fine-grained structure

• Range: -1 to 1
• Interpretation:
  • 1.0: Perfect clustering
  • 0.0: Random clustering
  • <0.0: Worse than random

• Range: 0 to 1
• Interpretation:
  • 1.0: Perfect clustering
  • 0.0: No mutual information

• Sharpness: Generated images should be clear
• Diversity: Samples should show variety
• Realism: Images should resemble training data

• Similarity: Reconstructions should preserve key features
• Smoothness: Minor variations acceptable due to stochastic nature

• Understanding VAEs: Interactive exploration of key concepts
• Hyperparameter Effects: Observe impact of different settings
• Latent Space Properties: Visualize learned representations

• Architecture Comparison: Compare different VAE variants
• Regularization Effects: Study β-VAE and other modifications
• Disentanglement Analysis: Evaluate factor separation in latent space

• Anomaly Detection: Identify outliers in latent space
• Data Augmentation: Generate synthetic training samples
• Feature Learning: Extract meaningful representations

git clone <repository-url>
cd VAE
pip install -e .  # Editable installation

VAE/
├── interactive_vae_explorer.py  # Main application
├── EM/                         # Custom EM implementation (optional)
│   └── EM_ForVae.py
├── data/                       # Dataset storage
└── README.md                   # This file

1. Fork the repository
2. Create a feature branch
3. Add tests for new functionality
4. Update documentation
5. Submit a pull request

• Support for additional datasets (CIFAR-10, CelebA)
• Advanced VAE architectures (β-VAE, WAE)
• 3D latent space visualization
• Model comparison tools
• Export functionality for trained models
• Batch experiment running

This project is open source. Please check the LICENSE file for details.

1. Auto-Encoding Variational Bayes - Kingma & Welling (2013)
2. β-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework - Higgins et al. (2017)
3. Understanding disentangling in β-VAE - Burgess et al. (2018)

• VAE Tutorial
• PyTorch VAE Examples
• Disentanglement in VAEs

For questions, issues, or contributions:

1. Check existing issues in the repository
2. Create a new issue with detailed description
3. Include error messages and system information
4. Provide minimal reproducible examples

Happy Exploring! 🚀

This interactive tool is designed to make VAE concepts accessible and engaging. Whether you're learning about generative models or conducting research, the visual feedback and real-time interaction provide valuable insights into how these powerful models work.