A Rigorous Factorial Comparison of Neural Network Architectures

Python 3.9+ PyTorch 2.0+ License: MIT CI Code style: black

Documentation | Quick Start | Results | Citation

Overview

Research Question

Architecture

┌─────────────────────────────────────────────────────────────────┐
│ ×ばつ3 FACTORIAL DESIGN │
├─────────────────────────────────────────────────────────────────┤
│ │
│ Feedforward Type Sequence Model │
│ ════════════════ ══════════════ │
│ │
│ ┌─────────────┐ ┌─────────────┐ │
│ │ MLP │──────────│ Transformer │ ─► mlp_transformer │
│ │ (ReLU/GELU)│ │ (Attention) │ │
│ └─────────────┘ └─────────────┘ │
│ │ │ │
│ │ │ │
│ ┌─────────────┐ ┌─────────────┐ │
│ │ MLP + │──────────│ Transformer │ ─► bspline_transformer│
│ │ B-spline │ │ (Attention) │ │
│ └─────────────┘ └─────────────┘ │
│ │ │ │
│ │ │ │
│ ┌─────────────┐ ┌─────────────┐ │
│ │ Full KAN │──────────│ Transformer │ ─► kan_transformer │
│ │ (Learnable)│ │ (Attention) │ │
│ └─────────────┘ └─────────────┘ │
│ │ │ │
│ │ ┌─────────────┐ │
│ └──────────────────│ Mamba │ ─► *_mamba variants │
│ │ (SSM) │ │
│ └─────────────┘ │
│ │
└─────────────────────────────────────────────────────────────────┘

Six Model Variants

Key Results

Results Visualization

Accuracy Comparison

Training Speed Comparison

Key Findings

Conclusions

Quick Start

Installation

# Clone the repository
git clone https://github.com/stchakwdev/Mamba_KAN.git
cd Mamba_KAN
# Create environment
conda create -n mamba_kan python=3.10
conda activate mamba_kan
# Install PyTorch (adjust CUDA version as needed)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
# Install package with development dependencies
pip install -e ".[dev]"

Run Quick Validation

# Quick test (1 model, 1 seed, 100 steps)
make train
# Or directly:
python scripts/run_experiment.py \
 --model mlp_transformer \
 --task symbolic \
 --seeds 1 \
 --max-steps 100 \
 --no-wandb

Run Full Comparison

# Full factorial experiment (all models, 10 seeds)
python scripts/run_full_comparison.py \
 --tasks symbolic special_functions timeseries language long_context \
 --seeds 10 \
 --max-steps 10000 \
 --output-dir ./results \
 --generate-report

Benchmark Tasks

1. Symbolic Regression (KAN-favorable)

2. Time Series Forecasting

3. Language Modeling (Mamba-favorable)

Project Structure

mamba_kan/
├── models/ # Model implementations
│ ├── base.py # Task-aware base class
│ ├── mlp_transformer.py # MLP-Transformer (baseline)
│ ├── bspline_transformer.py # B-spline activation baseline
│ ├── kan_transformer.py # Full KAN-Transformer
│ ├── *_mamba.py # Mamba variants
│ └── components/
│ ├── bspline_mlp.py # Learnable B-spline activation
│ ├── kan_layers.py # KAN building blocks
│ ├── mamba_layers.py # Mamba with B-spline support
│ └── transformer_layers.py
├── training/
│ ├── trainer.py # PyTorch Lightning module
│ ├── scheduler.py # Learning rate schedules
│ └── callbacks.py # Training monitoring
├── analysis/
│ └── statistics.py # Friedman, Wilcoxon, bootstrap CI
├── visualization/ # Plotting and dashboards
│ ├── plots.py # Training curves, comparisons
│ ├── heatmaps.py # Statistical visualizations
│ ├── animations.py # GIF generation
│ └── dashboard.py # Interactive HTML reports
├── data/
│ └── datasets.py # All benchmark datasets
└── configs/
 └── base_config.py # Configuration system
scripts/
├── run_experiment.py # Single experiment runner
├── run_full_comparison.py # Full factorial comparison
├── generate_assets.py # Generate README visualizations
└── runpod_setup.sh # Cloud GPU setup

Statistical Analysis

from mamba_kan.analysis import run_full_analysis, print_analysis_summary
results = {
 'mlp_transformer': [0.52, 0.51, 0.53, ...], # Loss per seed
 'bspline_transformer': [0.48, 0.47, 0.49, ...],
 'kan_transformer': [0.45, 0.44, 0.46, ...],
 # ... other models
}
analysis = run_full_analysis(results)
print_analysis_summary(analysis)

Hardware Requirements

Cloud Deployment

# RunPod H100 setup
chmod +x scripts/runpod_setup.sh
./scripts/runpod_setup.sh
# Run full experiment suite
python scripts/run_full_comparison.py --task all --seeds 10

Development

# Install dev dependencies
pip install -e ".[dev]"
# Run tests
make test
# Run linting
make lint
# Format code
make format
# Generate visualizations from results
make visualize

Documentation

References

Papers

Resources

Citation

@misc{mamba_kan_2025,
 title={Mamba-KAN: A Factorial Comparison of Neural Network Architectures},
 author={Samuel T. Chakwera},
 year={2025},
 url={https://github.com/stchakwdev/Mamba_KAN},
 note={Investigating whether KAN advantages stem from B-spline activations or network topology}
}

License

Back to Top

Made with PyTorch Lightning and scientific rigor