清華推出首個通用城市時空預測模型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)的兩個關鍵特性:
- 對豐富的時空數據具有強大的拓展能力;
- 像大語言模型一樣,展現出強大的通用性和泛化能力。
值得一提的是,純時空模型背后的直覺是:在人類干預下,城市運轉中產生的各種時空數據存在通用規律,可以通過類似于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
本文轉自 量子位 ,作者:量子位

















