[1000] Adding MoEMLP layer to the forecasting engine

[wael-mika]

Description

This draft PR introduces an experimental Mixture-of-Experts (MoE) MLP block as a drop-in replacement for the standard dense MLP in the Forecasting Engine.
The goal is to improve predictive skill and increase model capacity by encouraging expert specialization for different atmospheric regimes.

Key Changes:
Interface Preservation: The MoE block maintains the existing forward(*args) signature.
Architecture: Implements a lightweight top-k router and multiple small FFN experts (configurable num_experts, top_k).
Control: Enabled via a config flag (e.g., fe_mlp_type: "moe") for easy testing.

Issue Number

Closes #1000

This is a draft PR.

Checklist before asking for review

• I have performed a self-review of my code
• My changes comply with basic sanity checks:
  • I have fixed formatting issues with ./scripts/actions.sh lint
  • I have run unit tests with ./scripts/actions.sh unit-test
  • I have documented my code and I have updated the docstrings.
  • I have added unit tests, if relevant
• I have tried my changes with data and code:
  • I have run the integration tests with ./scripts/actions.sh integration-test
  • (bigger changes) I have run a full training and I have written in the comment the run_id(s): launch-slurm.py --time 60
  • (bigger changes and experiments) I have shared a hegdedoc in the github issue with all the configurations and runs for this experiments
• I have informed and aligned with people impacted by my change:
  • for config changes: the MatterMost channels and/or a design doc
  • for changes of dependencies: the MatterMost software development channel

[sophie-xhonneux]

Wael, thanks for this work, could you add some inference plots when you have them to the PR?

[wael-mika]

Hey @sophie-xhonneux, do you mean the actual physical output vs the target?

[clessig]

validation scores, please

[wael-mika]

RMSE, MAE, PSNR for arcsinh transformation with and without MoE on

Training plots

[clessig]

Thanks!

For me the difference are in the random seed noise range. How does it look for the other meterological variables (t2m, q_850, t_850, z_500)?