PRIME: A Knowledge Routing Framework for Multi‐Task Spatio‐Temporal Traffic Forecasting

ABSTRACT

Spatio‐temporal prediction is a fundamental component of intelligent transportation systems (ITS). While multi‐task learning provides a natural framework for exploiting shared dependencies across related traffic prediction tasks, existing methods face two key challenges: modelling heterogeneous spatio‐temporal patterns and lacking effective mechanisms for dynamic knowledge sharing. To address these issues, we propose PRIME, a unified multi‐task spatio‐temporal forecasting framework based on dynamic knowledge routing. PRIME is optimised jointly over all tasks in each urban dataset, so that shared parameters receive concurrent supervision from heterogeneous prediction objectives. A spatio‐temporal‐aware prompt pool retrieves context‐relevant knowledge through multi‐modal queries, while an adaptive spatio‐temporal attention (ASTA) module equipped with mixture‐of‐experts (MoE)‐based conditional computation determines how the retrieved knowledge is processed. Together, these components form a hierarchical routing mechanism that controls both knowledge selection and conditional computation. Experiments on two real‐world multi‐task traffic datasets show that PRIME achieves competitive overall performance across heterogeneous tasks. Results over five independent random seeds further demonstrate stable performance under stochastic initialisation and noisy expert routing. The findings support dynamic knowledge routing as an effective approach to jointly modelling shared and task‐dependent spatio‐temporal patterns.