AAAI 2023 | 自动时空多任务学习AutoSTL

今天跟大家分享一篇香港城市大学，吉林大学，奥尔堡大学和京东智能城市研究院在AAAI2023发表的"AutoSTL: Automated Spatio-Temporal Multi-Task Learning"。

针对时空数据挖掘算法难以处理多任务场景的问题，作者提出一种基于自动机器学习的自适应时空多任务学习模型。具体地，模型基于多任务学习结构，设置专用和共享模块学习不同任务以及其间的关系，定义时空数据挖掘操作符集合，通过Gumbel-Softmax近似离散选择，并利用自动机器学习自适应地分配模块操作和权重。实验结果表明，所提出模型能够适应不同时空多任务场景，在多个时空多任务数据集上取得了最佳表现。

1 动机

目前关于时空预测的研究中，绝大多数工作只关注在单一任务上预测精度的提升，如交通速度，交通流等。但构建智慧交通系统往往需要同时预测多个时空任务，应用时空单任务的模型会造成跷跷板现象，即模型训练过程中部分任务起到主导作用，其他任务效果很差。

此外，挖掘和利用时空多任务之间的相关性也能促进对各个任务的时空变化规律的学习，作者举例展示了时空多任务之间的时间与空间相关性：

以交通流预测中的输入交通流和输出交通流为例，图a展示了区域的输入和输出流高度相关；在图b中，区域的输入流和的输出流也表现出了相似的变化趋势和周期性。

本文的主要贡献可以被概括为以下三点：

1. 提出时空多任务模型AutoSTL，可以自适应地决定各个模块的时空操作符以及模块融合权重；
2. 提出多任务共享模块，有效挖掘时空多任务中公共的时空变化规律，促进各个任务的学习；
3. 本文是首个基于自动机器学习解决时空多任务的工作，在多个时空多任务场景中取得了最佳效果。

2 方法

如左图所示，AutoSTL自下而上由shared bottom，层hidden layer和tower组成。右图为hidden layer的内部结构，主要包括任务特定与任务共享模块和模块融合机制。

2.1 shared bottom

首先，我们使用多任务共享的MLP结构来将输入的时空多任务历史特征映射到嵌入式空间：

在shared bottom层之后，我们将不同任务的特征分开输入各自任务的模块，以及预测使用的tower。

2.2 hidden layer

在每层hidden layer中，我们首先定义一个时空操作符集合，其包含5个时间和空间特征挖掘的操作。基于此，我们设置了任务特定模块和任务共享模块，模块的操作在时空操作符集合中选择。任务特定和任务共享模块处理后，我们设置了模块融合机制。在此过程中，模块的时空操作符选择和融合机制的权重均由自动机器学习优化确定。

2.1.1 时空操作符集合

我们选择了5个时空预测模型中有效的时空操作符：diffusion convolution[1], 1-D dilated causal convolution[2], LSTM, informer[3] 和spatial-informer[4]来构成时空操作符集合。其中diffusion convolution和spatial-informer挖掘空间信息，1-D dilated causal convolution, LSTM和informer挖掘时间信息。

5个时空操作符的公式如下，其中spatial-informer与informer相比，需要将transformer操作的维度切换为空间维度，其余操作一致。

值得一提的是，时空操作符集合可以被轻易地压缩和扩展，以适应不同复杂度的任务和数据。

2.2.2 任务特定模块和任务共享模块

我们在每层hidden layer中设置任务特定模块与任务共享模块。举例来说，同时处理任务A和任务B时，第层hidden layer包括1个任务A特定模块，1个任务B特定模块，和1个多任务共享模块。其中，为了指定每个模块执行的具体操作，我们对每个模块设置权重来加权时空操作符中的操作。由此，每个模块执行的操作即为对时空操作符集合的加权和：

为了逼近hard selection，我们使用了Gumble-softmax：

给定两个任务的特征和为输入，输入任务A的特定模块以及任务共享模块，分别得到和，输入任务B的特定模块以及任务共享模块，分别得到和。

2.2.3 模块融合机制

为了融合多个模块的输出，我们设置了可学习权重。由此，第层hidden layer输出任务A的特征和任务B的特征分别为：

2.3 tower

经过了层hidden layer对多任务特征的表示学习，我们设置了MLP结构组成的任务特定tower，任务A的输出和任务B的输出分别为：

2.4模型优化

AutoSTL的优化过程分3阶段进行：Pretrain，search，和retrain。

1. Pretrain：初始化模型参数，固定结构参数和为均值，训练若干epoch。
2. Search：随着模型参数训练，以1 mini-batch的验证集数据梯度更新结构参数和，即Bi-level optimization。
3. Retrain：结构参数优化结束后，根据确定每个模块最优的时空操作符，根据固定模块融合机制的权重，进行完整的模型训练。

其中Bi-level optimization的损失函数为：我们用1 step的梯度来近似内循环的优化目标：

3 实验

我们设置了3中时空多任务实验场景，分别为NYCTaxi上的交通输入流&交通输出流 （In&Out） ，NYCTaxi上的交通需求流&行程时间（OD&Duration），和PEMSD4上的交通流&速度（Flow&Speed）。评价指标选择了RMSE和MAE。通过与目前先进的时空预测方法对比，我们的AutoSTL在多种设置上稳定地取得了最优效果：

我们还设置了消融实验验证所提结构的有效性：

1. w/o ：对于3层hidden layer结构自底向上分别设置为GCN，GCN和RNN，模拟手工建模
2. w/o ：固定融合权重为均值
3. w/o s-m：去掉每层hidden layer的任务共享模块

实验结果显示每层hidden layer内部的操作和模块融合权重自适应决定，并且删除任务共享模块后模型表现出了跷跷板现象。

4 总结

本文提出了一种自适应的时空多任务模型AutoSTL，以同时预测多个时空任务。它拥有一个灵活的架构，由包含先进时空操作的时空操作符集合组成。它能有效挖掘多个时空任务之间的内在关系，并基于自动机器学习自适应地分配结构参数。作为第一个自动解决时空多任务学习的模型，它在广泛的实验中取得了最先进的性能。

参考文献

[1] Li, Y.; Yu, R.; Shahabi, C.; and Liu, Y. 2018. Diffusion Convolutional Recurrent Neural Network: Data-Driven Traffic Forecasting. In International Conference on Learning Representations.

[2] Wu, Z.; Pan, S.; Long, G.; Jiang, J.; and Zhang, C. 2019. Graph WaveNet for Deep Spatial-Temporal Graph Modeling. In Proceedings of the Twenty-Eighth International Joint Conference on Artificial Intelligence, IJCAI-19, 1907–1913. International Joint Conferences on Artificial Intelligence Organization.

[3] Zhou, H.; Zhang, S.; Peng, J.; Zhang, S.; Li, J.; Xiong, H.; and Zhang, W. 2021. Informer: Beyond efficient transformer for long sequence time-series forecasting. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, 11106–11115.

[4] Wu, X.; Zhang, D.; Guo, C.; He, C.; Yang, B.; and Jensen, C. S. 2021. AutoCTS: Automated correlated time series forecasting. Proceedings of the VLDB Endowment, 15(4): 971–983.