本文詳細介紹了一篇發表于人工智能頂級會議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
一、研究背景
時空數據的時空特性理解對智能交通、城市規劃等動態系統至關重要,其中時空預測與補全是核心任務。精準的預測可優化交通管理與出行規劃,補全則能支撐時空模式分析及下游任務。然而,現有方法存在以下核心挑戰:
- 時空相關性建模不足:傳統PLM缺乏對時間序列動態性、空間拓撲結構(如節點連通性)及高階時空交互的顯式建模;
- 任務適配性差:現有PLM-based方法多聚焦預測任務,補全能力不足,且難以同時支持多任務;
- 計算效率瓶頸:PLM的高維嵌入與大量tokens導致處理時空數據時訓練和推理成本激增;
該論文的主要貢獻如下。
- 提出STD-PLM統一框架。設計基于預訓練語言模型的時空預測與補全統一框架,實現對時空數據的空間和時間特性的聯合建模,同時支持時空預測、補全、零樣本及少樣本學習任務。
- 設計拓撲感知節點嵌入,利用圖拉普拉斯矩陣的特征向量捕捉空間拓撲結構,使PLM能感知節點間的靜態關系與動態狀態。
- 開發空間和時間標記器,從時空雙維度將數據轉換為PLM可處理的標記序列,顯式建模時間相關性、空間連通性及高階時空交互。
- 設計沙漏注意力模塊(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對比了以下變體:
- 完整模型(STD-PLM):包含時空嵌入、時空標記器、沙漏注意力模塊(SGA)及PLM微調策略;
- w/o 時空嵌入:移除拓撲感知節點嵌入與時間周期嵌入,僅用原始數據特征;
- w/o 時空標記器:直接將原始數據輸入PLM,不進行時空維度解耦;
- w/o SGA 模塊:保留節點級標記,不進行區域壓縮;
- w/o 掩碼標記:在補全任務中移除缺失模式編碼。
表 3 消融分析
消融研究證實了STD-PLM的每個組件均不可替代:時空嵌入注入領域先驗,時空標記器激活PLM時空理解,SGA模塊優化計算效率,掩碼標記增強補全針對性。這一結果不僅驗證了模型設計的合理性,也為后續時空PLM的組件優化提供了方向。
四、總結
本文介紹了STD-PLM模型,旨在通過預訓練語言模型解決時空數據的預測與補全問題。模型通過時空嵌入模塊生成拓撲感知的節點嵌入與周期感知的時間嵌入,融合空間結構與時間規律;利用時空標記器解耦時空維度,構建包含動態狀態、靜態先驗及缺失掩碼的節點級標記與全局時間趨勢標記;通過沙漏注意力模塊壓縮節點級標記為區域級以降低預訓練語言模型的計算復雜度,同時通過約束損失保留拓撲信息;采用部分凍結預訓練語言模型結合低秩自適應?微調策略平衡預訓練知識與任務適配性。實驗表明,STD-PLM在交通數據集PEMS03/04/07/08的預測與補全任務中性能顯著優于傳統深度學習模型及其他PLM-based模型,尤其在70% 高缺失率場景下補全誤差降低超40%,消融研究驗證了各組件的必要性。該模型通過時空特性顯式建模、預訓練語言模型的推理能力與效率優化的結合,為時空智能任務提供了高效且泛化性強的解決方案,適用于智能交通、城市計算等數據稀缺或大規模場景。