MambaEP: Efficient Spatio-Temporal Forecasting with Mamba Block and Fourier-Enhanced Temporal Modeling
Spatio-temporal forecasting plays a vital role in a range of scientific and practical applications, from weather prediction to traffic forecasting. However, existing models often struggle with balancing accuracy, efficiency, and scalability. This paper presents MambaEP, an efficient framework for spatio-temporal forecasting that integrates the Mamba Block with Fourier-enhanced temporal modeling. By combining local and global spatial feature extraction with multi-scale temporal processing, MambaEP achieves superior prediction accuracy while maintaining computational efficiency. Extensive experiments on benchmark datasets, including SEVIR, RainNet, and MovingMNIST, demonstrate that MambaEP outperforms state-of-the-art methods in both accuracy and efficiency, achieving faster convergence and smaller model sizes. Our model's ability to efficiently handle long-range dependencies, identify extreme cases, and reduce computational costs positions it as a promising solution for spatio-temporal forecasting tasks.
Figure below present the evaluation outcomes of MambaEP alongside ten state-of-the-art (SOTA) baselines across six datasets. On the Moving MNIST dataset, MambaEP achieves the lowest MSE (scaled by 100), outperforming prominent methods such as EarthFarsser, Earthformer, and SimVP. Similarly, on the Component of Wind dataset, MambaEP achieves the lowest MSE and MAE scores, demonstrating superior accuracy compared to other approaches. Overall, MambaEP attains the best average performance ranking of 2.09, highlighting its consistent effectiveness across all datasets. These results validate the model's ability to capture and model spatial features effectively through the integration of the Mamba Block.
- Install Python 3.10. For convenience, execute the following command.
pip install -r requirements.txt- Prepare Data. You can obtain experimental datasets from the following links.
| Dataset | Task | Categories | Link |
|---|---|---|---|
| WeatherBench | Predict future natural representation | Natural Scene | [Google Cloud] |
| SEVIR | Predict future natural representation | Natural Scene | [AWS Opendata] |
| RainNet | Predict future natural representation | Natural Scene | [Github Link] |
| Moving MNIST | Predict future image | Synthetic | [Link] |
- Use the following instructions to quickly run the code. (Using Moving MNIST as a demonstration example)
python main.py --data_path dataset --num_epochs 100 --batch_size 4 --val_batch_size 4 --lr 0.001