GAS-PFN: Graph-Aware Spatiotemporal Prior-Data Fitted Networks for Zero-Shot Traffic Forecasting

Joseph Natter; Long Vu; Yifan Zhang; Bobin Deng; Kun Suo; Rui Wu

cs.LGOct 27, 2026

GAS-PFN: Graph-Aware Spatiotemporal Prior-Data Fitted Networks for Zero-Shot Traffic Forecasting

Authors: Joseph Natter, Long Vu, Yifan Zhang, Bobin Deng, Kun Suo, Rui Wu

Organizations: Kennesaw State University · Missouri State University

Abstract

Traffic forecasting requires modeling complex temporal dynamics and the spatial dependencies imposed by road topology. Supervised graph-based models achieve strong performance, but are often trained for a particular dataset and sensor graph, making deployment to new or evolving networks costly. Prior-Data Fitted Networks (PFNs) offer a zero-shot alternative by learning in-context inference from synthetic task priors, which provide an explicit mechanism for encoding domain-specific assumptions. Existing forecasting PFNs, however, do not explicitly represent the graph-structured dependencies governing traffic dynamics. To combine graph-aware traffic modeling with zero-shot adaptation, we introduce the Graph-Aware Spatiotemporal Prior-Data Fitted Network (GAS-PFN). GAS-PFN incorporates sampled road topology into its synthetic task prior through random-walk diffusion, teaching the model to learn graph-coupled dynamics during pretraining. Its co-designed inference architecture uses two mask-conditioned passes through a shared Transformer backbone to separately model sensor-specific temporal dynamics and spatiotemporal interactions among neighboring sensors before integrating their representations for prediction. This design enables GAS-PFN to use target-network structure and context without parameter updates. Across six benchmarks at the 90-minute horizon, GAS-PFN reduces MSE by 34.0% on average relative to TimePFN, its closest PFN-based predecessor. It also reduces MSE by 27.7% on average relative to four recent industry-developed general-purpose zero-shot forecasters. In this comparison, all models receive 90 minutes of contiguous recent history, and any additional context is drawn from each dataset's chronological training split. These results indicate that aligning PFN priors and inference architectures with graph structure improves in-context learning for zero-shot traffic forecasting.

Figures & tables

Figure 1: High-level overview of GAS-PFN. The method constructs graph-structured prior tasks, assembles mixed-context memory, embeds context and query patches as tokens, applies dual mask-conditioned backbone inference, and fuses the branch outputs to produce zero-shot forecasts.

Figure 2: Overview of the proposed prior data generation process. (A) A shared temporal base is constructed by composing trend, seasonality, and noise components. (B) A spatial graph prior is sampled from a real-world road network by selecting a root node and extracting a connected subgraph. (C) The temporal base is lifted to all nodes, modulated by node-wise gains, propagated via a random-walk diffusion operator, and perturbed with graph-correlated noise to yield the final spatiotemporal process.

Figure 3: Dual mask-conditioned backbone. A shared Transformer backbone evaluates

\mathbf{Z}^{(0)}

under separate multi-channel and self-channel attention masks. Each branch shows the mask and attention biases for the selected query token; gray blocks denote masked keys. Only the resulting query-token representations,

\mathbf{U}^{\mathrm{multi}}

and

\mathbf{U}^{\mathrm{self}}

, are retained.

Dataset	Horizon	GAS-PFN		Chronos-2		TimesFM-2.5		Moirai-2		Lag-Llama		TimePFN
Dataset	Horizon	MAE	MSE	MAE	MSE	MAE	MSE	MAE	MSE	MAE	MSE	MAE	MSE
METR-LA	30 min	3.97 $\pm$ 0.05	68.31 $\pm$ 1.92	3.68	70.62	3.80	72.30	4.23	83.92	4.86	86.23	4.86	69.41
	60 min	4.83 $\pm$ 0.06	97.61 $\pm$ 2.92	4.38	91.23	4.73	108.38	5.30	124.56	6.21	130.31	5.59	94.33
	90 min	5.50 $\pm$ 0.05	119.69 $\pm$ 4.15	4.95	107.42	5.45	136.04	6.35	163.35	7.41	178.32	6.13	112.73
PEMS-BAY	30 min	1.73 $\pm$ 0.01	14.85 $\pm$ 0.09	1.61	12.90	1.65	15.01	1.74	15.53	2.60	45.83	2.41	22.35
	60 min	2.21 $\pm$ 0.01	25.22 $\pm$ 0.09	2.24	27.25	2.18	27.70	2.36	29.99	2.99	52.64	2.85	32.51

Table 1: Zero-shot forecasting performance across datasets and prediction horizons. Lower values indicate better performance. Best results are shown in bold and second-best results are underlined. GAS-PFN reports mean

\pm

standard deviation over three independently trained checkpoints; each baseline uses a single released pretrained checkpoint.

Figure 4: MAE as a function of the per-source context budget

K

at the 90-minute horizon. Each setting uses

K

recent and

K

historical support patches per node. Values are computed using a single fixed GAS-PFN checkpoint.

Figure 5: MAE of GAS-PFN and SCPT at the 60-minute horizon as the fraction of roads available for SCPT training varies. Values average three runs with matched road-split and checkpoint seeds.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Signal	Time points	Sensors	Train / val. / test rows	Test windows	Graph input
METR-LA	Speed	34,272	207	23,990 / 3,428 / 6,854	6,837	Weighted, asymmetric
PEMS-BAY	Speed	52,116	325	36,481 / 5,212 / 10,423	10,406	Weighted, asymmetric
PEMS03	Flow	26,208	358	18,345 / 2,622 / 5,241	5,224	Binary, symmetric
PEMS04	Flow	16,992	307	11,894 / 1,700 / 3,398	3,381	Binary, symmetric
PEMS07	Flow	28,224	883	19,756 / 2,824 / 5,644	5,627	Weighted cost kernel, symmetric
PEMS08	Flow	17,856	170	12,499 / 1,786 / 3,571	3,554	Binary, symmetric

Appendix

Table 1: Dataset dimensions, chronological splits, and graph inputs. The row counts refer to raw time points. Test windows use an 18-step forecast horizon and stride one.

Category	Setting	Value
Architecture	Lookback / forecast length	$18/18$
Architecture	Transformer blocks / attention heads	$10/8$
Architecture	Hidden / feed-forward dimension	$512/2048$
Architecture	Attention-head dimension	$64$
Architecture	Dropout	$0.1$
Architecture	Temporal-distance initialization	$\gamma_{h}=0.01$

Appendix

Table 2: GAS-PFN architecture, pretraining, and evaluation hyperparameters.

Model	Implementation	Recent-only regime	History-augmented regime
Chronos-2	amazon/chronos-2 ; chronos-forecasting 2.2.2	18 recent observations	8,174 training-tail $+18$ recent $=8,192$
TimesFM-2.5	google/timesfm-2.5- 200m-pytorch ; timesfm 2.0.0	18 recent observations	78 training-tail $+18$ recent $=96$
Moirai-2	Salesforce/moirai-2.0- R-small ; uni2ts 2.0.0	18 recent observations	1,662 training-tail $+18$ recent $=1,680$
Lag-Llama	Released lag-llama.ckpt	Nominal context: 18 recent; 1,110 raw observations including lag buffer	Nominal context: 14 training-tail $+18$ recent $=32$ ; 1,124 raw observations including lag buffer
TimePFN	Released TimePFN checkpoint with benchmark adapters	18 recent observations, left-padded to length 96	78 training-tail $+18$ recent $=96$

Appendix

Table 3: External zero-shot baseline implementations and context regimes. “Training tail” denotes observations drawn contiguously from the end of the 70% training split. For Lag-Llama, the raw-history totals include the additional buffer required to construct its fixed lag features.

Model	Context	Total	30 min		60 min		90 min
Model	Context	Total	MAE	MSE	MAE	MSE	MAE	MSE
METR-LA
GAS-PFN	Mixed	1,170	$3.97\pm 0.07$	$\mathbf{68.31}\pm 2.35$	$4.83\pm 0.07$	$97.61\pm 3.58$	$5.50\pm 0.06$	$119.69\pm 5.08$
Chronos-2	Recent 18	18	4.31	91.60	5.74	152.61	6.66	193.68
Chronos-2	Recent + prefix	8,192	3.68	70.62	4.38	91.23	4.95	107.42
TimesFM-2.5	Recent 18	18	4.10	86.15	5.19	133.05	6.06	171.30

Appendix

Table 4: Zero-shot MAE and MSE under recent-only and history-augmented context regimes at cumulative 30-, 60-, and 90-minute horizons. Lower is better. Bold and underlined values denote the best and second-best results, respectively, using unrounded values.

Model	Context	Total	30 min		60 min		90 min
Model	Context	Total	MAE	MSE	MAE	MSE	MAE	MSE
PEMS04
GAS-PFN	Mixed	1,170	$\mathbf{22.52}\pm 0.28$	$\mathbf{1243.43}\pm 13.35$	$\mathbf{27.16}\pm 0.41$	$\mathbf{1782.98}\pm 26.45$	$\mathbf{32.32}\pm 0.53$	$\mathbf{2517.88}\pm 54.69$
Chronos-2	Recent 18	18	24.23	1402.54	31.26	2300.66	37.92	3370.61
Chronos-2	Recent + prefix	8,192	25.61	1402.39	32.89	2266.07	40.67	3466.71
TimesFM-2.5	Recent 18	18	25.33	1491.79	32.81	2476.17	40.18	3712.61

Appendix

Table 11

Model	Budget	Context	30 min		60 min		90 min
Model	Budget	Context	MAE	MSE	MAE	MSE	MAE	MSE
METR-LA
GAS-PFN	–	Mixed	$\mathbf{3.97}\pm 0.07$	$\mathbf{68.31}\pm 2.35$	$\mathbf{4.83}\pm 0.07$	$\mathbf{97.61}\pm 3.58$	$\mathbf{5.50}\pm 0.06$	$\mathbf{119.69}\pm 5.08$
TimePFN	FT-32	Recent 18	$5.12\pm 0.67$	$112.67\pm 24.70$	$6.07\pm 0.62$	$155.51\pm 26.45$	$6.78\pm 0.56$	$186.47\pm 29.54$
	FT-32	Recent + prefix	$5.14\pm 0.50$	$107.27\pm 16.52$	$6.06\pm 0.48$	$149.42\pm 19.17$	$6.76\pm 0.43$	$181.29\pm 21.37$
	FT-64	Recent 18	$4.68\pm 0.21$	$97.01\pm 7.15$	$5.58\pm 0.16$	$134.08\pm 10.82$	$6.33\pm 0.13$	$163.68\pm 14.66$

Appendix

Table 5: Comparison of zero-shot GAS-PFN with TimePFN fine-tuned using 32 or 64 target-domain patches. Results are mean

\pm

sample standard deviation over seeds 42, 43, and 44. Lower is better; bold and underlined values denote the best and second-best means, respectively, using unrounded results.

Train roads	Model	30 min		60 min		90 min
Train roads	Model	MAE	MSE	MAE	MSE	MAE	MSE
METR-LA
10%	GAS-PFN	$\mathbf{3.99}\pm 0.08$	$\mathbf{68.29}\pm 2.57$	$\mathbf{4.85}\pm 0.10$	$\mathbf{97.40}\pm 4.41$	$\mathbf{5.52}\pm 0.13$	$\mathbf{119.25}\pm 6.98$
10%	SCPT12	$4.32\pm 0.55$	$76.25\pm 7.20$	$5.48\pm 1.02$	$119.35\pm 19.44$	$6.36\pm 1.26$	$154.20\pm 26.70$
20%	GAS-PFN	$\mathbf{3.99}\pm 0.13$	$\mathbf{68.20}\pm 3.57$	$\mathbf{4.84}\pm 0.17$	$\mathbf{97.04}\pm 6.11$	$\mathbf{5.50}\pm 0.23$	$\mathbf{118.63}\pm 9.13$
20%	SCPT12	$4.26\pm 0.44$	$77.03\pm 6.41$	$5.31\pm 0.49$	$119.07\pm 8.59$	$6.13\pm 0.57$	$152.71\pm 12.12$

Appendix

Table 6: Cross-road comparison of zero-shot GAS-PFN and SCPT12 across supervised road fractions and cumulative 30-, 60-, and 90-minute horizons. Results are mean

\pm

sample standard deviation over paired seeds 42, 43, and 44. Lower is better; bold denotes the lower unrounded mean within each road fraction, horizon, and metric.

Train roads	Model	30 min		60 min		90 min
Train roads	Model	MAE	MSE	MAE	MSE	MAE	MSE
PEMS03
10%	GAS-PFN	$\mathbf{17.64}\pm 0.69$	$\mathbf{753.83}\pm 49.48$	$\mathbf{21.58}\pm 1.02$	$\mathbf{1167.46}\pm 95.74$	$\mathbf{26.05}\pm 1.46$	$\mathbf{1751.57}\pm 173.62$
10%	SCPT12	$23.76\pm 2.41$	$1474.14\pm 269.95$	$30.43\pm 4.14$	$2443.89\pm 623.04$	$36.68\pm 5.58$	$3559.57\pm 1016.24$
20%	GAS-PFN	$\mathbf{17.80}\pm 0.61$	$\mathbf{762.80}\pm 46.22$	$\mathbf{21.77}\pm 0.90$	$\mathbf{1181.56}\pm 91.13$	$\mathbf{26.27}\pm 1.31$	$\mathbf{1771.99}\pm 165.93$
20%	SCPT12	$22.57\pm 2.96$	$1305.09\pm 301.00$	$28.30\pm 4.85$	$2095.35\pm 735.58$	$34.13\pm 7.00$	$3121.79\pm 1310.33$

Appendix

Table 14

Train roads	Model	30 min		60 min		90 min
Train roads	Model	MAE	MSE	MAE	MSE	MAE	MSE
PEMS07
10%	GAS-PFN	$\mathbf{25.21}\pm 0.57$	$\mathbf{1502.58}\pm 44.38$	$\mathbf{30.89}\pm 0.80$	$\mathbf{2238.09}\pm 78.90$	$\mathbf{37.23}\pm 1.05$	$\mathbf{3257.79}\pm 133.56$
10%	SCPT12	$28.34\pm 0.99$	$1796.27\pm 117.07$	$35.22\pm 0.83$	$2763.04\pm 254.87$	$42.19\pm 2.30$	$3988.32\pm 571.39$
20%	GAS-PFN	$\mathbf{25.36}\pm 0.65$	$\mathbf{1522.72}\pm 51.96$	$\mathbf{31.05}\pm 0.90$	$\mathbf{2263.22}\pm 94.49$	$\mathbf{37.41}\pm 1.16$	$\mathbf{3288.00}\pm 156.24$
20%	SCPT12	$27.62\pm 2.68$	$1711.78\pm 330.57$	$35.67\pm 7.02$	$2854.59\pm 1022.40$	$43.91\pm 11.83$	$4411.63\pm 2085.96$

Appendix

Table 15

Variant	30 min		60 min		90 min
Variant	MAE	MSE	MAE	MSE	MAE	MSE
METR-LA
Learned mixture (full model)	$\mathbf{3.97}\pm 0.07$	$\underline{68.31}\pm 2.35$	$\mathbf{4.83}\pm 0.07$	$\underline{97.61}\pm 3.58$	$\mathbf{5.50}\pm 0.06$	$119.69\pm 5.08$
Self expert only	$4.09\pm 0.03$	$71.71\pm 1.58$	$4.97\pm 0.03$	$103.28\pm 2.50$	$5.66\pm 0.04$	$128.33\pm 3.91$
Multi expert only	$4.03\pm 0.10$	$69.13\pm 2.85$	$4.89\pm 0.11$	$98.14\pm 4.12$	$5.56\pm 0.10$	$\underline{119.46}\pm 5.42$
Fixed equal fusion	$3.98\pm 0.06$	$68.39\pm 1.87$	$4.84\pm 0.05$	$97.89\pm 2.88$	$5.51\pm 0.05$	$120.52\pm 4.47$

Appendix

Table 7: Frozen-checkpoint ablations of branch fusion, context construction, and neighborhood routing across six datasets and cumulative 30-, 60-, and 90-minute horizons. Results are mean

\pm

sample standard deviation over seeds 42, 43, and 44. Lower is better; bold and underlined values denote the best and second-best unrounded means, respectively.

Traffic Forecasting

Traffic forecasting predicts future flows, speeds, demand, travel times, and congestion from spatial, temporal, and contextual data.

See all 87 papers

Graph-Conditioned Mixture of Graph Neural Network Experts for Traffic Forecasting
May 28, 2026Amirhossein Ghaffari, Saeid Sheikhi, Ekaterina GilmanAdaptive Model Routing Mixture-of-Experts Models
Future Computing Group, University of Oulu
General Semantic Knowledge Infusion for Spatio-Temporal Traffic Forecasting
Aug 18, 2026Mattis Thor Straten, Yannick Wolker, Steffen Strohm +3Graph Neural Networks Traffic Flow Estimation
Department of Computer Science, Kiel University, Germany · GEOMAR Helmholtz Centre for Ocean Research Kiel, Germany · ZBW – Leibniz Information Centre for Economics, Kiel, Germany
Spectral Adaptive Graph Transformer for Efficient and Scalable Traffic Forecasting
Oct 18, 2026Amirhossein Ghaffari, Susanna Pirttikangas, Ekaterina GilmanSpatiotemporal Forecasting Traffic Flow Estimation
Future Computing Group, University of Oulu, 90570 Oulu, Finland

May 28, 2026J/K move · Enter open · S save

Graph-Conditioned Mixture of Graph Neural Network Experts for Traffic Forecasting

Amirhossein Ghaffari, Saeid Sheikhi, Ekaterina Gilman

Future Computing Group, University of Oulu

Adaptive Model Routing Mixture-of-Experts Models Traffic Forecasting Spatiotemporal GNNs

Open paper AskPDF arXiv

Spatio-temporal forecasting on sensor graphs is commonly tackled with a single backbone architecture applied uniformly across all nodes, although graph regions can exhibit different dynamics. Road segments differ in functional class, structure, and traffic behavior, suggesting that node-wise expert specialization can be useful. We propose GC-MoE, a graph-conditioned mixture of experts framework that assigns each node a personalized combination of frozen forecasting experts based on graph topology and the recent traffic input window. GC-MoE combines frozen pretrained spatio-temporal GNN experts with an input-aware, spatially contextualized router while training only a lightweight routing module. We also study a bounded graph-conditioned output refinement layer as an optional extension and include node-adaptive ST-LoRA adapters only as an ablation diagnostic. Across four standard benchmarks (PEMS04, PEMS07, METR-LA, and PEMS-BAY), GC-MoE improves MAE over a zero-parameter ensemble baseline, with competitive RMSE and MAPE, while training only ~17K parameters on top of 1.5M frozen expert weights. The implementation is available at https://github.com/Ahghaffari/gc_moe.

Abstract

Traffic Forecasting

Explore similar work

Graph-Conditioned Mixture of Graph Neural Network Experts for Traffic Forecasting

General Semantic Knowledge Infusion for Spatio-Temporal Traffic Forecasting

Spectral Adaptive Graph Transformer for Efficient and Scalable Traffic Forecasting