【论文阅读】One For All: Toward Training One Graph Model for All Classification Tasks

0、基本信息

• 会议：2024-ICLR-UNDER_REVIEW
• 评分：6，6，6，10
• 作者：Anonymous authors
• 文章链接：ONE FOR ALL: TOWARDS TRAINING ONE GRAPHMODEL FOR ALL CLASSIFICATION TASKS
• 代码链接：ONE FOR ALL: TOWARDS TRAINING ONE GRAPHMODEL FOR ALL CLASSIFICATION TASKS

1、研究动机

 设计一个能够解决多个任务的模型是人工智能长期发展的一个目标。最近，LLMs因其强大的学习能力，在执行跨领域的下游任务上取巨大的成功，但是在图领域上还没有一个统一的模型解决图结构数据上不同的任务。

面临的挑战——
（1）来自不同领域的图数据集通常在特征表示上不同。

 具体来说，来自不同领域的图数据特征有不同的维度，大小，分布和语义意义，几乎不可能使用一个相同的模型直接学习这些数据的表征。

（2）图领域中不同的下游任务涉及图的不同的部分，以及需要指定任务的知识和方法。

 图上任务可以分为，结点级，边级和图级的任务，图神经网络在不同类任务上成功背后的原因是不同的。结点级任务——节点特征恰当的平滑；边级or图级任务——编码局部结构很重要。一个优秀的结点级任务模型不一定能够在边级or图级任务上表现优秀。
总之，现有的模型不能够有效地学习不同的任务。

（3）如何设计一个统一的方法执行跨领域和上下文学习的图任务模型仍然是不清楚的

 针对上述挑战，作者提出一个通用的解决方案——one for all(OFA)，用于建立和训练具有跨领域上下文学习能力的模型。

2、创新点——One For All ：unique features

（1）文本属性图（TAG）
 OFA使用文本属性图（TAGs）将来自不同领域的图数据整合到一个大的TAG数据集，并利用LLMs的能力从所有领域中共同学习。

 9个常见的来自不同领域的数据集，用文本描述图中所有的结点和边，然后使用单个LLM将来自不同领域的文本嵌入到同一个嵌入空间。

（2）NOI子图 和NOI提示结点
 NOI（nodes-of-interest）子图 和NOI提示结点，不仅统一了不同类型的图任务，而且提高了基础模型学习图中结构信息的能力。

（3）图提示范式 (GPP)

 图提示范式 (graph prompting paradigm，GPP)——插入一个提示图到原始的输入图中的特定任务的方式。提示图中的节点包含关于下游任务的所有相关信息（由文本描述，并由与输入图相同的LLM编码器编码）。然后，修改后的图形成为基础模型的实际输入。

 因此，OFA是自适应的，根据提示图以执行不同的任务。训练后，用户可以用自然文本描述任何图形，并应用OFA管道来预测可能看不见的类。

3、准备

文本属性图（TAGs）
 将TAG定义为一个图，图中的每个节点和每条边都与一个文本句子相关联。
TAG：$\mathcal{G}=(\mathcal{V},\mathcal{E},\mathcal{R})$， V = { v 1 , … , v ∣ V ∣ } \mathcal{V}=\{v_1,\dots ,v_{|\mathcal{V}|}\} V={v1​,…,v∣V∣​}是结点集合， R = { r 1 , … , r R } \mathcal{R}=\{r_1,\dots,r_{\mathcal{R}}\} R={r1​,…,rR​}是关系集合， E = { e 1 , … , e ∣ E ∣ } \mathcal{E}=\{e_1,\dots,e_{|\mathcal{E}|}\} E={e1​,…,e∣E∣​}是边集合。一条边$e_{ij}=(v_i,r,v_j)\in \mathcal{E}$由源结点$v_i\in \mathcal{V}$，关系$r\in \mathcal{R}$和目的结点$v_j\in \mathcal{V}$组成。
$s_{v_i}$定义为与结点$v_i$相关的文本句子，$s_{e_{ij}}$表示与边$e_{ij}\in \mathcal{E}$相关联的文本句子。
${\mathcal{N}}_k(v)$定义为结点$v$的$k$跳内所有邻居结点的集合。

学习情景——分类问题

 定义一个数据集$\mathcal{D}=\{(d_i,y_i){\}}_1^D$，$D$为数据集中数据的数量，$d_i$是一个数据样本，$y_i\in \mathcal{Y}$是$d_i$的标签，$\mathcal{Y}$是所有数据标签的集合。将数据集划分为训练集，验证集和测试集，即${\mathcal{D}}_{train},{\mathcal{D}}_{val},{\mathcal{D}}_{test}$.对应的标签集合为${\mathcal{Y}}_{train},{\mathcal{Y}}_{val},{\mathcal{Y}}_{test}$.

关注三种学习情景
（1）监督学习
 该模型将在$D_{train}$上进行训练，并在$D_{val}$上进行评估，来选择最好的模型。最后，使用$D_{test}$对模型的性能进行评估。验证集和测试集数据中的所有标签都在训练过程中已知，即${\mathcal{Y}}_{train}={\mathcal{Y}}_{test}$（？？？）
（2）少样本学习
 少样本学习的训练和评估过程与监督学习类似，但是，在少样本学习中，${\mathcal{Y}}_{train}\cap {\mathcal{Y}}_{test}=\varnothing$.

 少样本学习通常处理 N-way K-shot任务，其中，使用$N\ast K$数据$\{(d_i,y_i){\}}_1^{N\ast K}$作为支持样本，使得每个不同的类都提供有$K$个标记数据，其中，$N$是不同类的总数。接下来，给定这些支持样本，模型需要将查询集合 Q = { d i } 1 n \mathcal{Q} = \{d_i\}^n_1 Q={di​}1n​中的数据分类到这$N$个类中。

（3）零样本学习
 这可以被看作是少数学习的特殊情况。在零次学习中，对于任何 N-way K-shot任务，将$K=0$，因为没有支持样本。

zero-shot learning 是机器学习中的一种策略，其目标是使模型能够理解并识别它在训练阶段未曾遇到过的类别。
few-shot learning旨在在只给出非常少量已知类别的标注数据时，通过学习如何利用这些数据来泛化到未知类别上，完成分类任务。
one-shot learning在只给出单个样本时，通过学习如何利用这个样本来泛化到未知类别上，完成分类任务。

（4）语境学习

 语境学习是指通过给预先训练好的语言模型提供特定任务的例子或语境，对其进行微调的过程。这种方法利用了预先训练的模型的知识和泛化能力，并对其进行调整，以便更好地执行手头的具体任务。
 对于语言模型，这主要是通过提示机制来实现的。预训练的LLM模型将由提示文本$C$提供的提示作为输入，然后，通过生成以C为条件的句子的其余部分来给出任务的答案。

在语境学习过程中，模型可以通过：

• 使用语义先验知识，在给定的上下文样例下预测标签
• 从给定的样例里学习输入-标签映射

 ICL主要思路是，给出少量的标注样本，设计任务相关的指令形成提示模板，用于指导待测试样本生成相应的结果。

4、具体实现

overview
（1）来自不同领域的图数据聚合到具有相同形式的文本属性图TAG，并使用单个LLM模型将所有的TAGs嵌入到相同的空间中；
（2）通过NOI子图和NOI提示结点统一图域中不同类型的任务，这样图模型能够自动的考虑到任务相关的信息；
（3）图提示范式（GPP）将任务信息有机地注入到图数据中，从而实现上下文学习。

下面来分别介绍：

4.1、用TAGs统一来自不同领域的图数据

 尽管不同领域图的属性不同，但所有属性都可以用自然语言来描述。这种优势在于，通过使用文本来描述节点和边，我们可以应用LLM将不同的图形属性编码到同一空间中。

用于任何结点和边文本特征生成的标准化形式：

给定一个文本属性图$\mathcal{G}$，文本特征$s_{v_i}$总是从文本特征结点（Feature node）开始，表示该结点是具有原始图特征的输入结点，而不是提示结点。接下来是特征类型的文本描述（feature description）和特征的内容（feature content），如果一个结点具有多个特征，则将他们通过分号连接起来。

边$e_{ij}$的文本特征$s_{e_{ij}}$的构造是类似的，不过，文本的开始是特征边：

附录B

 将LLM作为编码器，将所有的文本特征编码为一个固定长度的向量作为结点or边的最终输入特征。
也就是说，对于结点$v_i$和边$e_{ij}$，他们的向量表征为：
$$x_i=\mathbf{LLM}(s_{v_i})$$
和
$$x_{ij}=\mathbf{LLM}(s_{e_{ij}})$$
 因为LLM编码的输入特征包含域信息，所以后续的工作可以捕获和利用该信息。通常，任何类型的LLM都可以用作编码器，并且更强的LLM可能会产生更好的整体性能。在OFA中，为了简单起见，本文采用了句子转换器Transformer（Reimers & Gurevych，2019），它具有很强的学习句子嵌入的能力。

4.2、用NOI（NODES-OF-INTEREST）统一不同图任务

Nodes-of-Interest(NOI)是指任务中的一组目标结点，表示为$\mathcal{T}$。NOI不局限于列出任务级别，其大小取决于预测目标。

4.2.1、NOI子图

NOI子图为由NOI结点组成的子图。
定义： S h ( v ) = { V v h , E v h , R v h } S_h(v)=\{\mathcal{V}_v^h,\mathcal{E}_v^h,\mathcal{R}_v^h\} Sh​(v)={Vvh​,Evh​,Rvh​}为结点$v$周围$h$跳ego-subgraphs，包含结点$v$的$h$跳邻居结点和所有互连的边。NOI子图${\mathcal{G}}_h(\mathcal{T})$包含NOI中所有结点的ego-subgraphs：
G h ( T ) = ⋃ v ∈ T S h ( v ) = { ⋃ v ∈ T V v h , ⋃ v ∈ T E v h , ⋃ v ∈ T R v h } \mathcal{G}_h(\mathcal{T})=\bigcup_{v\in\mathcal{T}}\mathcal{S}_h(v)=\{\bigcup_{v\in\mathcal{T}}\mathcal{V}_v^h,\bigcup_{v\in\mathcal{T}}\mathcal{E}_v^h,\bigcup_{v\in\mathcal{T}}\mathcal{R}_v^h\} Gh​(T)=v∈T⋃​Sh​(v)={v∈T⋃​Vvh​,v∈T⋃​Evh​,v∈T⋃​Rvh​}

• 结点级任务——NOI就是结点本身， T = { v } \mathcal{T}=\{v\} T={v}且${\mathcal{G}}_h(\mathcal{T})=S_h(v)$
• 链接级任务—— T = { v i , v j } \mathcal{T}=\{v_i,v_j\} T={vi​,vj​}且${\mathcal{G}}_h(\mathcal{T})=S_h(v_i)\cup S_h{v}_j$
• 图级任务——NOI包含图上所有结点，NOI子图${\mathcal{G}}_h(\mathcal{T})=(\mathcal{V},\mathcal{E},\mathcal{R})$

4.2.2、NOI提示结点

NOI提示结点统一不同任务中处理和readout过程。NOI提示节点与任务提示文本相关联：

文本由相同的LLM进行编码。NOI提示结点与NOI中所有结点相连，并通过信息传递汇总NOI和任务描述中的信息。然后将类结点附加到下游任务的提示结点中。

3.3节会详细介绍class node。每个class node保存与特定类相关的文本信息：

 NOI子图和NOI提示节点的结合实现了对所有节点级、边级和图级任务的统一读取和处理。此外，连接到NOI的NOI提示节点可以被视为标记技巧，其提升了原始图模型的表达能力，以更好地学习NOI周围的结构信息。同时，NOI提示节点上的任务提示文本允许图模型根据特定任务调整读出参数，这在现有工作中是不可行的。

4.3、用于图的上下文学习（ICL）的图提示范式（GPP）

 LLMs通过提示进行情境学习，使模型不需要微调就能够执行不同学习情景的各种下游任务。

 上下文学习的核心原则是使输入数据与下游任务保持一致。因此，图提示范式（GPP）就是操作输入图，使模型能够从输入本身获取任务相关的信息。这种范式赋予图模型对可见和不可见类的上下文学习能力。

 提示图，$p=({\mathcal{V}}_p,{\mathcal{E}}_p,{\mathcal{R}}_p)$，有两种结点，第一种结点为NOI提示节点，假设我们查询目标NOI子图${\mathcal{G}}_h^q({\mathcal{T}}^q)=({\mathcal{V}}_q^h,{\mathcal{E}}_q^h,{\mathcal{R}}_p)$，则NOI提示结点为$p_q$。GPP在NOI提示结点和NOI中的每个结点之间添加一条边，定义为${\mathcal{E}}_{cross}^q=\{(t,r_{t2p},p_q),(p_q,r_{p2t},t)|t\in {\mathcal{T}}^q\}$，$r_{p2t}$和$r_{p2t}$分别是从NOI到NOI提示结点的边和反向边的关系。

提示图中第二种结点类型是class node。每个class node保存与特定类相关的文本信息：

类节点$i$定义为$c_i$，在每个类节点和NOI提示结点添加边，定义为：
${\mathcal{E}}_{query}^q=\{(p_q,r_{q2c},c_i),(c_i,r_{c2q},p_q)|i\in [N]\}$，$N$是类的数量，$r_{q2c}$和$r_{c2q}$分别是从NOI提示结点到类结点的边和反向边的关系。

一个提示图为：
V p = { p q } ⋃ { c i ∣ i ∈ [ N ] } , E p = E q u e r y ⋃ E c r o s s q , R p = { r t 2 p , r p 2 t , r q 2 c , r c 2 q } . \mathcal{V}_{p}=\{p_{q}\}\bigcup\{c_{i}|i\in[N]\},\mathcal{E}_{p}=\mathcal{E}_{query}\bigcup\mathcal{E}_{cross}^{q},\mathcal{R}_{p}=\{r_{t2p},r_{p2t},r_{q2c},r_{c2q}\}. Vp​={pq​}⋃{ci​∣i∈[N]},Ep​=Equery​⋃Ecrossq​,Rp​={rt2p​,rp2t​,rq2c​,rc2q​}.

然后，输入到后续图模型的提示图是输入图和提示图的组合，表示为${\mathcal{G}}_m=({\mathcal{V}}_q^h\bigcup {\hat{\mathcal{V}}}_p,{\mathcal{E}}_q^h\bigcup {\hat{\mathcal{L}}}_p,{\mathcal{R}}_q^h\bigcup {\mathcal{R}}_p).$

 使用图学习模型来处理提示图，并使用类节点的嵌入来进行二进制分类。具体地，令$h_{c_i}$是来自图学习模型的类节点$c_i$的向量表示。我们预测NOI属于类别$i$的可能性通过：
$$P[\text{NOI belongs to class }i]=\sigma (\text{MLP}(h_{c_i}))$$
MLP是一个多层感知机，其一维输出表示$h_{c_i}$的分类分数。

注意，因为NOI提示节点和类节点连接到NOI并包含任务文本描述，所以固定长度的向量$h_{c_i}$包含关于输入图和任务的信息，使得预测依赖于任务。虽然现有的图学习方法需要为不同的任务使用不同的池化机制，但这种公式化将不同级别的任务转化为类节点上相同的二进制分类任务，因此所有任务都可以一起训练。对于多类问题，我们比较不同类的预测得分来做出决策。
$$l={\mathrm{argmax}}_i\left(\mathrm{MLP}(h_{c_i})|i\in [N]\right)$$
$l$是NOI的预测类别。

 GPP也可以解决少样本问题，未看见类的支持NOI子图帮助更好的分类查询子图，支持NOI子图定义为${\mathcal{G}}_h^{i,k}({\mathcal{T}}_k^i)$，对于第$i$类和第$k$个支持样本，NOI $T_k^i$属于第$i$类。

 对于查询NOI提示节点，我们通过以下方式将每个支持NOI子图连接到其对应的支持NOI提示节点$p_{i,k}$：

$${\mathcal{E}}_{cross}^s=\bigcup _{i\in [N],k\in [K]}{\mathcal{E}}_{cross}^{i,k}=\bigcup _{i\in [N],k\in [K]}\{(t,r_{t2p},p_{i,k}),(p_{i,k},r_{p2t},t)|t\in {\mathcal{T}}_k^i\}.$$

 为了增强分类，将支持NOI提示结点连接到其NOI所属的类节点，即，${\mathcal{E}}_{supp}=\{(p_{i,k},r_{s2c},c_i)|i\in [N],k\in [K]\}$。由于查询NOI提示节点和类节点之间的关系类型$r_{s2c}$和$r_{q2c}$不同，该模型能够区分查询和信息，支持NOI。少镜头提示图$\mathcal{P}$的总体组件是

V p = { p q } ⋃ { p i , k ∣ i ∈ [ N ] , k ∈ [ K ] } ⋃ { c i ∣ i ∈ [ N ] } , \mathcal{V}_p=\{p_q\}\bigcup\{p_{i,k}|i\in[N],k\in[K]\}\bigcup\{c_i|i\in[N]\}, Vp​={pq​}⋃{pi,k​∣i∈[N],k∈[K]}⋃{ci​∣i∈[N]},
E p = E c r o s s q ⋃ E q u e r y ⋃ E c r o s s s ⋃ E s u p p , R p = { r t 2 p , r p 2 t , r q 2 c , r c 2 q , r s 2 c } \mathcal{E}_p=\mathcal{E}_{cross}^q\bigcup\mathcal{E}_{query}\bigcup\mathcal{E}_{cross}^s\bigcup\mathcal{E}_{supp},\quad\mathcal{R}_p=\{r_{t2p},r_{p2t},r_{q2c},r_{c2q},r_{s2c}\} Ep​=Ecrossq​⋃Equery​⋃Ecrosss​⋃Esupp​,Rp​={rt2p​,rp2t​,rq2c​,rc2q​,rs2c​}

 提示图可以以与上面讨论的类似的方式构造。与前面提到的场景一样，类节点的输出嵌入用于进行二进制分类。OFA通过将支持NOI提示节点连接到相应的类节点来利用少量支持示例，并且模型综合类节点上的示例信息和任务语义信息以进行更准确的预测。请注意，少数样本情景中类节点表示仍然与零样本情景中的表示一致，因此它们也可以一起训练。

 总而言之，NOI表示与任务相关的节点集合，并且提取的NOI子图包括NOI的邻域信息。然后，NOI提示节点通过图学习模型总结NOI中的信息，因为所有NOI节点都连接到NOI提示节点。NOI提示符节点稍后连接到一组带有文本描述的类节点。在图模型处理之后，类节点表示包含类信息、任务信息和NOI信息，它们可以用于独立地进行预测，就像LLM的提示输入包含输入、目标和任务描述一样。

设计了一个能够识别不同边缘类型的GNN作为OFA中的图模型。

5、训练和评估过程

 不管任务类型，所有输入都被标准化为提示文本属性图$G_m$，OFA首先使用LLM将图中的所有文本嵌入到向量表示中：
$$x_{ij}=LLM(s_{e_{ij}}),\forall e_{ij}\in {\mathcal{E}}_m,x_i=LLM(s_{v_i}),\forall v_i\in {\mathcal{V}}_m$$

 然后，通过一个多个的信息传递层处理嵌入文本，得到${\mathcal{G}}_m$中每个结点的最终嵌入：
$$h_{v_i}^{l+1}={\mathbf{W}}_{\mathrm{self}}{h}_{v_i}^l+\sum _{r\in {\mathcal{R}}_m}\sum _{v_j\in {\mathcal{N}}_1^r(v_i)}\frac{1}{|{\mathcal{N}}_1^r(v_i)|}{\mathbf{W}}_r(\mathrm{ReLU}(h_{v_j}^l+x_{ij}))$$
其中，$h_{v_i}^0=x_i$，$W_{self}$和$W_r$是用于自环和关系的可训练的变换矩阵，${\mathcal{N}}_1^r(v_i)$是连接到具有关系类型$r$的$v_i$的直接邻居。GNN是R-GCN的拓展。

 设计了一个简单的注意力层来总结从$h_1$到$h_L$的所有层输出。最终的结点表征为：
$$h_{v_i}=\mathbf{H}\cdot Softmax(({\mathbf{W}}_k\mathbf{H}{)}^{\prime }\cdot {\mathbf{W}}_q{x}_i{)}^{\prime },\mathbf{H}=[h_{v_i}^1...h_{v_i}^l...h_{v_i}^L]$$
$W_k$和$W_q$是key和query可训练的权重。

最后，我们收集每个类节点 { h c i ∣ ∀ v ∈ [ N ] } \{ h_{c_i}| \forall v \in[N]\} {hci​​∣∀v∈[N]}的输出节点嵌入并且使用MLP预测头对每个类节点执行二进制分类，类$i$的可能性为：

$$p_i=\sigma (\mathsf{MLP}(h_{c_i}))$$
$\sigma$为sigmoid函数。对于多分类问题，所有类别的概率为：

$$l_p=\mathrm{argmax}\left((p_i|i\in [N])\right),$$
$l_p$为模型的最终预测。
对于${\mathcal{G}}_m$的训练损失函数为：
$${\mathcal{L}}_{{\mathcal{G}}_m}=-\frac{1}{N}\sum _{i=1}^N{y}_i\cdot log(p_i)+(1-y_i)\cdot (1-log(p_i))$$
其中，$N$为${\mathcal{G}}_m$中类的数量。

未完待续