今天小编分享的科学经验:图领網域首个通用框架来了!入选ICLR‘24 Spotlight,任意数据集、分类问题都可搞定,欢迎阅读。
能不能有一种通用的图模型——
它既能够根据分子结构预测毒性,又能够给出社交网络的朋友推荐?
或者既能预测不同作者的论文引用,还可以发现基因网络中的人类衰老机制?
你还真别说,被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
— 完 —
点这里关注我,记得标星哦~
一键三连「分享」、「点赞」和「在看」
科技前沿进展日日相见 ~
>