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

2026-05-29

AI Summary

This study addresses the limitations of existing energy consumption forecasting methods, which often neglect spatial dependencies among regions and struggle to provide reliable uncertainty estimates under extreme weather conditions. To overcome these challenges, this work proposes the EnergyMamba framework, which introduces a graph-enhanced selective state space model (GE-Mamba) that integrates power grid topology with temporal dynamics. Additionally, an adaptive sequential conformal quantile regression module (AS-CQR) is designed, combining local adaptive normalization with online feedback to dynamically calibrate prediction intervals under distributional shifts. Evaluated on four large-scale real-world datasets, the proposed approach improves point forecasting accuracy by approximately 5% and achieves a roughly 6% gain in uncertainty quantification performance over 15 state-of-the-art baselines.

Abstract

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.

Problem

• energy consumption prediction
• spatiotemporal modeling
• uncertainty quantification
• spatial dependencies
• abnormal situations

Innovation

• Graph-Enhanced Selective State Space Model
• Uncertainty-Aware Prediction
• Spatiotemporal Modeling
• Adaptive Conformal Quantile Regression
• Energy Consumption Forecasting