论文精读 |2024[KDD]UniST: One-for-all城市时空预测模型，预训练+提示统一多种时空场景

论文标题：UniST: A Prompt-Empowered Universal Model for Urban Spatio-Temporal Prediction

作者：Yuan Yuan（苑苑）, Jingtao Ding（丁璟韬）, Jie Feng（冯杰）, Depeng Jin（金德鹏）, Yong Li（李勇）

机构：清华大学

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

代码和数据开源地址：https://github.com/tsinghua-fib-lab/UniST

量子位：清华推出首个通用城市时空预测模型UniST，零样本场景开箱即用｜KDD2024

点击文末阅读原文跳转本文arXiv链接。

引

言

最近，清华大学团队推出了第一个无需自然语言的纯时空通用模型 — UniST，第一次展示了纯时空模型本身的通用性和可扩展性，研究成果已被KDD2024接收。

研究团队利用超过20个时空数据集，包括超过1.3亿个时空样本点，统一了多个城市、多领域、不同空间划分和时间分辨率等丰富的城市时空数据，构建并训练了UniST这样一个「one-for-all」的时空通用模型。这是目前覆盖范围最广、统一性最强的城市时空通用模型。值得一提的是，UniST相比于当前的大语言模型具有更轻量级的优势，仅用20M的参数规模就展现出了很强的零样本学习能力。

UniST通过整合多个城市和多领域丰富的时空数据，利用基于Transformer的架构，以及独特的时空掩码策略（mask）和知识引导的时空提示（prompt），实现了对城市多样化时空动态性的统一建模和精准预测。在实验中，UniST展示了其在交通管理、资源优化等多个城市应用场景中的卓越表现，尤其是在跨场景零样本预测（zero-shot）中，其性能超过了少样本（few-shot）基线方法，展现出强大的通用性和泛化能力。

UniST的发布标志着通用基础模型在城市时空领域的重要突破，推动该领域进入一个全新阶段，有望促进全球智慧城市的发展。该成果的论文、代码和数据均已公开，供研究和应用者使用。

纯时空模型的逆袭

时空预测在城市中无处不在，它不仅关注交通和人流的流动，还涵盖了资源分配等多个维度。然而时空预测并不容易，模型需要处理复杂的且动态的时空关联。传统的AI方法需要大量的训练数据和领域知识，一般只能针对特定的数据集对模型进行训练，多个时空场景就需要训练多个模型，这在城市数据不足的情况下显得尤为困难。

与此同时，随着大语言模型的爆发，研究者们开始尝试使用「文本」来完成时空相关的任务，将文本描述与时空多模态数据结合。然而，在面对复杂的时空场景时，这种方法容易忽略大量的时空耦合和动态信息。事实上，时空数据的产生本质上并不依赖语言。因此，清华大学的研究团队选择了一条不同于大语言模型的方向：仅依靠时空数据，我们能走多远？时空通用模型是否能像自然语言大模型一样存在？

具体来说，清华大学的研究团队致力于训练一个纯时空通用模型，能够模仿大语言模型（LLM）的两个关键特性：

1. 对丰富的时空数据具有强大的拓展能力；

2. 像大语言模型一样，展现出强大的通用性和泛化能力。

值得一提的是，纯时空模型背后的直觉是：在人类干预下，城市运转的各种时空数据存在通用规律，可以通过类似于GPT的方式进行训练。

通用时空建模的挑战

挑战1：时空数据格式不统一

在自然语言处理中，数据通常是统一的1D序列格式；在计算机视觉中，无论是图片还是视频，也都遵循较为标准的格式。然而，时空数据在不同时空场景下，例如不同城市、不同领域，由于数据收集者和收集方式的不同，其数据形状以及时空分辨率存在明显的差异。这种多样性使得统一处理和分析时空数据变得异常困难。

挑战2：不同时空场景数据分布差异大

不同城市、不同地理空间、不同时间段的时空数据往往展现出显著的分布差异。此外，不同领域的数据，例如空气污染数据、交通数据、人流数据和网络基站数据，也存在显著的分布差异。这些差异增加了模型的复杂性，要求模型具备强大的泛化能力以适应各种数据分布。

如何构建纯时空通用模型

尽管不直接使用大语言模型， 但LLM的成功经验不可忽视。研究团队从LLM的思想出发， 实现了以下几个关键特性：

·     具有在多样数据下的拓展能力；

·     自监督预训练充分捕捉复杂时空关联；

·     通过提示（prompt）灵活进行泛化。

与已有的时空模型不同的是，UniST在以下几方面实现了突破：

·     灵活适应多样的时空数据特征：UniST能够处理不同城市、不同领域的多样化时空数据，实现真正的统一和通用模型。无论是交通数据、人群流动数据还是城市资源分布数据，UniST都能灵活应对 ，展现出强大的可拓展性。

·     高效的生成式预训练：通过巧妙设计的掩码策略，UniST能够捕捉复杂的时空关系，实现全面多维度的时空建模。

·     时空知识引导的提示：利用知识引导的时空提示，UniST能够对不同场景的内在和共享知识进行对齐和利用，提升预测性能。通过这种提示机制，UniST可以在数据稀缺或全新的场景中依然保持高效的预测能力。

时空数据的序列建模

为了有效处理不同来源、不同特征的时空数据，UniST提出了一种名为「时空序列」的建模方法。具体来说，时空数据被表示为一个四维张量：，其中 表示时间段数量， 表示变量数量， 和 分别表示空间划分中的纬度和经度网格数。为了统一处理这些不同形状的时空数据，UniST引入了时空编码器，将这些四维张量转换为小的三维向量，然后按照位置展开成序列。一旦将丰富的时空数据表征为为这种「时空序列」的通用格式，就可以利用Transformer强大的序列建模能力，进行模型的训练，全面捕捉复杂的时空关系。

得到「时空序列」之后，UniST通过生成式预训练进一步提升其建模能力。预训练过程中，研究团队采用了多种掩码策略，希望帮助模型更好地理解和捕捉时空关系。具体来说，UniST引入了以下几种掩码策略：

·     随机掩码（Random Masking）：类似于MAE中的随机掩码策略，通过随机遮蔽时空数据块来捕捉细粒度的时空关系。

·     管状掩码（Tube Masking）：模拟某些空间单元在所有时间段内的数据缺失情况，提升模型的空间外推能力。

·     块状掩码（Block Masking）：一种更具挑战性的掩码方式，通过遮蔽整个空间单元块在所有时间段内的数据，增强模型在有限上下文信息下的空间迁移能力。

·     时间掩码（Temporal Masking）：遮蔽未来的数据，仅依赖历史信息进行重建，旨在提升模型捕捉从过去到未来时间依赖关系的能力。

通过这些掩码策略，UniST在预训练阶段系统地增强了其从多角度捕捉时空关系的能力，不仅提高了模型的泛化性能，还显著减少了对大量标记数据的依赖。

知识引导的时空提示

在UniST中，提示机制（prompt）是提升模型泛化能力的关键。为了在不同的时空场景中保持高效预测，UniST设计了基于时空知识的提示网络（prompt network）。该提示网络利用已知的时空领域知识，生成有助于模型理解和预测的提示信息。

具体来说，提示网络基于以下四个方面的时空知识进行提示生成：

·     空间临近性：临近的空间单元可能相互影响；

·     空间层次结构：城市结构的层次组织影响时空动态；

·     时间临近性：近期的动态会影响未来结果；

·     时间周期性：每天或每周的相似模式会影响未来周期性结果。

如下图所示，提示网络从记忆池中提取有用的提示，这些记忆池存储了优化后的时空领域知识。提示的生成过程利用时空特征表示作为查询，提取相应的记忆向量，这些提示向量再集成到Transformer架构的输入空间，提升模型的预测能力。

实验结果

在15个城市和6个领域的广泛实验中，UniST展示了其卓越的通用性和强大的预测能力。特别是在少样本和零样本场景下，UniST表现出色，大幅提升了时空预测的准确性。实验结果表明，UniST在多个任务上的表现均超越了当前最先进的基线模型，证明了其在不同城市、不同数据集上的强大适应能力。

·     少样本学习：在训练数据有限的情况下，UniST依然能够提供高精度的预测。

·     零样本学习：在模型从未见过的时空场景下，UniST依然能实现出色的预测性能，甚至超过了大多数监督学习方法。

·     广泛适用性：在交通预测、人群流动预测、资源分配等多个任务中，UniST均展示了其强大的预测能力和适用性。

研究人员深入分析了提示（prompt）机制的作用。在时间记忆池中，他们研究了每个向量的记忆模式，根据该向量被数据集索引的权重高低，聚合数据集样本值在该向量上的结果。图 7(a) 和图 7(b) 展示了两个数据集（Crowd 和 TrafficSH）的结果。可以看到，提示机制中展现的记忆模式在不同的城市场景中表现出显著的一致性。这不仅证实了每个记忆向量都被很好地优化以记忆独特的时空模式，还证明了空间和时间记忆池在不同场景中的稳健性。

进一步，研究团队分析了两个不同场景对记忆向量的利用情况（获得的时空提示）。具体来说，他们计算了在每个数据集上下文中每个向量的平均注意力权重。图 7(c) 和图 7(d) 展示了比较结果。两个数据集的注意力权重分布显示出明显的不同。观察到的注意力权重分布的独特性表明，模型能够根据输入数据的特征动态调整其关注的记忆模式。这种动态调整注意力权重的能力增强了 UniST 模型在不同数据集上的适配性和泛化性。

这些实验结果表明，UniST在提示机制的帮助下，能够在不同的时空场景中有效地捕捉和利用重要的时空关系，这种创新方法使得UniST在应对复杂多变的时空数据时表现出色，展示了其强大的适应能力和广泛的应用潜力。

结语

UniST的发布不仅代表了城市时空领域的重要突破，也标志着通用大模型在复杂城市计算中的广泛应用前景。通过整合多城市、多领域的时空数据，UniST展示了其在少样本和零样本学习场景中的卓越性能，以及在交通管理、人群流动预测和资源分配等多个实际应用中的广泛适用性。随着智能城市建设的推进，UniST有望在全球范围内推动智慧城市的发展，为城市管理者提供更加精准的数据支持和决策依据。

研究团队将继续探索UniST的潜力，期待未来的研究能够进一步提升模型的性能和适应性，推动城市时空进入一个更加智能和高效的新阶段。

李勇老师团队最新在arXiv上挂出来两篇关于城市科学+LLM的最新研究：

1. CityGPT: Empowering Urban Spatial Cognition of Large Language Models

论文：https://arxiv.org/abs/2406.13948

TL; DR: 这篇论文提出并解决了大型语言模型（Large Language Models, LLMs）在理解和解决城市空间相关问题时面临的挑战。具体来说，论文中提到LLMs在经过训练后虽然在多个领域（如数学和代码生成）取得了成功，但由于缺乏物理世界城市空间的语料库和知识，它们在处理城市尺度上的真实任务时常常失败。为了解决这个问题，论文提出了一个名为CityGPT的系统框架，旨在通过构建城市规模的世界模型来增强LLMs对城市空间的理解能力，并解决相关的城市任务。

2. CityBench: Evaluating the Capabilities of Large Language Model as World Model

论文：https://arxiv.org/abs/2406.13945

TL; DR：这篇论文提出了一个名为CityBench的评估平台，旨在解决对大型语言模型（LLMs）在城市领域作为世界模型（world model）的能力进行系统性和可靠性评估的问题。具体来说，论文提出了CityBench平台，它包含一个名为CitySim的交互式模拟器，用于集成多源数据并模拟城市动态。基于CitySim，论文设计了7个任务，分为两个类别：感知理解（perception-understanding）和决策制定（decision-making），以评估LLMs作为城市规模世界模型在城市领域的表现。此外，CityBench平台可以轻松扩展到世界任何城市，为评估提供了灵活性和易用性。

欢迎各位作者投稿近期有关时空数据和时间序列录用的IJCAI、SIGIR、ICML以及KDD等顶级会议和期刊的优秀文章解读，我们将竭诚为您宣传,共同学习进步。如有意愿,请通过后台私信与我们联系。