UniST-Pred: Unified Traffic Forecasting

• The paper introduces a decouple-then-fuse framework that separately models temporal dynamics and spatial structures to enhance modularity and robustness.
• It employs TSMixer-style temporal mixing and GTN-style adaptive graph construction to effectively capture long-range dependencies and network changes.
• Extensive evaluations on benchmark datasets show that UniST-Pred achieves parameter efficiency and maintains low error rates even under severe bridge-removal scenarios.

Searching arXiv for the specified paper and closely related UniST variants to ground the article in the provided literature. UniST-Pred is a unified spatio-temporal traffic forecasting framework for transportation networks under disruptions, introduced to address forecasting in settings where the network may be partially disconnected, structurally changed, or observed under uncertainty (Wang et al., 15 Feb 2026). The framework is defined by a modular decouple-then-fuse design: it first separates temporal modeling from spatial representation learning, then combines both through adaptive representation-level fusion. In the paper’s formulation, this design is motivated by the limitations of tightly coupled spatio-temporal architectures, which often achieve strong short-term predictive accuracy at the cost of increased complexity, limited modularity, and potential brittleness when topology shifts (Wang et al., 15 Feb 2026).

1. Problem formulation and scope

UniST-Pred is formulated on a weighted directed road graph

$$\mathcal{G} = (\mathcal{V}, \mathcal{E}, \mathbf{A}),$$

where $\mathcal{V} = \{v_1,\dots,v_N\}$ denotes road segments or intersections, $$\mathcal{E} \subseteq \mathcal{V} \times \mathcal{V}$$ denotes directed connectivity, and $\mathbf{A}\in\mathbb{R}^{N\times N}$ is the weighted adjacency matrix (Wang et al., 15 Feb 2026). Each node $v_i$ has a time series

$\mathbf{x}_i=[x_i^1,x_i^2,\dots,x_i^T],$

and the network state at time $t$ is

$$\mathbf{X}^t=[x_1^t,\dots,x_N^t]^\top\in\mathbb{R}^N.$$

Nodes may also have static features

$\mathbf{X}'\in\mathbb{R}^{N\times d},$

such as lane count or capacity (Wang et al., 15 Feb 2026).

Given a history window of length $H$, the forecasting task is to predict the next $\mathcal{V} = \{v_1,\dots,v_N\}$0 steps through

$\mathcal{V} = \{v_1,\dots,v_N\}$1

The intended dependencies are both temporal and spatial: rush-hour structure, congestion buildup, and lagged effects on the one hand, and interactions among connected roads, rerouting, and corridor spillovers on the other (Wang et al., 15 Feb 2026).

A defining aspect of the framework is its emphasis on robustness under structural and observational uncertainty. The paper explicitly situates the method in scenarios such as bridge closures, flooded segments, capacity reductions, incidents, and disconnected roads. This suggests a broader design target than conventional benchmark optimization: accurate forecasting precisely when network structure changes and prediction is most operationally consequential (Wang et al., 15 Feb 2026).

2. Architectural principle: decoupling temporal and spatial learning

UniST-Pred is organized into three explicit stages: temporal modeling, spatial representation learning, and adaptive fusion (Wang et al., 15 Feb 2026). Its central claim is that temporal dynamics and spatial structure need not be interleaved layer by layer. Instead, the model computes a temporal representation and a spatial representation independently, and only then combines them.

The temporal branch produces

$\mathcal{V} = \{v_1,\dots,v_N\}$2

while the spatial branch produces

$\mathcal{V} = \{v_1,\dots,v_N\}$3

This separation is the paper’s main methodological contrast with standard spatio-temporal graph neural networks and graph-transformer-style models, where temporal and graph propagation operators are repeatedly intertwined through the network depth (Wang et al., 15 Feb 2026).

The paper attributes three consequences to this decoupling. First, it improves modularity, because each branch can be modified or analyzed independently. Second, it may improve robustness, because temporal dynamics are not forced to depend on a fixed topology at every layer. Third, it improves interpretability, because temporal and structural information remain explicitly separable in the learned representation (Wang et al., 15 Feb 2026).

A plausible implication is that UniST-Pred treats topology as one informative input among several, rather than the organizing principle of all hidden-state evolution. That interpretation is consistent with the framework’s emphasis on changing or degraded network structure.

3. Temporal branch, spatial branch, and fusion mechanism

The temporal branch uses TSMixer-style feature/time mixing rather than RNNs or attention (Wang et al., 15 Feb 2026). The paper defines learnable mixing operators $\mathcal{V} = \{v_1,\dots,v_N\}$4: $\mathcal{V} = \{v_1,\dots,v_N\}$5 mixes along the temporal dimension independently for each node and feature, while $\mathcal{V} = \{v_1,\dots,v_N\}$6 mixes across feature channels at each time step. A stack of $\mathcal{V} = \{v_1,\dots,v_N\}$7 temporal mixer blocks is defined as

$\mathcal{V} = \{v_1,\dots,v_N\}$8

with

$\mathcal{V} = \{v_1,\dots,v_N\}$9

The resulting temporal representation is

$$\mathcal{E} \subseteq \mathcal{V} \times \mathcal{V}$$0

The paper characterizes this branch as capturing long-range temporal dependencies efficiently and independently of graph propagation (Wang et al., 15 Feb 2026).

The spatial branch uses GTN-style task-adaptive graph construction followed by lightweight graph convolution (Wang et al., 15 Feb 2026). For each channel $$\mathcal{E} \subseteq \mathcal{V} \times \mathcal{V}$$1, an adjacency matrix is generated by soft composition of two selected adjacency matrices: $$\mathcal{E} \subseteq \mathcal{V} \times \mathcal{V}$$2 After adding self-loops,

$$\mathcal{E} \subseteq \mathcal{V} \times \mathcal{V}$$3

the model computes channel-wise graph propagation as

$$\mathcal{E} \subseteq \mathcal{V} \times \mathcal{V}$$4

concatenates the channel outputs,

$$\mathcal{E} \subseteq \mathcal{V} \times \mathcal{V}$$5

and projects them to a spatial embedding

$$\mathcal{E} \subseteq \mathcal{V} \times \mathcal{V}$$6

This branch uses static node features such as lane count or capacity, and its task-adaptive graph generation is intended to reduce reliance on a single fixed adjacency (Wang et al., 15 Feb 2026).

The fusion stage combines $$\mathcal{E} \subseteq \mathcal{V} \times \mathcal{V}$$7 and $$\mathcal{E} \subseteq \mathcal{V} \times \mathcal{V}$$8 through a squeeze-and-excitation residual integration block. First,

$$\mathcal{E} \subseteq \mathcal{V} \times \mathcal{V}$$9

The final fused representation is

$\mathbf{A}\in\mathbb{R}^{N\times N}$0

where, letting $\mathbf{A}\in\mathbb{R}^{N\times N}$1,

$\mathbf{A}\in\mathbb{R}^{N\times N}$2

The paper interprets this as adaptive channel reweighting over a residual transformation, allowing the model to decide how much to trust spatial and temporal information for a given prediction (Wang et al., 15 Feb 2026).

4. End-to-end forecasting objective and implementation settings

The paper writes the final prediction as

$\mathbf{A}\in\mathbb{R}^{N\times N}$3

and trains the model end to end by minimizing

$\mathbf{A}\in\mathbb{R}^{N\times N}$4

where $\mathbf{A}\in\mathbb{R}^{N\times N}$5 includes the parameters of the temporal branch, spatial branch, residual mapping, and squeeze-and-excitation fusion (Wang et al., 15 Feb 2026).

The implementation uses task-dependent loss functions. For SimSF-Bay, the paper specifies MAE loss; for PEMS-Bay, Smoothed L1 loss; and for NYCTaxi, MAE loss (Wang et al., 15 Feb 2026). Evaluation uses RMSE, MAE, and MAPE, with

$\mathbf{A}\in\mathbb{R}^{N\times N}$6

and

$\mathbf{A}\in\mathbb{R}^{N\times N}$7

A mask filters NaNs and very small true values (Wang et al., 15 Feb 2026).

The paper also reports concrete hyperparameter choices. The temporal branch uses $\mathbf{A}\in\mathbb{R}^{N\times N}$8 TSMixer layers for SimSF-Bay and NYCTaxi, and $\mathbf{A}\in\mathbb{R}^{N\times N}$9 for PEMS-Bay. Feature-mixing MLP width is 100 and dropout is 0.2. The spatial branch uses $v_i$0 channels, $v_i$1 GT layers, and GCN out channel 100; GCN in channel is 2 for SimSF-Bay and PEMS-Bay, and 1 for NYCTaxi. The squeeze-and-excitation reduction ratio is $v_i$2 for SimSF-Bay and NYCTaxi, and $v_i$3 for PEMS-Bay (Wang et al., 15 Feb 2026).

5. Datasets, disruption scenarios, and empirical performance

UniST-Pred is evaluated on two standard benchmark datasets and one disruption-oriented simulation dataset (Wang et al., 15 Feb 2026).

On PEMS-Bay, the dataset contains 325 sensors, 2369 edges, 5-minute intervals, and a fixed sensor graph, with data spanning Jan 1, 2017 to May 31, 2017. The input is the previous 7 days around the prediction time, corresponding to

$v_i$4

with a train/validation/test split of $v_i$5 (Wang et al., 15 Feb 2026).

On NYCTaxi, the dataset is a taxi trajectory-based traffic flow dataset for New York City with 200 nodes and 712 edges, at 30-minute intervals, using the previous 17.5 hours as input and the next 1 step as target: $v_i$6 The train/validation/test split is again $v_i$7 (Wang et al., 15 Feb 2026).

The distinctive benchmark is SimSF-Bay, a MATSim-generated microscopic agent-based traffic simulation dataset (Wang et al., 15 Feb 2026). It contains 7,709 road segments, 8,781 edges, 5-minute intervals, and 168 time steps per scenario, covering 6 AM to 8 PM. The traffic variable is flow, and the network topology varies across scenarios. The MATSim model includes a sample population of 463,938 commuters, 1454 Traffic Analysis Zones, and road capacities scaled to 8% of original to match sample population scale (Wang et al., 15 Feb 2026). Training uses 5 different scenarios with 8 random seeds; test uses 8 different scenarios with a fixed random seed. The paper emphasizes that train and test involve different scenarios, making the benchmark a test of structural variation rather than interpolation within a fixed environment (Wang et al., 15 Feb 2026).

The core robustness evaluation uses three severe bridge-removal scenarios: removal of the Oakland Bay Bridge; removal of the Carquinez Bridge and San Rafael Bridge; and removal of the Golden Gate Bridge and Dumbarton Bridge (Wang et al., 15 Feb 2026). In the scenario RMSE table, UniST-Pred obtains the lowest RMSE in every setting:

Setting	RMSE
Scenario 1	3.61
Scenario 2	3.59
Scenario 3	3.60
Fully connected	3.62

This is the paper’s strongest evidence that the framework remains effective under severe topological changes (Wang et al., 15 Feb 2026).

On the full benchmarks, UniST-Pred reports the following headline results (Wang et al., 15 Feb 2026):

Dataset	RMSE	MAE	MAPE
SimSF-Bay	3.6	2.37	34.2
PEMS-Bay	4.20	1.91	4.52
NYCTaxi	13.39	4.70	41.6

For SimSF-Bay, the method is best on all three metrics. For PEMS-Bay, it ties STEP in RMSE and is second-best in MAE and MAPE. For NYCTaxi, it is best in RMSE and MAE, while STEP achieves lower MAPE (Wang et al., 15 Feb 2026). The paper also reports that UniST-Pred is especially strong on high-flow locations, which it describes as operationally more important, whereas STEP is stronger on low-flow locations on PEMS-Bay and NYCTaxi (Wang et al., 15 Feb 2026).

6. Ablations, interpretability, and limitations

The ablation study on SimSF-Bay identifies all three components as important (Wang et al., 15 Feb 2026). The full model obtains RMSE 3.61, MAE 2.37, and MAPE 34.2. Removing the spatial branch degrades performance to RMSE 6.06, MAE 3.43, and MAPE 65.5. Removing the temporal branch yields RMSE 5.64, MAE 4.11, and MAPE 49.8. Removing fusion yields RMSE 4.06, MAE 2.57, and MAPE 36.2. Replacing the spatial learner with a standard GCN produces RMSE 4.17, MAE 2.65, and MAPE 71.2, while replacing the temporal mixer with a fully connected temporal block gives RMSE 3.96, MAE 2.58, and MAPE 38.1 (Wang et al., 15 Feb 2026). These numbers support the paper’s claim that the GTN-style spatial learner, TSMixer-style temporal modeling, and adaptive fusion are synergistic rather than interchangeable.

Interpretability is presented as both an architectural property and an empirical one. Because temporal and spatial processing are separated, the framework allows temporal embeddings to be read as intrinsic traffic evolution and spatial embeddings as structural context, with fusion making their relative importance explicit (Wang et al., 15 Feb 2026). To examine this under disruptions, the paper applies Integrated Gradients: $v_i$8 The reported road-level attribution maps show that, under bridge removal, the model reduces reliance on disconnected links and shifts importance toward alternative corridors (Wang et al., 15 Feb 2026). This suggests that the learned representation adapts to changed connectivity rather than rigidly encoding the original network.

The efficiency analysis is more nuanced than the term “lightweight” might imply. UniST-Pred is substantially smaller than STEP: 3,208,208 parameters versus 12,952,470 on SimSF-Bay, 16,558,467 versus 61,138,261 on PEMS-Bay, and 167,994 versus 5,670,166 on NYCTaxi (Wang et al., 15 Feb 2026). On SimSF-Bay, its training time is 21.5 hr, much lower than STEP’s 112 hr, STGCN’s 95 hr, ASTGCN’s 110 hr, ST-SSL’s 950 hr, and DCRNN’s 280 hr. However, its inference time is listed as 20 min, slower than STEP’s 4 min and much slower than TSMixer’s 0.05 min (Wang et al., 15 Feb 2026). The paper therefore supports “lightweight” more strongly in terms of parameter efficiency and architectural simplicity than in raw inference latency.

Several limitations are explicit or directly evidenced in the results (Wang et al., 15 Feb 2026). UniST-Pred is not always best on MAPE, especially on PEMS-Bay and NYCTaxi. The disruption case study indicates that short-lived spikes at individual roads remain difficult to capture exactly. The paper also leaves some implementation internals less explicit than ideal, especially for tensor dimensions and residual block details. In addition, the spatial branch relies on static node features $v_i$9, which may constrain expressiveness relative to models that repeatedly propagate dynamic traffic features across the graph. These caveats do not negate the paper’s main conclusion, but they delimit the circumstances under which the framework’s decoupled design is most advantageous.

From a broader research perspective, UniST-Pred occupies a different position from the earlier universal urban spatio-temporal prediction model “UniST” (Yuan et al., 2024). Whereas that earlier work pursued a one-for-all pretraining-and-prompting paradigm across more than 20 urban scenarios, UniST-Pred focuses specifically on transportation networks under disruptions and emphasizes modular decoupling, adaptive graph construction, and interpretable robustness (Yuan et al., 2024, Wang et al., 15 Feb 2026). This suggests two distinct strands of “unified” spatio-temporal modeling in recent arXiv literature: universal pretraining across domains, and robust decouple-then-fuse forecasting under structural change.