通过记忆增强提示学习,提升城市流量预测在分布偏移下的鲁棒性。
Memory-enhanced Invariant Prompt Learning for Urban Flow Prediction under Distribution Shifts
- 构建可学习的记忆库,存储时空图中的因果特征以生成不变与变化提示。
- 在两种公开数据集上对分布外数据的预测误差降低32%以上。
- 适合需要应对突发交通变化的城市交通系统建模者使用。
城市流量预测是经典的时空预报任务,旨在估计特定位置未来的交通流量。尽管基于时空图神经网络(STGNNs)的模型已具备较强预测能力,但在面对城市流量数据常见的分布偏移时仍表现不佳,这源于时空事件的动态性和不可预测性。然而,在时空应用中,动态环境难以用固定参数量化,而为每个时间-位置学习环境则计算成本过高。本文提出一种名为记忆增强不变提示学习(MIP)的新框架,用于应对恒定分布偏移下的城市流量预测。MIP配备可学习的记忆库,用于存储时空图中的因果特征。通过查询该可训练记忆库,可自适应地提取每个时间步、每个位置的不变与变化提示。相较于基于模拟环境干预原始数据,本方法直接对变化提示进行跨时空干预。利用不变学习最小化预测方差,确保预测仅依赖于不变特征。在两个公开城市流量数据集上的大量对比实验表明,MIP在分布外数据下具有显著鲁棒性。
原文摘要 · Abstract (English)
Urban flow prediction is a classic spatial-temporal forecasting task that estimates the amount of future traffic flow for a given location. Though models represented by Spatial-Temporal Graph Neural Networks (STGNNs) have established themselves as capable predictors, they tend to suffer from distribution shifts that are common with the urban flow data due to the dynamics and unpredictability of spatial-temporal events. Unfortunately, in spatial-temporal applications, the dynamic environments can hardly be quantified via a fixed number of parameters, whereas learning time- and location-specific environments can quickly become computationally prohibitive. In this paper, we propose a novel framework named Memory-enhanced Invariant Prompt learning (MIP) for urban flow prediction under constant distribution shifts. Specifically, MIP is equipped with a learnable memory bank that is trained to memorize the causal features within the spatial-temporal graph. By querying a trainable memory bank that stores the causal features, we adaptively extract invariant and variant prompts (i.e., patterns) for a given location at every time step. Then, instead of intervening the raw data based on simulated environments, we directly perform intervention on variant prompts across space and time. With the intervened variant prompts in place, we use invariant learning to minimize the variance of predictions, so as to ensure that the predictions are only made with invariant features. With extensive comparative experiments on two public urban flow datasets, we thoroughly demonstrate the robustness of MIP against OOD data.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。