AAAI 2025论文分享│STD-PLM：基于预训练语言模型的时空数据预测与补全方法

本文详细介绍了一篇发表于人工智能顶级会议AAAI 2025的论文《STD-PLM: Understanding Both Spatial and Temporal Properties of Spatial-Temporal Data with PLM》。该论文提出了一种基于预训练语言模型（Pre-trained Language Model‌，PLM）的时空预测与补全统一框架STD-PLM。通过显式设计的时空标记器和时空嵌入，STD-PLM 能够有效理解时空数据的空间和时间特性。此外，还引入沙漏注意力模块（Sandglass Attention.SGA），通过构建区域级空间标记显著降低计算成本。大量实验表明，STD-PLM表现出具有竞争力的性能。论文的工作表明，构建基于PLM的统一预训练时空模型具有广阔前景。

本推文由韩煦撰写，审核为邓镝。

原文链接：https://arxiv.org/html/2407.09096v4

代码链接：https://github.com/Hyheng/STD-PLM

一、研究背景

时空数据的时空特性理解对智能交通、城市规划等动态系统至关重要，其中时空预测与补全是核心任务。精准的预测可优化交通管理与出行规划，补全则能支撑时空模式分析及下游任务。然而，现有方法存在以下核心挑战：

1. 时空相关性建模不足：传统PLM缺乏对时间序列动态性、空间拓扑结构（如节点连通性）及高阶时空交互的显式建模；
2. 任务适配性差：现有PLM-based方法多聚焦预测任务，补全能力不足，且难以同时支持多任务；
3. 计算效率瓶颈：PLM的高维嵌入与大量tokens导致处理时空数据时训练和推理成本激增；

该论文的主要贡献如下。

1. 提出STD-PLM统一框架。设计基于预训练语言模型的时空预测与补全统一框架，实现对时空数据的空间和时间特性的联合建模，同时支持时空预测、补全、零样本及少样本学习任务。
2. 设计拓扑感知节点嵌入，利用图拉普拉斯矩阵的特征向量捕捉空间拓扑结构，使PLM能感知节点间的静态关系与动态状态。
3. 开发空间和时间标记器，从时空双维度将数据转换为PLM可处理的标记序列，显式建模时间相关性、空间连通性及高阶时空交互。
4. 设计沙漏注意力模块（SGA），通过区域级标记聚合减少计算量，结合约束损失函数强制模型学习符合原始图结构的特征，在保证性能的同时显著提升计算效率。

 二、方法

1、方法概述

图 1 STD-PLM 模型架构

图1展示了STD-PLM模型架构，图中将模型分为了四部分：模块A为节点和时间信息提供时空嵌入。时空标记器B从不同角度构建时间标记和空间标记。模块C基于B-2输出的节点级空间标记构建区域级空间标记。模块D将隐藏表示投影到目标输出。对于预训练语言模型，采用了部分冻结策略，并结合低秩适应技术对多头注意力、位置嵌入和层归一化层进行微调。

2、时空嵌入

拓扑感知节点嵌入：利用图拉普拉斯矩阵的特征向量捕捉空间节点的拓扑结构，通过选取前 K大特征值对应的特征向量，生成包含节点连通性信息的嵌入表示，其作用为捕捉节点在空间网络中的拓扑属性（如枢纽节点与边缘节点的差异），支持归纳式学习（即模型可泛化到未见过的图结构，如不同城市的交通网络）。

周期感知时间嵌入：结合 “小时级” 和 “星期级” 周期性粒度，构建时间嵌入字典，通过查表与拼接操作生成包含时间上下文的向量，捕捉数据的周期性波动模式。

3、时空标记器

空间标记器为每个节点生成包含 “静态属性”（节点嵌入与时间嵌入）、“动态状态”（历史数据特征）和 “缺失掩码” 的标记，支持 PLM 建模节点间的复杂空间依赖；

内在状态：拼接时间与节点嵌入，通过 MLP 提取静态时空特征（如节点在特定时段的拓扑角色）。

动态状态：直接对原始数据应用多层感知机（ Multilayer Perceptron‌，MLP），捕捉实时动态特征（如节点的流量、速度）。

掩码标记：将二进制掩码转换为标记，显式提示缺失位置（补全任务关键）。
融合操作：通过逐元素相加与层归一化，生成包含 “动态 + 静态 + 缺失” 信息的复合标记。

时间标记器通过全局聚合节点信息，生成 “整体状态” 与 “变化趋势” 标记，帮助 PLM 捕捉时空数据的宏观演化规律，避免仅依赖局部节点信息。

整体状态：融合全局平均状态与最后时刻的时间嵌入，体现当前系统整体状态。

整体趋势：融合全局趋势与时间嵌入，表征 “系统随时间的变化方向”

4、沙漏注意力模块

编码器：通过可学习查询矩阵将节点级空间标记压缩为更少的区域级标记，减少计算量的同时捕捉非成对、高阶时空相关性；

解码器：将区域级表示还原为节点级，结合约束损失函数（强制区域划分符合原始图结构），确保压缩过程中拓扑信息不丢失；

效率优化：通过稀疏注意力机制与区域级聚合，显著降低 PLM 处理大规模节点时的时间与内存开销。

三、实验

1.实验设置

数据集：在交通领域的四个真实数据集（PEMS03/04/07/08）上，STD-PLM通过多维度实验验证了其在时空预测、补全及少样本场景下的领先性能。

基线对比：选取LSTM、ASTGCN等 7 种传统深度学习模型，以及OFA、STGLLM等 3 种 PLM-based 模型作为基线；

训练策略：采用6:2:2的数据集划分，滑动窗口构建样本，使用AdaW优化器及LoRA技术微调 PLM 参数，平衡效率与性能。

2、与最先进方法的比较

基线模型有两种，包括六种传统深度学习模型和三种基于预训练语言模型的模型。

表 1 预测性能

表1展示了 STD-PLM 与六种传统深度学习模型（如LSTM、ASTGCN等）及三种PLM-based模型（如OFA、STGLLM等）在四个交通数据集上的预测性能对比，指标包括MAE（平均绝对误差）、RMSE（均方根误差）和MAPE（平均绝对百分比误差）。STD-PLM在所有数据集上实现最优或次优性能。

表 2 在PEMS08数据集上的补全性能

表2展示了STD-PLM与BRITS、E2GAN等4种基线模型在PEMS08数据集高缺失率（70%）场景下的补全性能，结果显示：在随机缺失（RM）和时空连续缺失（CM）两种模式中，STD-PLM均以显著优势领先——RM 场景下 MAE为14.36，CM 场景下MAE为22.69。其核心优势源于掩码标记显式定位缺失位置、时空嵌入融合拓扑与周期先验及PLM 生成能力支持长程推理，使模型在高缺失率下仍能结合全局趋势（如早晚高峰规律）与局部拓扑关联（生成合理补全值，突破了传统方法依赖局部观测或数据拟合的局限，为传感器故障、数据采集不全等实际场景提供了高效可靠的解决方案。

3、消融研究

消融研究通过移除STD-PLM的关键组件，验证各模块对性能的贡献。假设表3对比了以下变体：

1. 完整模型（STD-PLM）：包含时空嵌入、时空标记器、沙漏注意力模块（SGA）及PLM微调策略；
2. w/o 时空嵌入：移除拓扑感知节点嵌入与时间周期嵌入，仅用原始数据特征；
3. w/o 时空标记器：直接将原始数据输入PLM，不进行时空维度解耦；
4. w/o SGA 模块：保留节点级标记，不进行区域压缩；
5. w/o 掩码标记：在补全任务中移除缺失模式编码。

表 3 消融分析

消融研究证实了STD-PLM的每个组件均不可替代：时空嵌入注入领域先验，时空标记器激活PLM时空理解，SGA模块优化计算效率，掩码标记增强补全针对性。这一结果不仅验证了模型设计的合理性，也为后续时空PLM的组件优化提供了方向。

四、总结

本文介绍了STD-PLM模型，旨在通过预训练语言模型解决时空数据的预测与补全问题。模型通过时空嵌入模块生成拓扑感知的节点嵌入与周期感知的时间嵌入，融合空间结构与时间规律；利用时空标记器解耦时空维度，构建包含动态状态、静态先验及缺失掩码的节点级标记与全局时间趋势标记；通过沙漏注意力模块压缩节点级标记为区域级以降低预训练语言模型的计算复杂度，同时通过约束损失保留拓扑信息；采用部分冻结预训练语言模型结合低秩自适应‌微调策略平衡预训练知识与任务适配性。实验表明，STD-PLM在交通数据集PEMS03/04/07/08的预测与补全任务中性能显著优于传统深度学习模型及其他PLM-based模型，尤其在70% 高缺失率场景下补全误差降低超40%，消融研究验证了各组件的必要性。该模型通过时空特性显式建模、预训练语言模型的推理能力与效率优化的结合，为时空智能任务提供了高效且泛化性强的解决方案，适用于智能交通、城市计算等数据稀缺或大规模场景。