Addressing Spatial Indistinguishability in Spatiotemporal Prediction via Optimal Transport-Guided Masking

Guangyu Wang; Jiawei Tong

cs.LGSep 28, 2026

Addressing Spatial Indistinguishability in Spatiotemporal Prediction via Optimal Transport-Guided Masking

Authors: Guangyu Wang, Jiawei Tong

Organizations: Data Science and Artificial Intelligence, Dongbei University of Finance and Economics, Dalian, Liaoning, China · Institute of Systems, Molecular & Integrative Biology, University of Liverpool, Liverpool, United Kingdom

Abstract

Spatiotemporal prediction aims to learn discriminative representations from correlated temporal signals over spatial structures for accurate future inference. A central challenge is \emph{spatial indistinguishability}: different nodes may share similar historical patterns yet evolve toward divergent futures, severely degrading forecasting performance in real-world sensor networks. Existing embedding-based and graph neural network (GNN)-based approaches can partially detect such ambiguous nodes but rely on historical similarity, struggling to capture \emph{future behavioral divergence}. We propose \textbf{STOT} (\textbf{S}patio\textbf{T}emporal \textbf{O}ptimal \textbf{T}ransport), a self-supervised framework that resolves spatiotemporal ambiguity via structured masking guided by optimal transport. Our key idea treats indistinguishability as a \emph{disambiguation} problem: future states are inferred by exploiting concurrent spatial correlations and their time-varying similarity. We design a similarity-aware metric for dynamic inter-node relationships and an optimal transport-based masking strategy to emphasize ambiguous positions during pre-training. A batch consistency constraint preserves semantic coherence, while a random-walk masking mechanism promotes structured context exploration. Experiments on six real-world datasets show that STOT performs competitively with state-of-the-art baselines on the evaluated benchmarks and improved interpretability through transport-plan visualizations.

Figures & tables

Figure 1: Architecture of STOT: (a) the embedding layer aggregates raw long-term traffic data with temporal and spatial positional features; (b) the encoder learns spatially indistinguishable node representations through masked reconstruction; (c) the decoder reconstructs masked positions from SI-enhanced representations; (d) mask generation uses SI-Mask, Batch Consistency, and Random Walk, where SI-Mask guides spatial indistinguishability and the latter two constrain representation continuity and richness; (e) patch embedding segments spatiotemporal data into local-pattern units; and (f) the STOT-enhanced downstream predictor can be integrated with various downstream architectures.

Datasets	#Sensors	#Edges	#Time Period	#Type	#Place
PeMS03	358	574	2018/09–2018/11	Flow	North Central California, USA
PeMS04	307	340	2018/01–2018/02	Flow	San Francisco Bay Area, USA
PeMS07	883	866	2017/05–2017/08	Flow	Los Angeles area, USA
PeMS08	170	295	2016/07–2016/08	Flow	San Bernardino area, USA
PeMS-BAY	325	2369	2017/01–2017/06	Speed	San Francisco Bay Area, USA
METR-LA	207	1515	2012/03–2012/06	Speed	Los Angeles area, USA

Table 1: Dataset description and statistics.

Setting	STOT	STD-MAE	STEP	ST-WA	STID	ST-Norm
Learning Rate	0.001	0.001	0.001	0.001	0.002	0.002
Hidden Dimension	96	96	96	128	32	–
Number of Heads	4	4	4	–	–	–
Similarity Thresholds $(\epsilon_{u},\epsilon_{l})$	(0.9, 0.5)	–	–	–	–	–
Mask Ratio	0.25	0.25	0.75	–	–	–
Sinkhorn Temperature $\mathbf{\epsilon}$	0.8	–	–	–	–	–

Table 2: Detailed configurations of STOT and main baselines.

Figure 2: KDE distributions of node-sample forecasting MAE grouped by the spatial indistinguishability score in Equation 5. Dashed lines indicate the group-wise mean MAE.

Model	PeMS03			PeMS04			PeMS07
Model	MAE	RMSE	MAPE	MAE	RMSE	MAPE	MAE	RMSE	MAPE
ARIMA [ 41 ]	35.31	47.59	33.78	33.73	48.80	24.18	38.17	59.27	19.46
VAR [ 6 ]	23.65	38.26	24.51	23.75	36.66	18.09	75.63	115.24	32.22
DCRNN [ 26 ]	18.18	30.31	18.91	24.70	38.12	17.12	25.30	38.58	11.66
STGCN [ 44 ]	17.49	30.12	17.15	22.70	35.55	14.59	25.38	38.78	11.08
ASTGCN [ 16 ]	17.69	29.66	19.40	22.93	35.22	16.56	28.05	42.57	13.92

Table 3: Performance comparison on PeMS03, PeMS04, PeMS07, and PeMS08 benchmarks.

Model	PeMS08			METR-LA			PeMS-BAY
Model	MAE	RMSE	MAPE	MAE	RMSE	MAPE	MAE	RMSE	MAPE
ARIMA [ 41 ]	31.09	44.32	22.73	5.15	10.45	12.70	2.33	4.76	5.40
VAR [ 6 ]	23.46	36.33	15.42	5.28	9.06	12.50	2.24	4.96	4.83
DCRNN [ 26 ]	17.86	27.83	11.45	3.59	7.61	10.44	1.96	4.59	4.68
STGCN [ 44 ]	18.02	27.83	11.40	3.60	7.50	10.56	1.99	4.51	4.66
ASTGCN [ 16 ]	18.61	28.16	13.08	3.57	7.19	10.32	1.86	4.07	4.27

Table 4: Performance comparison on PeMS08, METR-LA, and PeMS-BAY benchmarks.

Figure 3: Multi-step prediction comparison on PeMS03, PeMS04, PeMS07, and PeMS08.

Figure 4: Mask pre-training component ablation on PeMS04 and PeMS08.

Figure 5: STOT reconstruction visualization on four datasets during pre-training.

Figure 6: STOT forecasting visualization on four datasets after pre-training.

Figure 7: Mask pre-training component ablation on PeMS03 and PeMS07.

Datasets	#Sensors	#Edges	#Training	#Validation	#Testing
PeMS04 (IN)	307	340	[10173, 12, 307,1]	[3375, 12, 307,1]	[3375, 12, 307,1]
PeMS04 (OUT)	307	340	[10173, 12, 307,1]	[3375, 12, 307,1]	[3375, 12, 307,1]
PeMS04-s (IN)	161	188	[10173, 12, 161,1]	[3375, 12, 161,1]	[3375, 12, 161,1]
PeMS04-s (OUT)	161	188	[10173, 12, 161,1]	[3375, 12, 161,1]	[3375, 12, 161,1]

Table 5: Dataset segmentation details.

Model	Metrics	#1 step	#3 steps	#6 steps	#12 steps
STID	MAE	16.88±0.03	18.12±0.02	19.12±0.00	21.10±0.25
	RMSE	27.29±0.05	29.31±0.03	30.82±0.01	33.32±0.27
	MAPE (%)	11.19±0.25	12.07±0.23	12.54±0.11	13.76±0.05
ST-Norm	MAE	16.26±0.02	18.23±0.03	19.58±0.07	21.66±0.14
	RMSE	26.20±0.02	30.00±0.03	32.38±0.07	35.37±0.17
	MAPE (%)	7.18±0.21	8.01±0.23	8.41±0.27	9.55±0.32

Table 6: High spatial indistinguishability prediction on PeMS04.

Figure 8: STOT parameter sensitivity on PeMS04 under MAE, RMSE, and MAPE.

Figure 9: Time cost of different approaches.

Dataset	Nodes	Mask generation	Sinkhorn
PeMS08	170	$1.3466\pm 0.0852$	$0.2138\pm 0.0036$
METR-LA	207	$1.4644\pm 0.1238$	$0.2150\pm 0.0035$
PeMS04	307	$1.3837\pm 0.0891$	$0.2137\pm 0.0009$
PeMS-BAY	325	$1.3890\pm 0.0963$	$0.2149\pm 0.0034$
PeMS03	358	$1.4294\pm 0.0556$	$0.2154\pm 0.0029$
PeMS07	883	$5.0419\pm 0.3514$	$0.2157\pm 0.0038$

Table 7: Mask generation time and the Sinkhorn substep time, in milliseconds per batch, across datasets with different numbers of nodes.

Figure 10: Visualization of OT and batch consistency: (a) and (b) cost matrices before and after batch consistency for the first 72 timesteps (6 patches), (c) and (d) the final cost matrix and OT matrix, (e) transport-matrix heatmap, and (f) batch-mask correlation matrix showing the effect of batch consistency on mask generation.

Figure 11: Consistency, diversity, and balance scores as functions of the batch consistency weight $\lambda$ on the six datasets. The dashed line marks the default $\lambda=0.7$ .

Figure 12: Comparison of random walk mechanism variants on PeMS04.

Explore similar work

Dimensional Balance Improves Large Scale Spatiotemporal Prediction Performance
May 11, 2026Jing Chen, Shixiang Pan, Yujie Fan +3Spatiotemporal Prediction Spatiotemporal
School of Computer Science and Technology, Hangzhou Dianzi University, Hangzhou, 310018, China · Key Laboratory of New Industrial Internet Control Technology, Hangzhou, 310018, China · College of Economics, China Jiliang University, Hangzhou, 314423, China
UniST-Pred: A Robust Unified Framework for Spatio-Temporal Traffic Forecasting in Transportation Networks Under Disruptions
Feb 15, 2026Yue Wang, Areg Karapetyan, Djellel Difallah +1Accurate Traffic Forecasting Spatio-Temporal Forecasting
Efficient Traffic State Prediction With Dynamic Joint Spatio-Temporal Relation Inference
Apr 10, 2025Zhifeng Hao, Kai Hu, Juncai Zhang +2Accurate Traffic Forecasting Spatial Encoding
School of Cyberspace Security, Hainan University, Haikou, China

May 11, 2026, cs.LGJ/K move · Enter open · S save

Dimensional Balance Improves Large Scale Spatiotemporal Prediction Performance

Jing Chen, Shixiang Pan, Yujie Fan, Haocheng Ye+2

School of Computer Science and Technology, Hangzhou Dianzi University, Hangzhou, 310018, China · Key Laboratory of New Industrial Internet Control Technology, Hangzhou, 310018, China · College of Economics, China Jiliang University, Hangzhou, 314423, China

Spatiotemporal Prediction Spatiotemporal Epidemiological Models Intrinsic Dimensionality Balance

Open paper AskPDF arXiv

Accurate spatiotemporal pattern analysis is critical in fields such as urban traffic, meteorology, and public health monitoring. However, existing methods face performance bottlenecks, typically yielding only incremental gains and often exhibiting limited cross-domain transferability. We analyze this bottleneck through spatial and temporal entropy measures, which are used as diagnostic indicators of spatiotemporal complexity mismatch rather than as guarantees that entropy alignment alone yields better forecasting. Empirically, larger mismatch is often accompanied by higher prediction uncertainty, especially under a fixed model-capacity budget. Guided by this diagnostic, we propose a scalable, adaptive framework that harmonizes spatial and temporal feature representations. Spatial dimensionality is compressed via low-rank matrix embedding to preserve essential structure, while an extended temporal horizon captures long-range dependencies and mitigates cumulative errors arising from temporal heterogeneity. Extensive experiments on urban traffic, meteorological, and epidemic datasets demonstrate substantial accuracy gains and broad applicability across the evaluated domains, suggesting that the framework is promising for a wide range of spatiotemporal tasks beyond the current study. The code is available on GitHub at https://github.com/ST-Balance/ST-Balance.

Abstract

Explore similar work

Dimensional Balance Improves Large Scale Spatiotemporal Prediction Performance

UniST-Pred: A Robust Unified Framework for Spatio-Temporal Traffic Forecasting in Transportation Networks Under Disruptions

Efficient Traffic State Prediction With Dynamic Joint Spatio-Temporal Relation Inference