破解三大數據庫AI難題!北大&亞馬遜推出全球首個圖中心RDB基礎模型
在企業系統和科學研究中普遍存在、結構復雜的關系型數據庫(Relational DataBase, RDB)場景中,基礎模型的探索仍處于早期階段。
這是因為RDB中的多表交互和異質特征,使傳統通用大模型在此類結構化環境下難以直接發揮效能。
基于此,北京大學張牧涵團隊聯合亞馬遜云科技共同提出了Griffin:一個具有開創性的、以圖為中心的RDB基礎模型。
Griffin將RDB視為動態異質圖進行建模與推理,通過在超過1.5億行的表格數據上進行預訓練和監督微調,構建了一個具備可遷移性與強泛化能力的基礎模型,相關成果已被國際頂級會議ICML 2025正式接收。
挑戰:復雜的表間關系及豐富的表內語義信息
關系數據庫通過明確的模式(Schema)定義數據結構,廣泛服務于金融、電商、科研、物流、政府信息系統等關鍵領域,是現代信息社會的核心數字基礎設施。
根據市場預測,到2028年全球數據庫管理系統(DBMS)市場將超過1330億美元。
然而,RDB智能建模所面臨的挑戰極為復雜,集中體現在以下三方面:
- 拓撲結構高度復雜數據以多表形式存儲,并通過主鍵外鍵等約束關系構成復雜的圖結構,傳統單表范式難以捕捉全局上下文。
- 特征高度異構表內字段涵蓋文本、數值、類別、時間序列等多種類型,信息表現形態各異,模型需具備統一表征能力。
- 深層次語義關系表內外存在豐富的顯性與隱性邏輯關系,對模型的關系理解和推理能力構成巨大挑戰。
上圖展示了一個典型的RDB,綠色的Purchase Table記錄了交易數據(每一行包括用戶ID、購買的商品ID、用戶對商品的評分、以及購買日期)。而每一行又可通過User ID這一外鍵鏈接到User Table里的對應行,或通過Item ID這一外鍵鏈接到Product Table里的對應行,來獲取用戶或商品的特定信息。
相比普通表格(單表)數據,RDB往往具有非常復雜的表間關系以及豐富的表內語義信息,對建模和基礎模型訓練提出了挑戰。同時,社區長期缺乏能真實反映生產場景的標準化基準。
諸如4DBInfer(arXiv:2404.18209)等數據集正緩慢填補空白,為新模型(包括Griffin)提供了統一的評測土壤。
方法:以圖為中心的數據庫建模
Griffin的核心思想是將關系型數據庫整體抽象為時序異質圖,再在此圖上進行統一的編碼、消息傳遞和解碼,以此捕捉跨表、跨時間的深層依賴。具體來說,它的創新設計可以拆解為以下幾點:
RDB的數據建模:結構化圖表示與時序感知
首先,Griffin把每張數據表中的一行記錄映射為圖中的一個節點,而主鍵-外鍵(PK-FK)約束被建模為帶類型的有向邊。這樣,原本分散在多張表中的記錄就組成了一張異質圖,其節點/邊類型天然反映了模式信息。
為了避免未來信息泄漏并符合生產預測任務的因果約束,模型在訓練和推理時會圍繞目標節點采樣“局部時序子圖”:僅納入時間戳早于目標節點的鄰域。
該采樣流程借鑒了4DBInfer等基準的成熟做法,可在保證效率的同時顯式注入時間方向。
統一數據編碼器:異構信息的規范化表征
RDB中既包含文本/類別字段,也有數值、時間序列等多模態特征。Griffin設計了一套統一編碼機制,把不同類型轉換為同一語義空間中的向量:
- 類別&文本:先將類別值映射為其自然語言描述,再同原生文本一起輸入經過預訓練的文本編碼器(如Nomic Embeddings),得到語義豐富的高維嵌入。
- 數值:歸一化后的數值輸入給預訓練的浮點編碼器(ENC)。ENC與其配對的解碼器DEC通過聯合重構任務訓練:編碼后必須能夠無損地解碼回原始浮點值,重構誤差被最小化后這兩個組件參數即被凍結。
- 元數據&任務上下文:表名、列名以及邊類型被同樣送入文本編碼器;同時,根據當前預測目標列名生成的任務描述會在后續所有層次參與注意力計算,指導模型聚焦目標。
經過上述步驟,原始多態信息被規范化為一組高語義的向量,為后續圖消息傳遞奠定基礎。
先進MPNN架構:深度關系推理網絡
統一編碼后的圖被送入Griffin定制的Message Passing Neural Network (MPNN),其核心由兩個互補模塊構成:
交叉注意力(Cross-Attention)列內聚合:對每個節點,模型利用當前節點嵌入和任務嵌入生成查詢向量,與列元數據和列特征進行交互,動態評估不同列對當前任務的重要性并加權聚合。該設計天然滿足列置換不變性,且可處理列數可變的表。
層級聚合(Hierarchical Aggregation)跨表推理:在消息傳遞的每一層,先對同一邊類型的鄰居消息做均值聚合,再在不同邊類型間做最大池化。這種兩階段層級策略提升了模型在處理具有復雜拓撲結構和多變鄰居數量的表間關聯時的穩定性。
通過多層迭代,MPNN能夠捕獲從近鄰到遠程的復合依賴,為下游任務提供信息豐富的節點表征。
統一任務解碼器:多任務輸出的一體化方案
MPNN輸出的節點向量隨后進入統一解碼器,使Griffin能夠在不改動架構的前提下同時處理多種預測任務。
分類任務:把候選類別標簽本身的文本嵌入當作可學習的動態分類頭,與節點向量做內積得到概率分布,能夠拓展到可變類比數量的任務。
回歸任務:直接將節點向量輸入預訓練DEC,反解得到最終的預測數值。
訓練:三階段優化方案
Griffin通過“自監督預訓練→聯合監督微調→下游任務微調”的三級管線,逐步注入從通用表格語義到特定RDB任務知識的能力層次。
第一階段:補全預訓練(Completion Pretraining)
Griffin首先在海量且多樣化的單表數據集上進行自監督學習,任務形式類似“完形填空”。模型根據一行數據中已知列信息來預測被遮蔽單元的嵌入表示,并最小化預測嵌入與真實嵌入間的余弦距離,從而建立對表格結構與語義的基礎理解。
第二階段:聯合監督微調(Joint Supervised Fine-Tuning, SFT)
在完成自監督預訓練后,Griffin使用單表任務或RDB任務的數據集進行監督微調,使模型進一步貼合真實場景中的預測需求與數據特性。
第三階段:下游任務微調(Downstream Task Fine-Tuning)
最后,經過預訓練和SFT的Griffin會針對具體下游RDB基準任務進行精細化微調,以在特定應用場景中取得最佳性能。
驗證:三階段訓練的優越性
為全面評估各訓練階段對模型性能的具體貢獻,對Griffin的三個關鍵變體進行深入分析:Griffin-unpretrained(僅采用Griffin的基礎架構,未經任何預訓練)、Griffin-pretrained(僅進行單表預訓練及單表SFT)以及Griffin-RDB-SFT(經歷完整的三階段訓練流程)。
上圖比較了四個 GNN 基線模型、四個使用 DFS 的單表基線模型以及兩個 Griffin 變體的性能,每個模型均在單個任務上進行了微調。
最左側的子圖展示了所有任務上的平均排名,其余子圖按評估指標對任務進行分組,結果相應地進行了平均。
通過系統實驗驗證了Griffin在架構設計和預訓練策略上的有效性,發現Griffin在多個RDB基準測試(如4DBInfer和RelBench)中表現優異,并進一步分析了其在少樣本場景下的跨任務遷移能力與數據領域間關系的影響。
Griffin的核心優勢可歸納為以下三點:
1.強大的基礎架構性能
即便完全未預訓練(Griffin-unpretrained),憑借統一編碼、交叉注意力和層級化MPNN等設計,模型在各下游RDB任務微調后的表現仍優于GNN基線方法及結合深度特征合成(DFS)的傳統單表模型,體現了架構本身的先進性。
2.單表預訓練的普適性增益
僅在大規模、多樣化單表數據上完成預訓練的Griffin-pretrained,相較未預訓練版本取得性能提升,驗證單表場景中學習到的知識可遷移至復雜的RDB任務,增強模型泛化能力。
3.RDB-SFT驅動的遷移
當進一步在針對性RDB數據上進行監督微調(Griffin-RDB-SFT)后,模型在一定情況下展現出跨任務遷移能力,尤其在小樣本場景下更為突出,取決于以下2個因素:
- 數據相似性:SFT數據與目標任務領域具有較高相似性(例如,同為電商領域的跨任務遷移),模型性能將獲得提升;
- 數據多樣性:在更多樣化的SFT數據上訓練(例如,使用包含體育、社交、醫療等多個其他領域的混合數據進行SFT,再遷移至電商領域任務),同樣能有效提升模型性能。
論文鏈接:https://arxiv.org/abs/2505.05568代碼鏈接:https://github.com/yanxwb/griffin