能不能有一種通用的圖模型——
它既能夠根據分子結構預測毒性，又能夠給出社交網路的朋友推薦？
或既能預測不同作者的論文引用，還可以發現基因網路中的人類老化機制？
你還真別說，被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://www.php.cn/link/dd4729902a3476b2bc9675e3530a852chttps://github.com/ LechengKong/OneForAll
