圖領(lǐng)域通用框架來(lái)了！入選ICLR'24 Spotlight，任意數(shù)據(jù)集、分類(lèi)問(wèn)題都可搞定

聲明:本文來(lái)自于微信公眾號(hào) 量子位（ID:QbitAI），作者:豐色，授權(quán)轉(zhuǎn)載發(fā)布。

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

它既能夠根據(jù)分子結(jié)構(gòu)預(yù)測(cè)毒性，又能夠給出社交網(wǎng)絡(luò)的朋友推薦?

或者既能預(yù)測(cè)不同作者的論文引用，還可以發(fā)現(xiàn)基因網(wǎng)絡(luò)中的人類(lèi)衰老機(jī)制?

你還真別說(shuō)，被ICLR2024接收為Spotlight的“One for All（OFA）”框架就實(shí)現(xiàn)了這個(gè)“精髓”。

它由圣路易斯華盛頓大學(xué)陳一昕教授團(tuán)隊(duì)、北京大學(xué)張牧涵以及京東研究院陶大程等研究者們聯(lián)合提出。

作為圖領(lǐng)域首個(gè)通用框架，OFA實(shí)現(xiàn)了訓(xùn)練單一GNN模型即可解決圖領(lǐng)域內(nèi)任意數(shù)據(jù)集、任意任務(wù)類(lèi)型、任意場(chǎng)景的分類(lèi)任務(wù)。

具體如何實(shí)現(xiàn)，以下為作者投稿。

圖領(lǐng)域通用模型設(shè)計(jì)面臨三大難

設(shè)計(jì)一個(gè)通用的基礎(chǔ)模型來(lái)解決多種任務(wù)是人工智能領(lǐng)域的一個(gè)長(zhǎng)期目標(biāo)。近年來(lái)，基礎(chǔ)大語(yǔ)言模型（LLMs）在處理自然語(yǔ)言任務(wù)方面表現(xiàn)出色。

然而，在圖領(lǐng)域，雖然圖神經(jīng)網(wǎng)絡(luò)（GNNs）在不同的圖數(shù)據(jù)中都有著不俗的表現(xiàn)，但如何設(shè)計(jì)與訓(xùn)練一個(gè)能同時(shí)處理多種圖任務(wù)的基礎(chǔ)圖模型依然前路茫茫。

與自然語(yǔ)言領(lǐng)域相比，圖領(lǐng)域的通用模型設(shè)計(jì)面臨著許多獨(dú)有的困難。

首先，區(qū)別于自然語(yǔ)言，不同的圖數(shù)據(jù)有著截然不同的屬性與分布。

比如分子圖描述了多個(gè)原子如何通過(guò)不同的作用力關(guān)系形成不同的化學(xué)物質(zhì)。而引用關(guān)系圖則描述了文章與文章之間相互引用的關(guān)系網(wǎng)。

這些不同的圖數(shù)據(jù)很難被統(tǒng)一在一個(gè)訓(xùn)練框架下。

其次，不同于LLMs中所有任務(wù)都可以被轉(zhuǎn)化成統(tǒng)一的下文生成任務(wù)，圖任務(wù)包含了多種子任務(wù)，比如節(jié)點(diǎn)任務(wù)，鏈路任務(wù)，全圖任務(wù)等。

不同的子任務(wù)通常需要不同的任務(wù)表示形式與不同的圖模型。

最后，大語(yǔ)言模型的成功離不開(kāi)通過(guò)提示范式而實(shí)現(xiàn)的上下文學(xué)習(xí)（in-context learning）。

在大語(yǔ)言模型中，提示范式通常為對(duì)于下游任務(wù)的可讀文字描述。

但是對(duì)于非結(jié)構(gòu)化且難以用語(yǔ)言描述的圖數(shù)據(jù)，如何設(shè)計(jì)有效的圖提示范式來(lái)實(shí)現(xiàn)in-context learning依然是個(gè)未解之謎。

用“文本圖”概念等來(lái)解決

下圖給出了OFA的整體框架:

具體而言，OFA的團(tuán)隊(duì)通過(guò)巧妙的設(shè)計(jì)來(lái)解決上述所提到的三個(gè)主要問(wèn)題。

對(duì)于不同圖數(shù)據(jù)屬性與分布不同的問(wèn)題，OFA通過(guò)提出文本圖（Text-Attributed Graph， TAGs）的概念來(lái)統(tǒng)一所有圖數(shù)據(jù)。利用文本圖，OFA將所有的圖數(shù)據(jù)中的節(jié)點(diǎn)信息與邊信息用統(tǒng)一的自然語(yǔ)言框架來(lái)描述，具體如下圖所示:

接著，OFA通過(guò)單一LLM模型對(duì)所有數(shù)據(jù)中的文本進(jìn)行表示學(xué)習(xí)得到其嵌入向量。

這些嵌入向量將作為圖模型的輸入特征。這樣，來(lái)自不同領(lǐng)域的圖數(shù)據(jù)將被映射到相同的特征空間，使得訓(xùn)練一個(gè)統(tǒng)一的GNN模型可行。

OFA收集了9個(gè)來(lái)自不同領(lǐng)域，不同規(guī)模的圖數(shù)據(jù)集，包括引用關(guān)系圖，Web鏈接圖，知識(shí)圖譜，分子圖，如下圖所示:

此外，OFA提出Nodes-of-Interest（NOI）子圖與NOI提示節(jié)點(diǎn)(NOI Prompt Node)來(lái)統(tǒng)一圖領(lǐng)域內(nèi)不同的子任務(wù)類(lèi)型。這里NOI代表參與到相應(yīng)任務(wù)的一組目標(biāo)節(jié)點(diǎn)。

比如，在節(jié)點(diǎn)預(yù)測(cè)任務(wù)中，NOI是指需要預(yù)測(cè)的單個(gè)節(jié)點(diǎn);而在鏈路任務(wù)中，NOI包括需要預(yù)測(cè)鏈路的兩個(gè)節(jié)點(diǎn)。NOI子圖是指圍繞著這些NOI節(jié)點(diǎn)擴(kuò)展出的一個(gè)包含h-hop鄰域的子圖。

然后，NOI提示節(jié)點(diǎn)為一個(gè)新引入的節(jié)點(diǎn)類(lèi)型，直接連接到所有的NOI上。

重要的是，每個(gè)NOI提示節(jié)點(diǎn)包含了當(dāng)前任務(wù)的描述信息，這些信息以自然語(yǔ)言的形式存在，并和文本圖被同一個(gè)LLM所表示。

由于NOI中節(jié)點(diǎn)所包含的信息在經(jīng)過(guò)GNNs的消息傳遞后將被NOI提示節(jié)點(diǎn)所收集，GNN模型僅需通過(guò)NOI提示節(jié)點(diǎn)來(lái)進(jìn)行預(yù)測(cè)。

這樣，所有不同的任務(wù)類(lèi)型將擁有統(tǒng)一的任務(wù)表示。具體實(shí)例如下圖所示:

最后，為了實(shí)現(xiàn)圖領(lǐng)域的in-context learning，OFA引入統(tǒng)一的提示子圖。

在一個(gè)有監(jiān)督的k-way分類(lèi)任務(wù)場(chǎng)景下，這個(gè)提示子圖包含了兩類(lèi)節(jié)點(diǎn):一類(lèi)是上文提到的NOI提示節(jié)點(diǎn)，另一類(lèi)是代表k個(gè)不同類(lèi)別的類(lèi)別節(jié)點(diǎn)（Class Node）。

每個(gè)類(lèi)別節(jié)點(diǎn)的文本將描述此類(lèi)別的相關(guān)信息。

NOI提示節(jié)點(diǎn)將會(huì)單向連接到所有類(lèi)別節(jié)點(diǎn)上。通過(guò)這個(gè)方式構(gòu)建好的圖將被輸入進(jìn)圖神經(jīng)網(wǎng)路模型進(jìn)行消息傳遞與學(xué)習(xí)。

最終，OFA將對(duì)每個(gè)類(lèi)別節(jié)點(diǎn)分別進(jìn)行二分類(lèi)任務(wù)，并取概率最高的類(lèi)別節(jié)點(diǎn)作為最終的預(yù)測(cè)結(jié)果。

由于類(lèi)別信息存在于提示子圖中，即使遇到全新的分類(lèi)問(wèn)題，OFA通過(guò)構(gòu)建相應(yīng)的提示子圖即可直接進(jìn)行預(yù)測(cè)而無(wú)需任何微調(diào)，從而實(shí)現(xiàn)了零樣本學(xué)習(xí)。

對(duì)于少樣本學(xué)習(xí)場(chǎng)景，一個(gè)分類(lèi)任務(wù)將包含一個(gè)query輸入圖和多個(gè)support輸入圖，OFA的提示圖范式會(huì)將每個(gè)support輸入圖的NOI提示節(jié)點(diǎn)與其所對(duì)應(yīng)的類(lèi)別節(jié)點(diǎn)相連，同時(shí)將query輸入圖的NOI提示節(jié)點(diǎn)與所有類(lèi)別節(jié)點(diǎn)相連。

后續(xù)的預(yù)測(cè)步驟與上文所述一致。這樣每個(gè)類(lèi)別節(jié)點(diǎn)將會(huì)額外得到support輸入圖的信息，從而在統(tǒng)一的范式下實(shí)現(xiàn)少樣本學(xué)習(xí)。

OFA的主要貢獻(xiàn)總結(jié)如下:

統(tǒng)一的圖數(shù)據(jù)分布:通過(guò)提出文本圖并用LLM轉(zhuǎn)化文本信息，OFA實(shí)現(xiàn)了圖數(shù)據(jù)的分布對(duì)齊與統(tǒng)一。

統(tǒng)一的圖任務(wù)形式:通過(guò)NOI子圖與NOI提示節(jié)點(diǎn)，OFA實(shí)現(xiàn)了多種圖領(lǐng)域子任務(wù)的統(tǒng)一表示。

統(tǒng)一的圖提示范式:通過(guò)提出新穎的圖提示范式，OFA實(shí)現(xiàn)了圖領(lǐng)域內(nèi)的多場(chǎng)景in-context learning。

超強(qiáng)泛化能力

文章在所收集的9個(gè)數(shù)據(jù)集上對(duì)OFA框架進(jìn)行了測(cè)試，這些測(cè)試覆蓋了在有監(jiān)督學(xué)習(xí)場(chǎng)景下的十種不同任務(wù)，包括節(jié)點(diǎn)預(yù)測(cè)、鏈路預(yù)測(cè)和圖分類(lèi)。

實(shí)驗(yàn)的目的是驗(yàn)證單一的OFA模型處理多任務(wù)的能力，其中作者對(duì)比使用不同LLM（OFA-{LLM}）和每個(gè)任務(wù)訓(xùn)練單獨(dú)模型(OFA-ind-{LLM})的效果。

比較結(jié)果如下表所示:

可以看到，基于OFA強(qiáng)大的泛化能力，一個(gè)單獨(dú)的圖模型（OFA-st，OFA-e5，OFA-llama2-7b，OFA-llama2-13b）即能夠在所有的任務(wù)上都具有與傳統(tǒng)的單獨(dú)訓(xùn)練模型(GCN， GAT， OFA-ind-st)相近或更好的表現(xiàn)。

同時(shí)，使用更強(qiáng)大的LLM可以帶來(lái)一定的性能提升。文章進(jìn)一步繪制了訓(xùn)練完成的OFA模型對(duì)于不同任務(wù)的NOI提示節(jié)點(diǎn)的表示。

可以看到不同的任務(wù)被模型嵌入到不同的子空間，從而使得OFA可以對(duì)于不同的任務(wù)進(jìn)行分別的學(xué)習(xí)而不會(huì)相互影響。

在少樣本以及零樣本的場(chǎng)景下，OFA在ogbn-arxiv（引用關(guān)系圖），F(xiàn)B15K237(知識(shí)圖譜)以及Chemble(分子圖)上使用單一模型進(jìn)行預(yù)訓(xùn)練，并測(cè)試其在不同下游任務(wù)及數(shù)據(jù)集上的表現(xiàn)。結(jié)果如下: