能不能有一種通用的圖模型——
它既能夠根據分子結構預測毒性，又能夠給出社交網絡的朋友推薦？
或者既能預測不同作者的論文引用，還可以發現基因網絡中的人類衰老機制？
你還真别說，被ICLR 2024接收爲 Spotlight 的 "One for All（OFA）" 框架就實現了這個 " 精髓 "。
它由聖路易斯華盛頓大學陳一昕教授團隊、北京大學張牧涵以及京東研究院陶大程等研究者們聯合提出。
作爲圖領域首個通用框架，OFA 實現了訓練單一 GNN 模型即可解決圖領域内任意數據集、任意任務類型、任意場景的分類任務。
具體如何實現，以下爲作者投稿。
圖領域通用模型設計面臨三大難
設計一個通用的基礎模型來解決多種任務是人工智能領域的一個長期目标。近年來，基礎大語言模型（LLMs）在處理自然語言任務方面表現出色。
然而，在圖領域，雖然圖神經網絡（GNNs）在不同的圖數據中都有着不俗的表現，但如何設計與訓練一個能同時處理多種圖任務的基礎圖模型依然前路茫茫。
與自然語言領域相比，圖領域的通用模型設計面臨着許多獨有的困難。
首先，區别于自然語言，不同的圖數據有着截然不同的屬性與分布。
比如分子圖描述了多個原子如何通過不同的作用力關系形成不同的化學物質。而引用關系圖則描述了文章與文章之間相互引用的關系網。
這些不同的圖數據很難被統一在一個訓練框架下。
其次，不同于 LLMs 中所有任務都可以被轉化成統一的下文生成任務，圖任務包含了多種子任務，比如節點任務，鏈路任務，全圖任務等。
不同的子任務通常需要不同的任務表示形式與不同的圖模型。
最後，大語言模型的成功離不開通過提示範式而實現的上下文學習（in-context learning）。
在大語言模型中，提示範式通常爲對于下遊任務的可讀文字描述。
但是對于非結構化且難以用語言描述的圖數據，如何設計有效的圖提示範式來實現 in-context learning 依然是個未解之謎。
用 " 文本圖 " 概念等來解決
下圖給出了 OFA 的整體框架：
具體而言，OFA 的團隊通過巧妙的設計來解決上述所提到的三個主要問題。
對于不同圖數據屬性與分布不同的問題，OFA 通過提出文本圖（Text-Attributed Graph, TAGs） 的概念來統一所有圖數據。利用文本圖，OFA 将所有的圖數據中的節點信息與邊信息用統一的自然語言框架來描述，具體如下圖所示：
接着，OFA 通過單一 LLM 模型對所有數據中的文本進行表示學習得到其嵌入向量。
這些嵌入向量将作爲圖模型的輸入特征。這樣，來自不同領域的圖數據将被映射到相同的特征空間，使得訓練一個統一的 GNN 模型可行。
OFA 收集了 9 個來自不同領域，不同規模的圖數據集，包括引用關系圖，Web 鏈接圖，知識圖譜，分子圖， 如下圖所示：
此外，OFA 提出 Nodes-of-Interest（NOI）子圖與 NOI 提示節點 （NOI Prompt Node）來統一圖領域内不同的子任務類型。這裏 NOI 代表參與到相應任務的一組目标節點。
比如，在節點預測任務中，NOI 是指需要預測的單個節點；而在鏈路任務中，NOI 包括需要預測鏈路的兩個節點。NOI 子圖是指圍繞着這些 NOI 節點擴展出的一個包含 h-hop 鄰域的子圖。
然後，NOI 提示節點爲一個新引入的節點類型，直接連接到所有的 NOI 上。
重要的是，每個 NOI 提示節點包含了當前任務的描述信息，這些信息以自然語言的形式存在，并和文本圖被同一個 LLM 所表示。
由于 NOI 中節點所包含的信息在經過 GNNs 的消息傳遞後将被 NOI 提示節點所收集，GNN 模型僅需通過 NOI 提示節點來進行預測。
這樣，所有不同的任務類型将擁有統一的任務表示。具體實例如下圖所示：
最後，爲了實現圖領域的 in-context learning，OFA 引入統一的提示子圖。
在一個有監督的 k-way 分類任務場景下，這個提示子圖包含了兩類節點：一類是上文提到的 NOI 提示節點，另一類是代表 k 個不同類别的類别節點 ( Class Node ) 。
每個類别節點的文本将描述此類别的相關信息。
NOI 提示節點将會單向連接到所有類别節點上。通過這個方式構建好的圖将被輸入進圖神經網路模型進行消息傳遞與學習。
最終，OFA 将對每個類别節點分别進行二分類任務，并取概率最高的類别節點作爲最終的預測結果。
由于類别信息存在于提示子圖中，即使遇到全新的分類問題，OFA 通過構建相應的提示子圖即可直接進行預測而無需任何微調，從而實現了零樣本學習。
對于少樣本學習場景，一個分類任務将包含一個 query 輸入圖和多個 support 輸入圖，OFA 的提示圖範式會将每個 support 輸入圖的 NOI 提示節點與其所對應的類别節點相連，同時将 query 輸入圖的 NOI 提示節點與所有類别節點相連。
後續的預測步驟與上文所述一緻。這樣每個類别節點将會額外得到 support 輸入圖的信息，從而在統一的範式下實現少樣本學習。
OFA 的主要貢獻總結如下：
統一的圖數據分布：通過提出文本圖并用 LLM 轉化文本信息，OFA 實現了圖數據的分布對齊與統一。
統一的圖任務形式：通過 NOI 子圖與 NOI 提示節點，OFA 實現了多種圖領域子任務的統一表示。
統一的圖提示範式：通過提出新穎的圖提示範式，OFA 實現了圖領域内的多場景 in-context learning。
超強泛化能力
文章在所收集的 9 個數據集上對 OFA 框架進行了測試，這些測試覆蓋了在有監督學習場景下的十種不同任務，包括節點預測、鏈路預測和圖分類。
實驗的目的是驗證單一的 OFA 模型處理多任務的能力，其中作者對比使用不同 LLM（OFA-{LLM}）和每個任務訓練單獨模型（OFA-ind-{LLM}）的效果。
比較結果如下表所示：
可以看到，基于 OFA 強大的泛化能力，一個單獨的圖模型（OFA-st，OFA-e5，OFA-llama2-7b，OFA-llama2-13b）即能夠在所有的任務上都具有與傳統的單獨訓練模型（GCN, GAT, OFA-ind-st）相近或更好的表現。
同時，使用更強大的 LLM 可以帶來一定的性能提升。文章進一步繪制了訓練完成的 OFA 模型對于不同任務的 NOI 提示節點的表示。
可以看到不同的任務被模型嵌入到不同的子空間，從而使得 OFA 可以對于不同的任務進行分别的學習而不會相互影響。
在少樣本以及零樣本的場景下，OFA 在 ogbn-arxiv（引用關系圖），FB15K237（知識圖譜）以及 Chemble（分子圖）上使用單一模型進行預訓練，并測試其在不同下遊任務及數據集上的表現。結果如下：
可以看到，即使在零樣本場景下，OFA 依舊可以取得不錯的效果。綜合來看，實驗結果很好的驗證了 OFA 強大的通用性能以及其作爲圖領域基礎模型的潛力。
更多研究細節，可參考原論文。
地址：
https://arxiv.org/abs/2310.00149
https://github.com/LechengKong/OneForAll
— 完 —
點這裏關注我，記得标星哦～
一鍵三連「分享」、「點贊」和「在看」
科技前沿進展日日相見 ~
