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

城市时空的预测，迎来GPT时刻。

清华大学电子系城市科学与计算研究中心推出了第一个无需自然语言的纯时空通用模型——UniST，首次展示了纯时空模型本身的通用性和可扩展性，研究成果已被KDD2024接收。

研究团队利用超过20个时空数据集、1.3亿+个时空样本点，构建了涵盖多个城市、不同领域、空间划分和时间分辨率等维度的城市时空数据，构建并训练了「one-for-all」的时空通用模型——UniST。

该模型是目前覆盖范围最广、统一性最强的城市时空通用模型。值得一提的是，UniST相较于当前的大语言模型具有更轻量级的优势，仅用20M的参数规模就展现出很强的零样本学习能力。

△图1. UniST实现时空通用建模（one-for-all）

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

该成果的论文、代码和数据均已公开，供研究和应用者使用。

纯时空模型的逆袭

时空预测在城市中无处不在，它不仅关注交通和人群的流动，还涉及资源分配、节能减排，公共卫生等多个场景。然而时空预测模型往往需要处理复杂且动态的时空关联，因此建模难度较大。传统的AI方法需要大量的训练数据和领域知识，且一般只能针对特定的数据集进行模型训练，不同时空场景就需要训练多个模型，这在城市数据不足的情况下显得尤为困难。

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

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

1. 对丰富的时空数据具有强大的拓展能力；
2. 像大语言模型一样，展现出强大的通用性和泛化能力。

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

通用时空建模的挑战

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

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

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

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

如何构建纯时空通用模型

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

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

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

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

△图2. UniST整体架构：时空预训练和知识引导的提示微调

时空数据的序列建模

为了有效处理不同来源、不同特征的时空数据，UniST提出了一种名为「时空序列」的建模方法。具体来说，时空数据首先被表示为一个四维张量：T×C×H×W，其中 T 表示时间段数量， C 表示变量数量， H 和 W 分别表示空间划分中的纬度和经度网格数。

为了统一处理不同形状的时空数据，UniST引入了时空编码器，将这些四维张量转换为小的三维向量，然后按照位置展开成序列。将丰富的时空数据表征为「时空序列」的通用格式后，就可以利用Transformer强大的序列建模能力，进行模型的训练，全面捕捉复杂的时空关系。

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

• 随机掩码（Random Masking）：类似于MAE中的随机掩码策略，通过随机遮蔽时空数据块来捕捉细粒度的时空关系。
• 管状掩码（Tube Masking）：模拟某些空间单元在所有时间段内的数据缺失情况，提升模型的空间外推能力。
• 块状掩码（Block Masking）：一种更具挑战性的掩码方式，通过遮蔽整个空间单元块在所有时间段内的数据，增强模型在有限上下文信息下的空间迁移能力。
• 时间掩码（Temporal Masking）：遮蔽未来的数据，仅依赖历史信息进行重建，旨在提升模型捕捉从过去到未来时间依赖关系的能力。

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

知识引导的时空提示

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

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

• 空间临近性：临近的空间单元可能相互影响；
• 空间层次结构：城市结构的层次组织会影响时空动态；
• 时间临近性：近期的动态会影响未来结果；
• 时间周期性：每天或每周的相似模式会影响未来周期性结果。

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

△图3. 时空提示网络

实验结果

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

• 少样本学习：在训练数据有限的情况下，UniST依然能够提供高精度的预测。
• 零样本学习：在模型从未见过的时空场景下，UniST依然能实现出色的预测性能，甚至超过了大多数监督学习方法。
• 广泛适用性：在交通预测、人群流动预测、资源分配等多个任务中，UniST均展示了其强大的预测能力和适用性。

△图4. 多个数据集与基线模型预测性能对比

△图5. （a）少样本场景（b）零样本场景性能

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

△图6. 记忆向量模式在不同数据集对比（一致性高）

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

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

△图7. 不同数据集时空提示结果对比（差异性大）

结语

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

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

论文地址：https://arxiv.org/abs/2402.11838 代码和数据开源地址：https://github.com/tsinghua-fib-lab/UniST