---
license: cc-by-nc-sa-4.0
---


# **DriftMoE** – A Mixture of Experts Approach to Handle Concept Drifts

Model weights for paper [DriftMoE](https://www.arxiv.org/abs/2507.18464)This repository hosts *weights only* so you can plug the model straight into your Python pipeline. These weights correspond to one training run on the LED_g stream.
Full training code & utilities live in a separate GitHub repo: **[https://github.com/miguel-ceadar/drift-moe](https://github.com/miguel-ceadar/drift-moe)**

---

## 📂 Files

We have two folders, one for the MoE-Task variant and other for the MoE-Data variant, both share this file structure:
```
router.pth           # PyTorch state_dict for the gating MLP
expert_0.pkl         # CapyMOA HoeffdingTree for each expert(pickled)
expert_1.pkl
…
expert_{N‑1}.pkl
```

---

## ⚡ Quick Start (CPU or GPU)

### 1 · Install runtime deps

You need to install Java and have a working java Runtime Environment to run Capymoa:
[https://openjdk.org/install/](https://openjdk.org/install/)
```bash
python -m pip install torch capymoa numpy river
# git clone training repo – needed so we can recreate RouterMLP & Expert wrappers
git clone https://github.com/miguel-ceadar/drift-moe drift_moe
```

### 2 · Load the router & experts

```python
import torch, pickle, numpy as np 
from capymoa.misc import load_model
from drift_moe.driftmoe.moe_model import RouterMLP, Expert

INPUT_DIM   = 24      # ↩ must match dimensions of LED_g stream
NUM_CLASSES = 10      # ↩ idem
N_EXPERTS   = 12      # ↩ number of expert_*.pkl files (12 for MoE_Data and 10 for MoE_Task)
DEVICE      = 'cpu'   # or 'cuda'

# 2‑a) Router
router = RouterMLP(input_dim=INPUT_DIM, hidden_dim=256, output_dim=N_EXPERTS)
router.load_state_dict(torch.load('path/to/router.pth', map_location=DEVICE))
router = router.to(DEVICE).eval()

# 2‑b) Experts (pickled CapyMOA trees)
experts = []
for i in range(N_EXPERTS):
    with open(f'path/to/expert_{i}.pkl', 'rb') as f:
        ex = load_model(f)               # HoeffdingTree object
    
    experts.append(ex)

```

*Reference: the official CapyMOA [save & load notebook](https://capymoa.org/notebooks/save_and_load_model.html).*

### 3 · Single‑sample inference helper

```python
def predict_one(instance) -> int:
    """Route a single feature vector through driftMoE and return class index."""
    x_vec = instance.x
    x_t = torch.tensor(x_vec, dtype=torch.float32).unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        logits = router(x_t)                 # shape [1, N_EXPERTS]
        eid    = int(torch.argmax(logits, 1).item())
    return experts[eid].predict(instance)


```

---

## 🚰 Streaming usage

```python
from capymoa.stream.generator import LEDGenerator
stream = LEDGenerator()
while stream.has_more_instances():
    inst = stream.next_instance()
    y_hat = predict_one(inst)
    print(y_hat)
    print(inst.y_index)
```

The experts are frozen; only the router runs every forward pass.


## ✏️ Citation

```bibtex
@misc{aspis2025driftmoemixtureexpertsapproach,
      title={DriftMoE: A Mixture of Experts Approach to Handle Concept Drifts}, 
      author={Miguel Aspis and Sebastián A. Cajas Ordónez and Andrés L. Suárez-Cetrulo and Ricardo Simón Carbajo},
      year={2025},
      eprint={2507.18464},
      archivePrefix={arXiv},
      primaryClass={stat.ML},
      url={https://arxiv.org/abs/2507.18464}, 
}
```


*Questions or issues?* Open an issue on the GitHub repo and we’ll be happy to help.