# Multitask Modality Routing: Supplementary Code

This repository contains the supplementary code for our NeurIPS 2025 submission on multitask learning with modality routing. The codebase includes scripts for training, analyzing, and visualizing a modular neural network that routes information across multiple modalities and tasks.

---

## Project Structure

```
.
├── experiments/
│   ├── train_routing.py         # Script to train the routing model
│   ├── analyze_routing.py       # Script to analyze and visualize routing behavior
│   └── results/
│       ├── scenario1_baseline/  # Results and plots for scenario 1
│       └── scenario2_case2/     # Results and plots for scenario 2
├── src/
│   ├── models/                  # Model definitions (routers, experts, etc.)
│   ├── utils/                   # Utility functions (data simulation, plotting, etc.)
│   └── configs/                 # (Empty or for future configs)
```

---

## Installation

1. **Clone the repository:**
   ```bash
   git clone <your-repo-url>
   cd <repo-folder>
   ```
2. **Install dependencies:**
   ```bash
   pip install -r requirements.txt
   ```
   - Python 3.8+ is recommended.

---

## Usage

### 1. Train the Routing Model

```bash
python experiments/train_routing.py --n_samples 1000 --batch_size 32 --n_epochs 100 --output_dir experiments/results/scenario1_baseline
```
- See `--help` for all options:
  ```bash
  python experiments/train_routing.py --help
  ```

### 2. Analyze Routing and Results

```bash
python experiments/analyze_routing.py --results_dir experiments/results/scenario1_baseline
```
- This will generate plots and analysis in the specified results directory.

---

## Results & Plots
- All key results and plots are in `experiments/results/scenario1_baseline/` and `experiments/results/scenario2_case2/`.
- Example plots:
  - `routing_heatmap.png`, `routing_distribution.png`, `prediction_analysis.png`, etc.
- Best model checkpoints: `best_model.pt`

---

## Code Overview

- **experiments/train_routing.py**: Main training script for the routing model.
- **experiments/analyze_routing.py**: Analysis and visualization of routing behavior and model predictions.
- **src/models/**: Model components (routers, experts, transforms).
- **src/utils/**: Data simulation, plotting, and common utilities.

---

## Reproducibility
- All random seeds are set for reproducibility.
- Synthetic data is generated on the fly; no external datasets are required.
- To reproduce results, run the training script followed by the analysis script as shown above.

---

## License
- MIT

---

## Contact
For questions or issues, please contact [redacted].

---

