能不能有一種通用的圖模型——

它既能夠根據分子結構預測毒性，又能夠給出社交網絡的朋友推薦？

或者既能預測不同作者的論文引用，還可以發現基因網絡中的人類衰老機制？

你還真別說，被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.00149https://github.com/LechengKong/OneForAll