HiSTM: Hierarchical SpatioTemporal Mamba for Spatiotemporal Forecasting of Cellular Traffic Networks

HiSTM (“High S T M” — /haɪ ɛs tiː ɛm/) is a hierarchical Mamba-based neural architecture for spatiotemporal forecasting tasks such as cellular network traffic prediction. It combines convolutional spatial encoding, state-space temporal modeling (Mamba), and temporal attention aggregation into a unified, lightweight framework.

---

Overview

HiSTM addresses the challenge of forecasting spatially distributed time series (e.g., traffic volumes over a 2D grid) by capturing:

• Spatial correlations via convolutional encoders
• Temporal dependencies via Mamba selective state-space models
• Dynamic relevance of time steps via a temporal attention mechanism

The model learns to predict the next-step traffic intensity at the center cell of each spatiotemporal patch.

---

Key Features

• Hierarchical Encoding – Stacked encoder layers extract multiscale spatiotemporal features.
• Mamba Temporal Modeling – Leverages the efficient state-space dynamics of Mamba (Gu et al., 2024) to replace recurrent or attention-based temporal modules.
• Temporal Attention – Aggregates all time-step features adaptively instead of using only the last hidden state.
• Lightweight & Efficient – Comparable or superior performance to transformer baselines with fewer parameters.

---

Problem Formulation

Given a sequence of past spatial observations

$$ M = {M_1, M_2, \dots, M_T}, \quad M_t \in \mathbb{R}^{H \times W} $$

the task is to predict the next-step traffic volume $M_{T+1}$.

For each spatial coordinate $(i, j)$, we extract a spatiotemporal patch

$$ X_t^{(i,j)} \in \mathbb{R}^{T \times K \times K} $$

and predict the next value at the center cell:

$$ \hat{y}_{t+1}^{(i,j)} = f_\theta(X_t^{(i,j)}) $$

where $f_\theta$ denotes the HiSTM model.

---

Model Architecture

1. Encoder Layers

Each encoder layer performs:

• 2D Convolution: Captures local spatial patterns at each time step.
• Mamba SSM: Models temporal dependencies across the sequence for each spatial position independently.

2. Temporal Attention

The encoded temporal features at the center spatial cell are combined using attention weights:

$$ \alpha_t = \frac{\exp(e_t)}{\sum_{j=1}^T \exp(e_j)}, \quad e_t = \text{Linear}_{att}(h_t) $$

producing a weighted context vector $c = \sum_t \alpha_t h_t$.

3. Prediction Head

A two-layer MLP maps the aggregated context vector to the final scalar prediction:

$$ \hat{y} = \text{MLP}(c) $$

---

Installation

1. Clone and install locally

git clone https://github.com/ZineddineBtc/histm.git
cd histm
pip install -e .

2. Dependencies

• Python ≥ 3.8
• PyTorch ≥ 2.0
• mamba-ssm

You can install dependencies manually:

pip install torch mamba-ssm

---

Usage

Example: Quick test

import torch
from histm import HiSTM

# Define input dimensions
B, T, H, W = 32, 6, 11, 11
x = torch.randn(B, T, H, W)

# Initialize model
model = HiSTM(device='cuda' if torch.cuda.is_available() else 'cpu')

# Forward pass
with torch.no_grad():
    y = model(x)

print("Output shape:", y.shape)  # (B, 1)

---

Baseline Comparisons

HiSTM was benchmarked against several state-of-the-art spatiotemporal forecasting models:

Model	Description
STN	CNN-based spatiotemporal forecasting network
xLSTM	Scalable exponential-gate LSTM
STTRE	Transformer with relative embeddings
VMRNN	Vision Mamba–LSTM hybrid
PredRNN++	Recurrent memory network with spatiotemporal flow
PatchTST	Patch-based Transformer for time series
TimesNet	Time–frequency decomposition model
Informer	ProbSparse Transformer for efficient long forecasting
Autoformer	Decomposition-based auto-correlation Transformer

---

Citation

If you use HiSTM in your research, please cite:

@article{Bettouche2025histm,
  title={HiSTM: Hierarchical SpatioTemporal Mamba for Cellular Traffic Forecasting},
  author={Zineddine Bettouche, Khalid Ali, Andreas Fischer, Andreas Kassler},
  year={2025},
  journal={arXiv preprint arXiv:2508.09184}
}

---

Acknowledgements

HiSTM builds upon:

• Mamba SSM
• PyTorch
• Foundational ideas from spatiotemporal modeling and time series forecasting research.

---

Repository Structure

histm/
├── histm/
│   ├── __init__.py          # Exports HiSTM
│   ├── model.py             # Core HiSTM architecture
│   ├── __version__.py
├── pyproject.toml
├── README.md
└── examples/
    └── test_forward.py

---

Developed for research in spatiotemporal forecasting.