ICML 2024 论文分享┆一个简单且通用的交通预测提示调优框架

论文简介

本推文介绍了2024 ICML的优秀论文之一《FlashST: A Simple and Universal Prompt-Tuning Framework for Traffic Prediction》。论文的核心目标是通过整合空间和时间因素，精准地预测和分析交通流量的动态变化。然而，在交通预测领域，分布偏移问题始终是一个重大的挑战。当前的模型在面对测试数据与训练数据存在显著差异的情况下，往往表现出较弱的泛化能力，难以适应复杂的实际交通情况。为了解决这一问题，论文提出了一种简单而通用的时空提示调优框架——FlashST，该框架能够根据多种下游数据集的特征调整预训练模型，从而改善在不同交通预测场景中的泛化能力。具体而言，FlashST框架采用了轻量级的时空提示网络，旨在进行上下文学习，捕捉时空不变知识，并促进有效适应。

推文作者为邱雪，审校为许东舟和黄星宇。

论文链接：https://arxiv.org/abs/2405.17898

1.会议介绍

ICML（International Conference on Machine Learning，国际机器学习会议）创办于1980年，由国际机器学习协会主办，每年举办一次。作为人工智能（Artificial Intelligence, AI）和机器学习领域中最具有影响力的会议之一，它的讨论主题领域广泛，包括通用机器学习、深度学习、学习理论、优化、概率推理等。

2.研究背景

近年来，交通预测已成为智能交通系统和城市规划中的一个重要研究领域。准确预测交通流量和模式不仅有助于缓解交通拥堵，还能优化资源配置，提高运输效率。然而，现有模型在面对分布偏移时往往表现不佳，难以适应测试数据与训练数据之间的显著差异。这种分布偏移使得交通预测面临着更大的挑战，因此，提高模型的泛化能力变得尤为重要。

基于上述背景，论文做出了以下的研究创新：

（1）时空提示调优

提出了一种简单且通用的时空提示调优框架——FlashST，旨在通过对预训练模型进行适应性调整，以应对多样化的下游数据集，最终提升模型在不同下游任务中的适应能力。

（2）通过上下文学习实现的适应

利用上下文蒸馏机制，捕捉来自未见数据的时间和空间上下文信号，帮助模型适应不同的时空场景。

（3）统一分布映射机制

引入了统一分布映射机制，通过对预训练阶段与下游数据集的分布进行对齐，减小分布差异，从而促进了跨不同城市交通数据集的有效知识迁移。

图1 FlashST提出的必要性。左图展示了不同时空数据集之间多样的数据分布，而右图则表明端到端模型的参数在训练集A上出现过拟合，无法很好地泛化到测试集B。

图1表明了FlashST框架通过提示调优和分布映射机制来解决这一问题的必要性。其中左图展示了不同时空数据集之间的分布差异；右图展示了现有端到端模型在训练集上表现良好，但在测试集上泛化能力差的问题。

3.方法

图2 我们提出的FlashST框架采用了一种集成时空上下文学习和统一分布映射机制的架构，为在不同场景下进行时空提示调优提供了一种高效且有效的解决方案。

FlashST框架通过三个主要模块来实现对交通数据的时空建模与预测。这些模块分别负责捕捉时空上下文、建模时空依赖关系以及优化数据分布映射，使得模型可以更好地在不同交通场景中进行泛化和适应。

（1）时空上下文学习

a)时空上下文蒸馏：首先将输入的时空数据嵌入到一个初始的时空表示中，然后使用线性层将时间特征和空间特征转化为上下文嵌入，捕捉基础的时空关系。

b)时空依赖建模：空间依赖通过图神经网络（GNNs）进行传播，帮助模型理解交通网络中位置间的关系。时间依赖通过多层感知机（MLP）和门控机制捕捉，建模时空数据中的动态变化。

（2）统一分布映射机制

首先通过提示嵌入（Prompt Embedding）和InfoNCE损失优化，将输入数据的分布映射为更标准、均匀的分布。同时，处理正负样本对，优化信息嵌入，使模型能够更好地适应不同区域的数据。

4.数据集

表1 预训练数据集的统计信息

表2 下游任务数据集的统计信息

表1和表2展示了FlashST模型在预训练和下游任务中所使用的数据集的统计信息。具体如下：

（1）预训练数据集

预训练阶段使用了PEMS系列交通流量数据集，分别是PEMS03、PEMS04、PEMS07和PEMS08。这些数据集的记录类型都是交通流量，涵盖了不同的区域数量和时间步长，采样时间从2016年到2018年不等。

（2）下游任务数据集

下游任务阶段使用了多种类型的数据集，包括PEMS07(M)交通速度数据集、CA-D5交通流量数据集、ChengDu-DIDI交通指数数据集以及NYC Citi Bike单车订单数据集。这些数据集用于测试模型在不同任务场景中的泛化能力，涵盖了不同的区域数和时间步长。

5.实验及结果

（1）对比实验

表3展示了在PEMS07(M)、CA-D5、ChengDu-DIDI和NYC Citi Bike等多个数据集上，各种模型在三个评价指标（MAE、RMSE和MAPE）下的整体性能表现。

表3 在PEMS07(M)、CA-D5、成都-DIDI和纽约市 CitiBike 数据集上的总体表现，评估指标包括 MAE、RMSE 和 MAPE

根据表3可以看出模型的对比结果如下：

a)传统模型（如TGCN、STGCN、ASTGCN等）在不同的数据集上的性能表现相对较弱，尤其是TGCN模型在多个数据集上显示出较高的误差值，说明其泛化能力不够强；

b)MTGNN模型在NYC Citi Bike数据集上表现出色，特别是 MAPE最低，表明其对交通流量预测的有效性较强；

c)我们的模型（Ours）在所有数据集上都显示了优异的性能。特别是在ChengDu-DIDI和NYC Citi Bike数据集上，其MAE、RMSE和MAPE均优于其他模型，说明模型在不同任务中的鲁棒性和泛化能力都非常突出。

总的来说，FlashST在所有数据集和评估指标上均取得了显著的优势，这表明该模型能够更好地处理不同交通场景中的预测任务，并具有较强的泛化能力。

（2）模型无关和模型微调

表4展示了在PEMS07(M)、CA-D5、ChengDu-DIDI和NYC Citi Bike数据集上的模型无关实验结果，分别比较了STGCN、GWN、MTGNN、PDFormer以及“我们的模型”（Ours）的性能表现，主要以三项指标MAE、RMSE和MAPE进行评估。同时，对比了模型在没有进行微调（w/o Finetune）和经过微调（w/ Finetune）后的表现。

表4 在 PEMS07(M)、CA-D5、成都-DIDI 和纽约市 CitiBike 数据集上进行的模型无关实验

根据表4可以看出模型的无关实验结果如下：

a)无微调模型的表现：可以看到，未经过微调的模型（w/o Finetune）在所有数据集上的误差（MAE、RMSE、MAPE）显著高于微调后和我们的模型。这表明，如果模型没有适应特定的数据分布，其性能会大幅下降，尤其是在CA-D5数据集上的表现极差；

b)经过微调后的表现：在微调后，所有模型的误差都明显下降，但仍未能超越“我们的模型”。例如，STGCN经过微调后，MAE在PEMS07(M)数据集上从8.07下降到3.18，但我们的模型仍然保持在2.68的较低值。

c)我们的模型的表现：不论是哪个数据集，我们的模型在所有指标上均表现出色。例如，在ChengDu-DIDI数据集中，经过微调的MTGNN模型的MAE为2.33，但我们的模型达到了更低的2.31。类似地，在NYC Citi Bike数据集上，我们的模型以最小的MAPE（50.06%）和RMSE（2.67）领先。

表4证明了FlashST模型在多种数据集上的优越性，即便在其他模型经过微调之后，仍能在MAE、RMSE和MAPE上保持较低的误差。这表明了我们提出的方法具有更强的泛化能力和稳健性，特别是在处理不同的数据分布时，能够显著减少预测误差。

（3）模型效率评估（训练时间）

表5展示了不同模型在计算时间上的开销（以秒为单位），以比较“我们的模型”与其他模型在进行预测任务时的效率。

表5 计算时间成本调查（秒）

根据表5可以看出传统模型和微调后的模型虽然在预测性能上有所提升，但计算开销较大。而结合FlashST框架的模型在保持良好预测性能的同时，极大地减少了计算时间，表现出更好的计算效率和实用性。

（4） FlashST收敛速度

图3展示了在PEMS07(M)和CA-D5两个数据集上不同模型的验证损失（Validation Loss）随训练轮次（Epochs）变化的对比情况。曲线图对比了原始的MTGNN模型、经过微调的MTGNN模型（Finetune），以及“我们的模型”（Ours）的性能表现。

图3 FlashST的收敛效率

从图中可以看出，在这两个数据集上都表现出色，验证损失在较少的训练轮次内迅速降低并趋于稳定，明显优于原始模型和经过微调的模型。这表明我们的模型在时空数据建模和预测任务中具有更强的泛化能力和稳定性。

（5）消融实验

图4展示了在PEMS07(M)和CA-D5数据集上，不同模型配置下的MAE、RMSE和MAPE的对比结果。每个图包含多个柱状图，分别表示不同特征的移除或变更对模型性能的影响，比较了多个配置和“我们的模型”（Ours）的表现。

图4 FlashST的消融实验

从图中可以看出，不同特征的移除或变更会对模型的性能产生不同程度的影响。其中，时间上下文和统一分布映射的移除对性能影响最大。“我们的模型”通过保留所有关键模块，在不同数据集上均表现出最优的性能。

6.总结及展望

本文介绍了FlashST，这是一种用于将时空预测模型适应于未见过数据的下游任务的框架。FlashST通过引入时空提示网络，包含时空上下文学习机制和时空依赖性建模方案，能够有效地捕捉上下文信号并建模时间和地点之间的复杂关系，从而适应不同的时空场景。为了解决分布差异问题，框架还整合了一个统一分布映射机制，促进了预训练数据与下游数据分布的对齐，从而实现了高效的知识转移。

在未来的研究中，FlashST将继续优化其在多种下游时空预测场景中的适应能力。一个有前景的研究方向是探索如何将大规模语言模型（LLMs）整合到FlashST框架中，进一步提升模型的泛化性与知识指导能力。