聲明:本文來自于微信公眾號 量子位 （ID:QbitAI），作者:豐色，授權(quán)轉(zhuǎn)載發(fā)布。
能不能有一種通用的圖模型——
它既能夠根據(jù)分子結(jié)構(gòu)預(yù)測毒性，又能夠給出社交網(wǎng)絡(luò)的朋友推薦?
或者既能預(yù)測不同作者的論文引用，還可以發(fā)現(xiàn)基因網(wǎng)絡(luò)中的人類衰老機制?
你還真別說，被ICLR2024接收為Spotlight的“One for All（OFA）”框架就實現(xiàn)了這個“精髓”。
它由圣路易斯華盛頓大學(xué)陳一昕教授團隊、北京大學(xué)張牧涵以及京東研究院陶大程等研究者們聯(lián)合提出。
作為圖領(lǐng)域首個通用框架，OFA實現(xiàn)了訓(xùn)練單一GNN模型即可解決圖領(lǐng)域內(nèi)任意數(shù)據(jù)集、任意任務(wù)類型、任意場景的分類任務(wù)。
具體如何實現(xiàn)，以下為作者投稿。
圖領(lǐng)域通用模型設(shè)計面臨三大難
設(shè)計一個通用的基礎(chǔ)模型來解決多種任務(wù)是人工智能領(lǐng)域的一個長期目標。近年來，基礎(chǔ)大語言模型（LLMs）在處理自然語言任務(wù)方面表現(xiàn)出色。
然而，在圖領(lǐng)域，雖然圖神經(jīng)網(wǎng)絡(luò)（GNNs）在不同的圖數(shù)據(jù)中都有著不俗的表現(xiàn)，但如何設(shè)計與訓(xùn)練一個能同時處理多種圖任務(wù)的基礎(chǔ)圖模型依然前路茫茫。
與自然語言領(lǐng)域相比，圖領(lǐng)域的通用模型設(shè)計面臨著許多獨有的困難。
首先，區(qū)別于自然語言，不同的圖數(shù)據(jù)有著截然不同的屬性與分布。
比如分子圖描述了多個原子如何通過不同的作用力關(guān)系形成不同的化學(xué)物質(zhì)。而引用關(guān)系圖則描述了文章與文章之間相互引用的關(guān)系網(wǎng)。
這些不同的圖數(shù)據(jù)很難被統(tǒng)一在一個訓(xùn)練框架下。
其次，不同于LLMs中所有任務(wù)都可以被轉(zhuǎn)化成統(tǒng)一的下文生成任務(wù)，圖任務(wù)包含了多種子任務(wù)，比如節(jié)點任務(wù)，鏈路任務(wù)，全圖任務(wù)等。
不同的子任務(wù)通常需要不同的任務(wù)表示形式與不同的圖模型。
最后，大語言模型的成功離不開通過提示范式而實現(xiàn)的上下文學(xué)習(xí)（in-context learning）。
在大語言模型中，提示范式通常為對于下游任務(wù)的可讀文字描述。
但是對于非結(jié)構(gòu)化且難以用語言描述的圖數(shù)據(jù)，如何設(shè)計有效的圖提示范式來實現(xiàn)in-context learning依然是個未解之謎。
用“文本圖”概念等來解決
下圖給出了OFA的整體框架:
具體而言，OFA的團隊通過巧妙的設(shè)計來解決上述所提到的三個主要問題。
對于不同圖數(shù)據(jù)屬性與分布不同的問題，OFA通過提出文本圖（Text-Attributed Graph， TAGs）的概念來統(tǒng)一所有圖數(shù)據(jù)。利用文本圖，OFA將所有的圖數(shù)據(jù)中的節(jié)點信息與邊信息用統(tǒng)一的自然語言框架來描述，具體如下圖所示:
接著，OFA通過單一LLM模型對所有數(shù)據(jù)中的文本進行表示學(xué)習(xí)得到其嵌入向量。
這些嵌入向量將作為圖模型的輸入特征。這樣，來自不同領(lǐng)域的圖數(shù)據(jù)將被映射到相同的特征空間，使得訓(xùn)練一個統(tǒng)一的GNN模型可行。
OFA收集了9個來自不同領(lǐng)域，不同規(guī)模的圖數(shù)據(jù)集，包括引用關(guān)系圖，Web鏈接圖，知識圖譜，分子圖， 如下圖所示:
此外，OFA提出Nodes-of-Interest（NOI）子圖與NOI提示節(jié)點(NOI Prompt Node)來統(tǒng)一圖領(lǐng)域內(nèi)不同的子任務(wù)類型。這里NOI代表參與到相應(yīng)任務(wù)的一組目標節(jié)點。
比如，在節(jié)點預(yù)測任務(wù)中，NOI是指需要預(yù)測的單個節(jié)點;而在鏈路任務(wù)中，NOI包括需要預(yù)測鏈路的兩個節(jié)點。NOI子圖是指圍繞著這些NOI節(jié)點擴展出的一個包含h-hop鄰域的子圖。
然后，NOI提示節(jié)點為一個新引入的節(jié)點類型，直接連接到所有的NOI上。
重要的是，每個NOI提示節(jié)點包含了當(dāng)前任務(wù)的描述信息，這些信息以自然語言的形式存在，并和文本圖被同一個LLM所表示。
由于NOI中節(jié)點所包含的信息在經(jīng)過GNNs的消息傳遞后將被NOI提示節(jié)點所收集，GNN模型僅需通過NOI提示節(jié)點來進行預(yù)測。
這樣，所有不同的任務(wù)類型將擁有統(tǒng)一的任務(wù)表示。具體實例如下圖所示:
最后，為了實現(xiàn)圖領(lǐng)域的in-context learning，OFA引入統(tǒng)一的提示子圖。
在一個有監(jiān)督的k-way分類任務(wù)場景下，這個提示子圖包含了兩類節(jié)點:一類是上文提到的NOI提示節(jié)點，另一類是代表k個不同類別的類別節(jié)點（Class Node）。
每個類別節(jié)點的文本將描述此類別的相關(guān)信息。
NOI提示節(jié)點將會單向連接到所有類別節(jié)點上。通過這個方式構(gòu)建好的圖將被輸入進圖神經(jīng)網(wǎng)路模型進行消息傳遞與學(xué)習(xí)。
最終，OFA將對每個類別節(jié)點分別進行二分類任務(wù)，并取概率最高的類別節(jié)點作為最終的預(yù)測結(jié)果。
由于類別信息存在于提示子圖中，即使遇到全新的分類問題，OFA通過構(gòu)建相應(yīng)的提示子圖即可直接進行預(yù)測而無需任何微調(diào)，從而實現(xiàn)了零樣本學(xué)習(xí)。
對于少樣本學(xué)習(xí)場景，一個分類任務(wù)將包含一個query輸入圖和多個support輸入圖，OFA的提示圖范式會將每個support輸入圖的NOI提示節(jié)點與其所對應(yīng)的類別節(jié)點相連，同時將query輸入圖的NOI提示節(jié)點與所有類別節(jié)點相連。
后續(xù)的預(yù)測步驟與上文所述一致。這樣每個類別節(jié)點將會額外得到support輸入圖的信息，從而在統(tǒng)一的范式下實現(xiàn)少樣本學(xué)習(xí)。
OFA的主要貢獻總結(jié)如下:
統(tǒng)一的圖數(shù)據(jù)分布:通過提出文本圖并用LLM轉(zhuǎn)化文本信息，OFA實現(xiàn)了圖數(shù)據(jù)的分布對齊與統(tǒng)一。
統(tǒng)一的圖任務(wù)形式:通過NOI子圖與NOI提示節(jié)點，OFA實現(xiàn)了多種圖領(lǐng)域子任務(wù)的統(tǒng)一表示。
統(tǒng)一的圖提示范式:通過提出新穎的圖提示范式，OFA實現(xiàn)了圖領(lǐng)域內(nèi)的多場景in-context learning。
超強泛化能力
文章在所收集的9個數(shù)據(jù)集上對OFA框架進行了測試，這些測試覆蓋了在有監(jiān)督學(xué)習(xí)場景下的十種不同任務(wù)，包括節(jié)點預(yù)測、鏈路預(yù)測和圖分類。
實驗的目的是驗證單一的OFA模型處理多任務(wù)的能力，其中作者對比使用不同LLM（OFA-{LLM}）和每個任務(wù)訓(xùn)練單獨模型(OFA-ind-{LLM})的效果。
比較結(jié)果如下表所示:
可以看到，基于OFA強大的泛化能力，一個單獨的圖模型（OFA-st，OFA-e5，OFA-llama2-7b，OFA-llama2-13b）即能夠在所有的任務(wù)上都具有與傳統(tǒng)的單獨訓(xùn)練模型(GCN， GAT， OFA-ind-st)相近或更好的表現(xiàn)。
同時，使用更強大的LLM可以帶來一定的性能提升。文章進一步繪制了訓(xùn)練完成的OFA模型對于不同任務(wù)的NOI提示節(jié)點的表示。
可以看到不同的任務(wù)被模型嵌入到不同的子空間，從而使得OFA可以對于不同的任務(wù)進行分別的學(xué)習(xí)而不會相互影響。
在少樣本以及零樣本的場景下，OFA在ogbn-arxiv（引用關(guān)系圖），F(xiàn)B15K237(知識圖譜)以及Chemble(分子圖)上使用單一模型進行預(yù)訓(xùn)練，并測試其在不同下游任務(wù)及數(shù)據(jù)集上的表現(xiàn)。結(jié)果如下:
可以看到，即使在零樣本場景下，OFA依舊可以取得不錯的效果。綜合來看，實驗結(jié)果很好的驗證了OFA強大的通用性能以及其作為圖領(lǐng)域基礎(chǔ)模型的潛力。
更多研究細節(jié)，可參考原論文。
地址:
https://arxiv.org/abs/2310.00149
https://github.com/LechengKong/OneForAll
—完—
