🤖 machine learning

STM3: Mixture of Multiscale Mamba for Long-Term Spatio-Temporal Time-Series Prediction

The paper proposes STM3, a novel framework combining a Multiscale Mamba architecture with a Disentangled Mixture-of-Experts structure and an adaptive graph causal network to efficiently capture complex long-term spatio-temporal dependencies, achieving state-of-the-art performance across ten real-world benchmarks.

Original authors: Haolong Chen, Liang Zhang, Zhengyuan Xin, Guangxu Zhu

Published 2026-05-21

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Haolong Chen, Liang Zhang, Zhengyuan Xin, Guangxu Zhu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to predict the future traffic flow of a massive city. You have data from thousands of sensors (nodes) recording speed and volume every few minutes. This isn't just a simple line graph; it's a complex web where what happens on a highway affects a downtown street, and patterns repeat at different speeds: some things change every minute (a red light turning green), some every hour (rush hour), and some every day (the weekend lull).

The paper introduces STM3, a new AI model designed to solve the headache of predicting this kind of "spatio-temporal" data over long periods. Here is how it works, broken down into simple concepts and analogies.

The Problem: The "Messy Kitchen" of Data

Existing AI models struggle with two main issues when looking at long-term data:

The Scale Problem: Data contains information at many different "scales" (speeds) mixed together. It's like trying to hear a whisper, a conversation, and a siren all at once. Old models often mash these together, losing the details.
The Routing Problem: Different parts of the city behave differently. A highway behaves like a fast river; a downtown area behaves like a crowded market. Old models try to use one "brain" to understand both, or they switch brains too quickly and chaotically, causing confusion.

The Solution: STM3 (The "Specialized Team")

The authors built STM3, which acts like a highly organized, specialized team of experts. It uses three main tricks to solve the problems above.

1. The Multiscale Mamba (The "Zoom Lens")

Think of standard AI models as having a camera with a fixed zoom. They can see the big picture or the close-up, but not both well at the same time.
STM3 uses a Multiscale Mamba. Imagine a camera that can instantly switch between wide-angle, medium, and telephoto lenses without stopping. It looks at the traffic data through different "lenses" (time scales) simultaneously.

How it works: It doesn't just look at the raw data; it creates multiple versions of the data, each smoothed out to highlight different time patterns (like daily rush hours vs. minute-by-minute fluctuations). It processes these layers efficiently, ensuring it doesn't get overwhelmed by the amount of data.

2. The Adaptive Graph Causal Network (The "Smart Traffic Controller")

In a city, a traffic jam on a highway might cause a ripple effect on a side street, but the side street doesn't usually cause the highway jam. The flow of information has a direction.
STM3 uses an Adaptive Graph Causal Network.

The Analogy: Imagine a traffic controller who doesn't just look at a static map. Instead, they learn the map as they go. Crucially, they enforce a rule: "Information can flow from the big picture (coarse scale) down to the details (fine scale), but not the other way around."
Why it matters: This prevents the "noise" of a tiny, random event from confusing the model's understanding of the big, daily trends. It keeps the different time scales distinct and organized.

3. Disentangled Mixture of Experts (The "Specialized Squad")

This is the most unique part. Instead of one giant brain trying to do everything, STM3 hires a team of Experts.

The Setup: There is one "Shared Expert" who handles general traffic rules, and several "Specialized Experts" who focus on specific types of patterns (e.g., one for highway speeds, one for city center congestion).
The Routing Strategy (The "ID Card" System): In older models, the system decides which expert to use based on the current data. If the data is noisy, it might flip-flop between experts, causing instability.
- STM3's Fix: It uses a Node-Embedding system. Think of every sensor in the city having a permanent "ID Card." The system looks at the ID card (the location) to decide which expert should handle it, rather than looking at the noisy data itself. This ensures that the same location always gets the same expert, creating a stable, smooth workflow.
The "Causal Contrastive Learning" (The "Training Drill"): To make sure the experts don't all end up doing the exact same job (which would be a waste), the model uses a special training drill. It forces the experts to learn different things. It's like a coach telling the "Highway Expert," "You focus on speed," and the "City Expert," "You focus on density," and punishing them if they try to copy each other's notes.

The Results: Why It Matters

The paper tested STM3 on 10 real-world datasets, including traffic in Los Angeles (METR-LA), air quality, and solar power generation.

The Win: STM3 consistently beat the current best models.
The Highlight: On the PEMSD8 traffic dataset, it didn't just win; it dominated. It improved prediction accuracy by 7.1% (MAE), 8.5% (RMSE), and a massive 15.9% (MAPE) compared to the second-best model.
Efficiency: Even though it uses a team of experts, it is designed to be fast. It only "wakes up" the specific expert needed for a job, keeping it efficient enough for real-time use.

Summary

STM3 is like upgrading from a single, tired generalist trying to manage a chaotic city to a well-organized, specialized task force.

It uses zoom lenses to see patterns at all speeds.
It uses a smart controller to ensure information flows in the right direction without mixing up the scales.
It uses stable ID cards to assign the right specialist to the right job, ensuring everyone does their unique part without stepping on each other's toes.

The result is a system that can look far into the future and predict complex, long-term patterns with much higher accuracy than anything before it.

Technical Summary: STM3: Mixture of Multiscale Mamba for Long-Term Spatio-Temporal Time-Series Prediction

Problem Definition

Spatio-temporal time-series prediction is critical for applications such as traffic forecasting and environmental monitoring. While deep learning methods have advanced this field, existing approaches struggle to efficiently learn complex long-term spatio-temporal dependencies. The paper identifies two specific challenges inherent to long-term sequences:

Multiscale Information Extraction: Long-term sequences naturally contain multiscale information (e.g., minute-level patterns vs. daily periodicities). Existing methods often rely on coarse aggregation mechanisms (like naive summation) that lead to the re-entanglement of these scales, or they suffer from high computational overhead when attempting to extract them.
Heterogeneous Spatial Correlations: Multiscale temporal information across different spatial nodes is highly correlated yet heterogeneous (e.g., traffic patterns differ between highways and business districts). Current models often fail to capture the fundamental origin of spatial correlations, which the authors argue is intrinsically driven by the similarity of temporal patterns, rather than just static graph structures.

Methodology

The authors propose STM3 (Spatio-Temporal Mixture of Multiscale Mamba), a framework designed to disentangle multiscale temporal patterns and model complex spatial dependencies. The architecture consists of three core components:

1. Multiscale Preprocessing

The input time series is projected into raw features and then decomposed into multiple temporal scales ( $Q$ ) using 1D temporal convolutions with varying kernel sizes. This creates a multiscale representation where finer and coarser scales are explicitly separated before entering the main backbone.

2. Adaptive Graph Causal Convolution Network (AGCCN)

To address the computational complexity of connecting all nodes across all scales and the lack of known ground-truth spatial graphs, STM3 introduces AGCCN.

Adaptive Graph Learning: It utilizes a learnable node embedding dictionary to infer a shared graph structure across all temporal scales, avoiding domain-specific heuristics.
Causal Attention: A causal attention mechanism aggregates cross-scale information. It enforces a hierarchical rule where fine-grained features can retrieve coarser-grained information, but not vice versa. This prevents information entanglement and preserves scale distinguishability.

3. Disentangled Mixture of Multiscale Mamba (DMMM)

This is the core innovation, combining a Multiscale Mamba architecture with a Disentangled Mixture-of-Experts (DMoE) framework.

Multiscale Mamba: Unlike standard Mamba, this module employs a "scale amplification" mechanism using 1D convolutions to expand the receptive field efficiently. It uses a Multiscale SSM that shares input-aware parameters across scales but introduces learnable bias terms ( $\Omega$ ) to guide the State Space Model (SSM) in performing scale-specific pattern extraction.
Disentangled Mixture-of-Experts (DMoE): To handle heterogeneous spatial patterns, the model employs multiple experts (one shared, several specialized).
- Node-Embedding-Based Routing: Instead of using dynamic input time series for gating (which causes unstable routing), STM3 uses static learnable node embeddings to determine expert assignment. This ensures that all time steps from the same spatial node are routed to the same expert, guaranteeing routing smoothness.
- Causal Contrastive Learning: To ensure experts learn distinct patterns, the authors introduce a causal contrastive loss. This loss treats coarser-scale outputs of the same sample as positive pairs and finer-scale outputs or outputs from other experts as negative pairs. This asymmetric similarity measure enforces pattern disentanglement, ensuring each expert specializes in specific temporal dynamics.

Key Contributions

The paper claims three primary contributions:

Theoretical Insight: The authors uncover that inter-node spatial correlations are intrinsically driven by the similarity of multiscale temporal patterns. Based on this, they propose a method for multiscale feature aggregation and model decoupling.
Architectural Innovation (STM3): A comprehensive solution featuring:
- Multiscale Mamba: Efficiently captures multiscale temporal patterns with low computational overhead via scale amplification and bias-guided SSMs.
- AGCCN: Performs causal scale-wise spatio-temporal fusion, ensuring hierarchical information flow.
- DMoE: Achieves expert disentanglement via node-embedding routing and causal contrastive learning.
Empirical Performance: Extensive experiments on 10 real-world benchmarks demonstrate that STM3 achieves State-of-the-Art (SOTA) results in long-term spatio-temporal prediction.

Experimental Results

The model was evaluated on 10 datasets across transportation (METR-LA, PEMSD4, PEMSD8), air quality (KnowAir), wireless traffic (Milan), solar power (NREL), and electricity (ETTh1, Electricity).

Performance: STM3 consistently outperformed 13 baselines, including Mamba-based models (DST-Mamba, SST, MixMamba) and Transformer-based models (iTransformer, PatchTST).
Notable Gains: On the PEMSD8 dataset, STM3 surpassed the second-best model by 7.1% in MAE, 8.5% in RMSE, and 15.9% in MAPE. On NREL, it improved MAE by 10.1% and RMSE by 7.1%.
Efficiency: While training costs are higher due to the multi-expert design, the Top-1 MoE mechanism ensures efficient inference. STM3 is significantly faster than complex baselines like MixMamba and AMD during inference.
Ablation Studies: Removing any core component (Multiscale mechanism, DMoE, Causal Attention, or the specific routing strategy) resulted in performance degradation, confirming the necessity of each module. The multiscale mechanism provided the largest average performance gain (9.72%).

Significance and Claims

The paper positions STM3 as a significant advancement in spatio-temporal modeling by addressing the "entanglement" of spatial and temporal dependencies.

Decoupling: It successfully decouples complex multiscale information, preventing the re-entanglement seen in previous aggregation methods.
Stability: The proposed node-embedding-based routing strategy is theoretically proven to achieve zero node-wise disagreement, solving the instability issues common in traditional input-based MoE gating for time series.
Generalizability: By demonstrating superior performance across diverse domains (traffic, energy, environment), the authors claim STM3 offers a robust framework for long-term prediction tasks where fine-grained, multi-scale dynamics are critical.

The work concludes that modeling long-term spatio-temporal dependencies requires a unified approach that respects the intrinsic link between spatial node similarity and temporal pattern similarity, a principle effectively realized through the STM3 architecture.