KDD 2024 | UniST：清华推出首个通用城市时空预测模型，代码数据均已公开

城市时空预测对于明智的决策至关重要，例如交通管理、资源优化和应急响应。尽管预训练自然语言模型取得了显著突破，使一个模型能够处理多种任务，但时空预测的通用解决方案仍然具有挑战性。现有的预测方法通常针对特定的时空场景量身定制，需要特定任务的模型设计和大量的特定领域训练数据。近期，清华大学团队推出了第一个无需自然语言的纯时空通用模型 UniST，第一次展示了纯时空模型本身的通用性和可扩展性，研究成果已被 KDD2024 接收。研究团队利用超过 20 个时空数据集，包括超过 1.3 亿个时空样本点，统一了多个城市、多领域、不同空间划分和时间分辨率等丰富的城市时空数据，构建并训练了 UniST 这样一个「one-for-all」的时空通用模型。这是目前覆盖范围最广、统一性最强的城市时空通用模型。值得一提的是，UniST 相比于当前的大语言模型具有更轻量级的优势，仅用 20M 的参数规模就展现出了很强的零样本学习能力。
UniST 的发布标志着通用基础模型在城市时空领域的重要突破，推动该领域进入一个全新阶段，有望促进全球智慧城市的发展。该成果的论文、代码和数据均已公开，供研究和应用者使用。

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

【论文地址】https://arxiv.org/abs/2402.11838

【代码&数据开源地址】https://github.com/tsinghua-fib-lab/UniST
 

纯时空模型的逆袭

时空预测在城市中无处不在，它不仅关注交通和人流的流动，还涵盖了资源分配等多个维度。然而时空预测并不容易，模型需要处理复杂的且动态的时空关联。传统的AI方法需要大量的训练数据和领域知识，一般只能针对特定的数据集对模型进行训练，多个时空场景就需要训练多个模型，这在城市数据不足的情况下显得尤为困难。
与此同时，随着大语言模型的爆发，研究者们开始尝试使用「文本」来完成时空相关的任务，将文本描述与时空多模态数据结合。然而，在面对复杂的时空场景时，这种方法容易忽略大量的时空耦合和动态信息。事实上，时空数据的产生本质上并不依赖语言。因此，清华大学的研究团队选择了一条不同于大语言模型的方向：仅依靠时空数据，我们能走多远？时空通用模型是否能像自然语言大模型一样存在？
具体来说，清华大学的研究团队致力于训练一个纯时空通用模型，能够模仿大语言模型（LLM）的两个关键特性：

1. 对丰富的时空数据具有强大的拓展能力；
2. 像大语言模型一样，展现出强大的通用性和泛化能力。
值得一提的是，纯时空模型背后的直觉是：在人类干预下，城市运转的各种时空数据存在通用规律，可以通过类似于GPT的方式进行训练。
通用时空建模的挑战

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

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

如何构建纯时空通用模型

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

• 具有在多样数据下的拓展能力；
• 自监督预训练充分捕捉复杂时空关联；
• 通过提示（prompt）灵活进行泛化。

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

• 灵活适应多样的时空数据特征：UniST 能够处理不同城市、不同领域的多样化时空数据，实现真正的统一和通用模型。无论是交通数据、人群流动数据还是城市资源分布数据，UniST 都能灵活应对 ，展现出强大的可拓展性。
• 高效的生成式预训练：通过巧妙设计的掩码策略，UniST 能够捕捉复杂的时空关系，实现全面多维度的时空建模。
• 时空知识引导的提示：利用知识引导的时空提示，UniST 能够对不同场景的内在和共享知识进行对齐和利用，提升预测性能。通过这种提示机制，UniST 可以在数据稀缺或全新的场景中依然保持高效的预测能力。

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

• 少样本学习：在训练数据有限的情况下，UniST 依然能够提供高精度的预测。
• 零样本学习：在模型从未见过的时空场景下，UniST 依然能实现出色的预测性能，甚至超过了大多数监督学习方法。
• 广泛适用性：在交通预测、人群流动预测、资源分配等多个任务中，UniST 均展示了其强大的预测能力和适用性。
  研究人员深入分析了提示（prompt）机制的作用。在时间记忆池中，他们研究了每个向量的记忆模式，根据该向量被数据集索引的权重高低，聚合数据集样本值在该向量上的结果。图5(a) 和图5(b) 展示了两个数据集（Crowd 和 TrafficSH）的结果。可以看到，提示机制中展现的记忆模式在不同的城市场景中表现出显著的一致性。这不仅证实了每个记忆向量都被很好地优化以记忆独特的时空模式，还证明了空间和时间记忆池在不同场景中的稳健性。
  进一步，研究团队分析了两个不同场景对记忆向量的利用情况（获得的时空提示）。具体来说，他们计算了在每个数据集上下文中每个向量的平均注意力权重。图6(c) 和图6(d) 展示了比较结果。两个数据集的注意力权重分布显示出明显的不同。观察到的注意力权重分布的独特性表明，模型能够根据输入数据的特征动态调整其关注的记忆模式。这种动态调整注意力权重的能力增强了 UniST 模型在不同数据集上的适配性和泛化性。
  这些实验结果表明，UniST 在提示机制的帮助下，能够在不同的时空场景中有效地捕捉和利用重要的时空关系，这种创新方法使得 UniST 在应对复杂多变的时空数据时表现出色，展示了其强大的适应能力和广泛的应用潜力。
  结语
  在这项工作中，研究者解决了一个重要问题，即构建一个用于城市时空预测的通用模型 UniST。
  UniST 通过整合多个城市和多领域丰富的时空数据，利用基于 Transformer 的架构，以及独特的时空掩码策略（mask）和知识引导的时空提示（prompt），实现了对城市多样化时空动态性的统一建模和精准预测。在实验中，UniST 展示了其在交通管理、资源优化等多个城市应用场景中的卓越表现，尤其是在跨场景零样本预测（zero-shot）中，其性能超过了少样本（few-shot）基线方法，展现出强大的通用性和泛化能力。
  研究团队将继续探索 UniST 的潜力，期待未来的研究能够进一步提升模型的性能和适应性，推动城市时空进入一个更加智能和高效的新阶段。