# EnergyMamba: An Uncertainty-Aware Graph-Enhanced Selective State Space Model for Energy Consumption Prediction

arXiv: https://arxiv.org/abs/2606.00506 · HTML: https://mlsanity.com/paper/2606.00506/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 60 references that are in the map.
- First posted on arXiv 30 May 2026.

## Builds on (most cited first, 27 of 27)

- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [A Time Series is Worth 64 Words: Long-term Forecasting with Transformers](https://mlsanity.com/paper/2211.14730/) (2022)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [Spatio-Temporal Graph Convolutional Networks: A Deep Learning Framework for Traffic Forecasting](https://mlsanity.com/paper/1709.04875/) (2017)
- [Temporal Fusion Transformers for Interpretable Multi-horizon Time Series Forecasting](https://mlsanity.com/paper/1912.09363/) (2019)
- [Conformalized Quantile Regression](https://mlsanity.com/paper/1905.03222/) (2019)
- [Adaptive Graph Convolutional Recurrent Network for Traffic Forecasting](https://mlsanity.com/paper/2007.02842/) (2020)
- [Adaptive Conformal Inference Under Distribution Shift](https://mlsanity.com/paper/2106.00170/) (2021)
- [ST-LLM: Large Language Models Are Effective Temporal Learners](https://mlsanity.com/paper/2404.00308/) (2024)
- [UrbanGPT: Spatio-Temporal Large Language Models](https://mlsanity.com/paper/2403.00813/) (2024)
- [Uncertainty Quantification of Sparse Travel Demand Prediction with Spatial-Temporal Graph Neural Networks](https://mlsanity.com/paper/2208.05908/) (2022)
- [GeoGen: A Two-stage Coarse-to-Fine Framework for Fine-grained Synthetic Location-based Social Network Trajectory Generation](https://mlsanity.com/paper/2510.07735/) (2025)
- [Regulating Model Reliance on Non-Robust Features by Smoothing Input Marginal Density](https://mlsanity.com/paper/2407.04370/) (2024)
- [PowerPM: Foundation Model for Power Systems](https://mlsanity.com/paper/2408.04057/) (2024)
- [UQGNN: Uncertainty Quantification of Graph Neural Networks for Multivariate Spatiotemporal Prediction](https://mlsanity.com/paper/2508.08551/) (2025)
- [HealthMamba: An Uncertainty-aware Spatiotemporal Graph State Space Model for Effective and Reliable Healthcare Facility Visit Prediction](https://mlsanity.com/paper/2602.05286/) (2026)
- [Uncertainty-aware Predict-Then-Optimize Framework for Equitable Post-Disaster Power Restoration](https://mlsanity.com/paper/2508.04780/) (2025)
- [HCRide: Harmonizing Passenger Fairness and Driver Preference for Human-Centered Ride-Hailing](https://mlsanity.com/paper/2508.04811/) (2025)
- [LLM as Clinical Graph Structure Refiner: Enhancing Representation Learning in EEG Seizure Diagnosis](https://mlsanity.com/paper/2604.28178/) (2026)
- [TyphoFormer: Language-Augmented Transformer for Accurate Typhoon Track Forecasting](https://mlsanity.com/paper/2506.17609/) (2025)
- [STS-CCL: Spatial-Temporal Synchronous Contextual Contrastive Learning for Urban Traffic Forecasting](https://mlsanity.com/paper/2307.02507/) (2023)
- [BTS: A Comprehensive Benchmark for Tie Strength Prediction](https://mlsanity.com/paper/2410.19214/) (2024)
- [Learning from the Storm: A Multivariate Machine Learning Approach to Predicting Hurricane-Induced Economic Losses](https://mlsanity.com/paper/2506.17964/) (2025)
- [TrustEnergy: A Unified Framework for Accurate and Reliable User-level Energy Usage Prediction](https://mlsanity.com/paper/2601.13422/) (2026)
- [CREDIT: Certified Ownership Verification of Deep Neural Networks Against Model Extraction Attacks](https://mlsanity.com/paper/2602.20419/) (2026)
- [Adaptive traffic signal control optimization using a novel road partition and multi-channel state representation method](https://mlsanity.com/paper/2602.12296/) (2026)
- [CITED: A Decision Boundary-Aware Signature for GNNs Towards Model Extraction Defense](https://mlsanity.com/paper/2602.20418/) (2026)

## Built on by (most cited first, 2 of 2)

- [NeuroGRIP: Retrieval-Augmented Graph Refinement for Knowledge-Grounded EEG Seizure Diagnosis](https://mlsanity.com/paper/2607.14314/) (2026)
- [Capability Scaling-Down Laws for LLM Compression](https://mlsanity.com/paper/2610.02462/) (2026)

## Abstract (arXiv)

Energy consumption prediction is essential for efficient grid management, demand-side optimization, and sustainable energy planning. Although advanced machine learning methods have been employed for better prediction performance, existing works have two key limitations: (1) they usually formulate this task as a purely time-series prediction problem without explicitly modeling the spatial dependencies among different regions, and (2) they fail to provide reliable predictions with uncertainty estimates under abnormal situations such as extreme weather events. To advance existing research, we propose EnergyMamba, an uncertainty-aware spatiotemporal learning framework for accurate and reliable energy consumption prediction, which comprises two key components: (i) a novel Graph-Enhanced Selective State Space Model (GE-Mamba) that injects spatial context learned from the grid topology into the temporal dynamics, enabling coupled spatiotemporal modeling, and (ii) an Adaptive Sequential Conformalized Quantile Regression (AS-CQR) module, which includes locally adaptive normalization and an online feedback mechanism to dynamically calibrate prediction intervals under potential distribution shifts. We evaluate EnergyMamba on four large-scale real-world datasets from Florida, New York, and California. Results show EnergyMamba achieves around 5% improvement in prediction accuracy and 6% improvement in uncertainty quantification over 15 state-of-the-art baselines.