FaST: Efficient and Effective Long-Horizon Forecasting for Large-Scale Spatial-Temporal Graphs via Mixture-of-Experts

AI summary

Overview

Research area: Machine learning for spatial-temporal graph (STG) forecasting, specifically long-horizon prediction on very large urban sensor networks.

Technical level: Advanced. The paper assumes familiarity with graph neural networks, transformer attention mechanisms, Mixture-of-Experts architectures, and the complexity analysis conventions used in deep learning systems research.

Scope: The paper introduces FaST, a Mixture-of-Experts framework that enables one-week-ahead traffic forecasting (672 steps at 15-minute granularity) on graphs with thousands of nodes while keeping computation and memory linear in the number of nodes.

What This Paper Is About

Existing STG forecasting models work well for short horizons and small graphs, but their cost grows quadratically with both the number of nodes and the sequence length, so they run out of GPU memory when asked to predict days ahead on networks of thousands of sensors. Naive efficiency fixes (neighbor sampling, linear attention, aggressive sequence compression) buy speed but discard long-range spatial dependencies or homogenize temporal patterns that differ across sensors and times of day. FaST's goal is to reach one-week-ahead forecasting on graphs with up to 8,600 nodes without giving up either accuracy or the ability to model heterogeneous spatial-temporal behavior.

Key Contributions

1. A heterogeneous parallelized Mixture-of-Experts framework (FaST) — the first approach, per the authors, to combine long-horizon (672-step) and large-scale (tens of thousands of nodes) STG forecasting within tolerable training time.

2. A Heterogeneity-aware Router (HA-Router) and parallel GLU-MoE module — the router injects adaptive spatial and temporal expert biases so different nodes and time periods select different experts, avoiding expert weight polarization without auxiliary load-balancing losses, while GLU experts replace feed-forward networks and are computed as a single fused linear layer for parallel execution.

3. Linear complexity via two modules — an MoE-based temporal compression input that projects the raw historical sequence into a compact embedding, and an Adaptive Graph Agent Attention (AGA-Att) mechanism that routes spatial information through a small set of learned agent tokens, reducing spatial interaction cost from O(N²) to O(N·a) with a ≪ N.

4. Empirical validation on a large-scale benchmark — superior accuracy across all metrics, horizons, and datasets compared to state-of-the-art baselines, with 4.4%–18.4% error reduction and 1.3×–2.2× faster inference on the LargeST benchmark.

Main Findings

• Accuracy gains are consistent across horizons and datasets: FaST outperforms all baselines on MAE, RMSE, MAPE, and R² at 48, 96, 192, and 672 steps. On San Diego it improves MAE by 4.73% and MAPE by 8.35% on average over the best baseline; on California, MAE/MAPE gains range from 3.67% to 10.06%. It won the best result on 16 separate prediction tasks.

• Scalability is the decisive advantage: STGCN and STPGNN cannot finish training on the GLA and CA datasets within reasonable time, and RPMixer hits memory limits on GLA (192-step) and CA (96-step). FaST trains on CA (8,600 nodes, 672-step horizon) with a single fixed configuration, while representative STGNNs exhaust 48 GB of GPU memory on the same setting.

• GLU experts match FFN experts at lower cost: Replacing feed-forward experts with Gated Linear Units produces essentially identical accuracy but delivers roughly a 1.4× wall-clock speed-up and substantially lower GPU memory (e.g., on CA, 17.2 GB vs 23.2 GB, and 158 s vs 225 s per training epoch).

• Every component contributes: Ablations show the full model beats all variants. Removing HA-MoE causes the largest degradation (most visible on GBA); replacing the MoE input layer with a plain linear layer hurts notably on SD and GLA; removing the HA-Router or AGA-Att also degrades performance measurably.

• Optimal hyperparameters depend on graph size: Eight experts is optimal for CA and SD, 32 for GBA and GLA; three backbone layers suit CA and SD, while GBA and GLA benefit from four or five. Embedding dimensions beyond 64 yield diminishing returns. Sixteen to 32 agent tokens provide a good accuracy–cost trade-off.

• Temporal-centric baselines remain weak on spatial structure: DLinear, NHITS, and CycleNet capture non-stationary and periodic patterns but cannot exploit spatial correlation, and they consistently trail graph-based methods where those methods are trainable.

Methodology in Plain English

FaST borrows the stacked residual block layout of a transformer but swaps its two main components for cheaper alternatives.

Input compression. Instead of feeding a full day of 96 historical time steps per node into the network, FaST compresses that history into a compact dense vector. Crucially, it does not use one compression rule for everyone. A router inspects the raw series and each node's spatial and temporal identity (which sensor, which time of day, which day of week) and decides how much weight to give each of several compression experts. Sensors with different behavior patterns therefore get different compressed representations rather than a single blended average. Trainable spatial, daily, and weekly embeddings are added so the model knows where and when each value came from.

Spatial mixing without all-pairs attention. Standard self-attention compares every node against every other node, which is what makes large graphs explode in cost. FaST instead introduces a small set of learnable "agent" tokens. In a first step, each agent queries all real nodes and summarizes long-range information into itself. In a second step, each node queries the agents and receives a redistributed summary. This is a low-rank bottleneck: information still travels between distant nodes, but it passes through a compressed set of intermediaries, cutting the cost from N² to N times the number of agents.

Parallel experts. Each block replaces the transformer's feed-forward network with a Mixture-of-Experts layer. Instead of running each expert separately and hoping the hardware scheduler parallelizes them, FaST concatenates all expert weight matrices into one large linear layer, splits the output into two halves, and applies the gating nonlinearity once. This produces all experts' outputs in a single matrix multiplication, which is why it runs efficiently on one GPU without custom parallel infrastructure.

Output. Features from all layers are concatenated and passed through a small MLP that directly emits all future time steps at once. Training uses Huber loss, which is less sensitive to outliers than squared error.

Why This Matters

Impact on research. The paper reframes the scalability question in STG forecasting: rather than choosing between structure-aware methods that discard long-range dependencies and structure-free methods that erase spatial semantics, it argues that heterogeneity-aware routing plus a low-rank spatial bottleneck can preserve both. If the results hold, the quadratic-complexity barrier that has limited STG research to small graphs and short horizons is substantially relaxed, opening up longer-horizon and larger-network experimental settings as a standard benchmark target.

Real-world applications.

• Electricity grid operations: A seven-day demand outlook lets operators schedule renewable generation, storage discharge, and maintenance windows ahead of peak loads rather than reacting to them.
• Urban traffic management: Week-ahead congestion forecasts on city-wide sensor networks support proactive signal timing, incident staging, and traveler advisories.
• Public transit and logistics planning: Multi-day ridership or fleet-demand forecasts on networks with thousands of stops inform vehicle allocation and crew scheduling.
• Environmental and utility monitoring: Long-horizon forecasting of air quality, water distribution, or district heating across dense sensor deployments supports preventive maintenance and regulatory reporting.

Industry relevance. The efficiency claims matter operationally. Training on a single 48 GB GPU and running inference 1.3×–2.2× faster than existing state-of-the-art models lowers the infrastructure threshold for deployment. Because FaST uses one fixed configuration across all four datasets, it also reduces the engineering burden of per-city hyperparameter tuning that typically accompanies production STG systems.

Future Directions

• Adaptive agent counts. The paper notes that when spatial redundancy is weak, more agent tokens are needed to compensate, at higher O(N·a) cost. A mechanism that grows or shrinks the agent set per layer or per graph could improve the accuracy–efficiency frontier.

• Graphs without reliable topology. FaST's agent attention avoids dependence on an explicit graph structure, but the experiments all use the LargeST benchmark. Testing on domains where no road network or sensor adjacency exists — and comparing directly against structure-aware methods on those domains — would clarify how general the approach is.

• Beyond traffic. All four datasets are traffic flow. Validating on electricity, water, or air-quality networks, which have different periodicity and spatial correlation structures, would test whether the heterogeneity-aware routing generalizes.

• Interaction between horizon length and compression. The temporal compression module is fixed per dataset rather than tuned per prediction horizon. Studying whether the compression pathway should adapt to the target horizon — or whether the expert count should scale with it — is a natural extension.

Target Audience

This paper is most useful to graduate students and researchers working on spatial-temporal graph neural networks, efficient transformer architectures, or Mixture-of-Experts systems. Practitioners building production forecasting systems for urban infrastructure — traffic operations, utility load forecasting, sensor network monitoring — will find the scalability and single-configuration deployment arguments directly relevant. Readers without background in graph neural networks or attention mechanisms will need to consult the cited baselines (DCRNN, STGCN, GWNet, BigST) first, since the paper's framing assumes familiarity with the complexity trade-offs those models embody.

Authors’ abstract

Spatial-Temporal Graph (STG) forecasting on large-scale networks has garnered significant attention. However, existing models predominantly focus on short-horizon predictions and suffer from notorious computational costs and memory consumption when scaling to long-horizon predictions and large graphs. Targeting the above challenges, we present FaST, an effective and efficient framework based on heterogeneity-aware Mixture-of-Experts (MoEs) for long-horizon and large-scale STG forecasting, which unlocks one-week-ahead (672 steps at a 15-minute granularity) prediction with thousands of nodes. FaST is underpinned by two key innovations. First, an adaptive graph agent attention mechanism is proposed to alleviate the computational burden inherent in conventional graph convolution and self-attention modules when applied to large-scale graphs. Second, we propose a new parallel MoE module that replaces traditional feed-forward networks with Gated Linear Units (GLUs), enabling an efficient and scalable parallel structure. Extensive experiments on real-world datasets demonstrate that FaST not only delivers superior long-horizon predictive accuracy but also achieves remarkable computational efficiency compared to state-of-the-art baselines. Our source code is available at: https://github.com/yijizhao/FaST.