Soft Gate Mixture of Experts for RUL Prediction

Reference implementation for the paper:

Soft Gate Mixture of Experts for Remaining Useful Life Prediction of Turbofan Engines Under Multiple Operating Conditions César Augusto Grisa, Instituto de Computação, Universidade Federal Fluminense (UFF).

The model combines a causal Temporal Convolutional Network (TCN), multi-head self-attention, and a differentiable soft gate that blends operational-regime experts continuously rather than routing each input to a single expert.

---

Quick start on Google Colab

Two self-contained notebooks reproduce the main results without any local setup (mount Google Drive, point to the data, run all cells):

• notebooks/Soft_Gate_MoE_CMAPSS_Colab.ipynb — CMAPSS blind test (FD001–FD004)
• notebooks/Soft_Gate_MoE_NCMAPSS_Colab.ipynb — N-CMAPSS DS02 (RMSE 5.85)

For the command-line / production version, use the scripts below.

Results (summary)

Dataset	Metric	Soft Gate MoE	Best baseline
CMAPSS FD002	NASA	459	563
CMAPSS FD004	NASA	814	877
N-CMAPSS DS02	RMSE	5.85	9.08

Cross-simulator fine-tuning of only 35% of parameters for 10 epochs reduces transfer RMSE from 70.86 (zero-shot) to 11.17 (-84%).

All CMAPSS numbers are the ensemble of 5 seeds [42, 123, 777, 999, 31415].

---

Installation

git clone https://github.com/cagrisa/soft-gate-moe.git
cd soft-gate-moe
pip install -r requirements.txt

Tested with Python 3.10 and PyTorch 2.x (CPU or single GPU).

---

Data

This repository does not redistribute the datasets. Download them from the original sources:

• CMAPSS — NASA Prognostics Data Repository. Place train_FDxxx.txt, test_FDxxx.txt, RUL_FDxxx.txt in ./data/cmapss/.
• N-CMAPSS DS02 — Figshare (doi:10.6084/m9.figshare.20436504). Place N-CMAPSS_DS02-006.h5 in ./data/ncmapss/.

---

Reproducing the paper

CMAPSS blind test (Tables 1-4, Figure 1)

python scripts/run_cmapss.py --data_dir ./data/cmapss --subset FD004
python scripts/run_cmapss.py --data_dir ./data/cmapss --subset FD002
python scripts/run_cmapss.py --data_dir ./data/cmapss --subset FD001
python scripts/run_cmapss.py --data_dir ./data/cmapss --subset FD003

Each run prints per-seed NASA/RMSE/R2, the 95% confidence interval of the mean, the best single seed, and the ensemble.

N-CMAPSS DS02 full training (Table 4, Figure 2)

python scripts/run_ncmapss.py --data_file ./data/ncmapss/N-CMAPSS_DS02-006.h5

Reproduces the full-training condition (3 seeds, ensemble): RMSE = 5.85, R2 = 90.5% over all test windows.

Figures (English-labelled)

python scripts/make_figures.py

---

Repository layout

soft-gate-moe/
├── src/
│   ├── model.py      # Soft Gate MoE architecture
│   ├── data.py       # loading, leakage-safe preprocessing, clustering
│   └── train.py      # training loop, NASA score, metrics, 95% CI
├── scripts/
│   ├── run_cmapss.py  # reproduce CMAPSS blind test
│   ├── run_ncmapss.py # reproduce N-CMAPSS DS02 full training
│   └── make_figures.py
├── notebooks/
│   ├── Soft_Gate_MoE_CMAPSS_Colab.ipynb
│   └── Soft_Gate_MoE_NCMAPSS_Colab.ipynb
├── configs/
│   └── default.yaml  # all fixed hyperparameters
├── requirements.txt
└── LICENSE

---

Leakage prevention

The preprocessing in src/data.py enforces:

1. Train/test split at the unit level before windowing.
2. StandardScaler fitted on training data only.
3. KMeans clustering fitted on training data only.
4. No test window is ever used during training (final-epoch weights are used for evaluation; there is no test-set checkpointing).
5. Evaluation over all test windows (N-CMAPSS) / last window per unit (CMAPSS), following the respective benchmark protocols.

---

Citation

@article{grisa_softgatemoe,
  title   = {Soft Gate Mixture of Experts for Remaining Useful Life
             Prediction of Turbofan Engines Under Multiple Operating
             Conditions},
  author  = {Grisa, C{\'e}sar Augusto},
  note    = {Manuscript},
  year    = {2026}
}

License

Released under the MIT License (see LICENSE).