张昊 李牧非 王敏捷 张峥作者

深入理解图注意力机制

图卷积网络 Graph Convolutional Network (GCN) 告诉我们将局部的图结构和节点特征结合可以在节点分类任务中获得不错的表现。美中不足的是 GCN 结合邻近节点特征的方式和图的结构依依相关,这局限了训练所得模型在其他图结构上的泛化能力。

Graph Attention Network (GAT) 提出了用注意力机制对邻近节点特征加权求和。邻近节点特征的权重完全取决于节点特征,独立于图结构。

在这个教程里我们将:

  • 解释什么是 Graph Attention Network

  • 演示用 DGL 实现这一模型

  • 深入理解学习所得的注意力权重

  • 初探归纳学习 (inductive learning)

难度:★★★★✩(需要对图神经网络训练和 Pytorch 有基本了解)

在 GCN 里引入注意力机制

GAT 和 GCN 的核心区别在于如何收集并累和距离为 1 的邻居节点的特征表示。

在 GCN 里,一次图卷积操作包含对邻节点特征的标准化求和:

其中 N(i) 是对节点 i 距离为 1 邻节点的集合。我们通常会加一条连接节点 i 和它自身的边使得 i 本身也被包括在 N(i) 里。 是一个基于图结构的标准化常数;σ是一个激活函数(GCN 使用了 ReLU);W^((l)) 是节点特征转换的权重矩阵,被所有节点共享。由于 c_ij 和图的机构相关,使得在一张图上学习到的 GCN 模型比较难直接应用到另一张图上。解决这一问题的方法有很多,比如 GraphSAGE 提出了一种采用相同节点特征更新规则的模型,唯一的区别是他们将 c_ij 设为了|N(i)|。

图注意力模型 GAT 用注意力机制替代了图卷积中固定的标准化操作。以下图和公式定义了如何对第 l 层节点特征做更新得到第 l+1 层节点特征:

图 1:图注意力网络示意图和更新公式。

对于上述公式的一些解释:

  • 公式(1)对 l 层节点嵌入做了线性变换,W^((l)) 是该变换可训练的参数

  • 公式(2)计算了成对节点间的原始注意力分数。它首先拼接了两个节点的 z 嵌入,注意 || 在这里表示拼接;随后对拼接好的嵌入以及一个可学习的权重向量 做点积;最后应用了一个 LeakyReLU 激活函数。这一形式的注意力机制通常被称为加性注意力,区别于 Transformer 里的点积注意力。

  • 公式(3)对于一个节点所有入边得到的原始注意力分数应用了一个 softmax 操作,得到了注意力权重

  • 公式(4)形似 GCN 的节点特征更新规则,对所有邻节点的特征做了基于注意力的加权求和。

出于简洁的考量,在本教程中,我们选择省略了一些论文中的细节,如 dropout, skip connection 等等。感兴趣的读者们欢迎参阅文末链接的模型完整实现。

本质上,GAT 只是将原本的标准化常数替换为使用注意力权重的邻居节点特征聚合函数。

GAT 的 DGL 实现

以下代码给读者提供了在 DGL 里实现一个 GAT 层的总体印象。别担心,我们会将以下代码拆分成三块,并逐块讲解每块代码是如何实现上面的一条公式。

import torch
import torch.nn as nn
import torch.nn.functional as F

class GATLayer(nn.Module):
    def __init__(self, g, in_dim, out_dim):
        super(GATLayer, self).__init__()
        self.g = g
        # 公式 (1)
        self.fc = nn.Linear(in_dim, out_dim, bias=False)
        # 公式 (2)
        self.attn_fc = nn.Linear(2 * out_dim, 1, bias=False)

    def edge_attention(self, edges):
        # 公式 (2) 所需,边上的用户定义函数
        z2 = torch.cat([edges.src['z'], edges.dst['z']], dim=1)
        a = self.attn_fc(z2)
        return {'e' : F.leaky_relu(a)}

    def message_func(self, edges):
        # 公式 (3), (4)所需,传递消息用的用户定义函数
        return {'z' : edges.src['z'], 'e' : edges.data['e']}

    def reduce_func(self, nodes):
        # 公式 (3), (4)所需, 归约用的用户定义函数
        # 公式 (3)
        alpha = F.softmax(nodes.mailbox['e'], dim=1)
        # 公式 (4)
        h = torch.sum(alpha * nodes.mailbox['z'], dim=1)
        return {'h' : h}

    def forward(self, h):
        # 公式 (1)
        z = self.fc(h)
        self.g.ndata['z'] = z
        # 公式 (2)
        self.g.apply_edges(self.edge_attention)
        # 公式 (3) & (4)
        self.g.update_all(self.message_func, self.reduce_func)
        return self.g.ndata.pop('h')

实现公式 (1) 

第一个公式相对比较简单。线性变换非常常见。在 PyTorch 里,我们可以通过 torch.nn.Linear 很方便地实现。

实现公式 (2) 

原始注意力权重 e_ij 是基于一对邻近节点 i 和 j 的表示计算得到。我们可以把注意力权重 e_ij 看成在 i->j 这条边的数据。因此,在 DGL 里,我们可以使用 g.apply_edges 这一 API 来调用边上的操作,用一个边上的用户定义函数来指定具体操作的内容。我们在用户定义函数里实现了公式(2)的操作:

 def edge_attention(self, edges):
        # 公式 (2) 所需,边上的用户定义函数
        z2 = torch.cat([edges.src['z'], edges.dst['z']], dim=1)
        a = self.attn_fc(z2)
        return {'e' : F.leaky_relu(a)}

公式中的点积同样借由 PyTorch 的一个线性变换 attn_fc 实现。注意 apply_edges 会把所有边上的数据打包为一个张量,这使得拼接和点积可以并行完成。

实现公式 (3) 和 (4)

类似 GCN,在 DGL 里我们使用 update_all API 来触发所有节点上的消息传递函数。update_all 接收两个用户自定义函数作为参数。message_function 发送了两种张量作为消息:消息原节点的 z 表示以及每条边上的原始注意力权重。reduce_function 随后进行了两项操作:

  1. 使用 softmax 归一化注意力权重(公式(3))。

  2. 使用注意力权重聚合邻节点特征(公式(4))。

这两项操作都先从节点的 mailbox 获取了数据,随后在数据的第二维(dim = 1 ) 上进行了运算。注意数据的第一维代表了节点的数量,第二维代表了每个节点收到消息的数量。

 def reduce_func(self, nodes):
        # 公式 (3), (4)所需, 归约用的用户定义函数
        # 公式 (3)
        alpha = F.softmax(nodes.mailbox['e'], dim=1)
        # 公式 (4)
        h = torch.sum(alpha * nodes.mailbox['z'], dim=1)
        return {'h' : h}

多头注意力 (Multi-head attention)

神似卷积神经网络里的多通道,GAT 引入了多头注意力来丰富模型的能力和稳定训练的过程。每一个注意力的头都有它自己的参数。如何整合多个注意力机制的输出结果一般有两种方式:

以上式子中 K 是注意力头的数量。作者们建议对中间层使用拼接对最后一层使用求平均。

我们之前有定义单头注意力的 GAT 层,它可作为多头注意力 GAT 层的组建单元:

class MultiHeadGATLayer(nn.Module):
    def __init__(self, g, in_dim, out_dim, num_heads, merge='cat'):
        super(MultiHeadGATLayer, self).__init__()
        self.heads = nn.ModuleList()
        for i in range(num_heads):
            self.heads.append(GATLayer(g, in_dim, out_dim))
        self.merge = merge

    def forward(self, h):
        head_outs = [attn_head(h) for attn_head in self.heads]
        if self.merge == 'cat':
            # 对输出特征维度(第1维)做拼接
            return torch.cat(head_outs, dim=1)
        else:
            # 用求平均整合多头结果
            return torch.mean(torch.stack(head_outs))

在 Cora 数据集上训练一个 GAT 模型

Cora 是经典的文章引用网络数据集。Cora 图上的每个节点是一篇文章,边代表文章和文章间的引用关系。每个节点的初始特征是文章的词袋(Bag of words)表示。其目标是根据引用关系预测文章的类别(比如机器学习还是遗传算法)。在这里,我们定义一个两层的 GAT 模型:

class GAT(nn.Module):
    def __init__(self, g, in_dim, hidden_dim, out_dim, num_heads):
        super(GAT, self).__init__()
        self.layer1 = MultiHeadGATLayer(g, in_dim, hidden_dim, num_heads)
        # 注意输入的维度是 hidden_dim * num_heads 因为多头的结果都被拼接在了
        # 一起。 此外输出层只有一个头。
        self.layer2 = MultiHeadGATLayer(g, hidden_dim * num_heads, out_dim, 1)

    def forward(self, h):
        h = self.layer1(h)
        h = F.elu(h)
        h = self.layer2(h)
        return h

我们使用 DGL 自带的数据模块加载 Cora 数据集。

from dgl import DGLGraph
from dgl.data import citation_graph as citegrh

def load_cora_data():
    data = citegrh.load_cora()
    features = torch.FloatTensor(data.features)
    labels = torch.LongTensor(data.labels)
    mask = torch.ByteTensor(data.train_mask)
    g = DGLGraph(data.graph)
    return g, features, labels, mask

模型训练的流程和 GCN 教程里的一样。

import time
import numpy as np
g, features, labels, mask = load_cora_data()

# 创建模型
net = GAT(g, 
          in_dim=features.size()[1], 
          hidden_dim=8, 
          out_dim=7, 
          num_heads=8)
print(net)

# 创建优化器
optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)

# 主流程
dur = []
for epoch in range(30):
    if epoch >=3:
        t0 = time.time()

    logits = net(features)
    logp = F.log_softmax(logits, 1)
    loss = F.nll_loss(logp[mask], labels[mask])

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if epoch >=3:
        dur.append(time.time() - t0)

    print("Epoch {:05d} | Loss {:.4f} | Time(s) {:.4f}".format(
            epoch, loss.item(), np.mean(dur)))

可视化并理解学到的注意力

Cora 数据集

以下表格总结了 GAT 论文以及 dgl 实现的模型在 Cora 数据集上的表现:

可以看到 DGL 能完全复现原论文中的实验结果。对比图卷积网络 GCN,GAT 在 Cora 上有 2~3 个百分点的提升。

不过,我们的模型究竟学到了怎样的注意力机制呢?

由于注意力权重与图上的边密切相关,我们可以通过给边着色来可视化注意力权重。以下图片中我们选取了 Cora 的一个子图并且在图上画出了 GAT 模型最后一层的注意力权重。我们根据图上节点的标签对节点进行了着色,根据注意力权重的大小对边进行了着色(可参考图右侧的色条)。 

图 2:Cora 数据集上学习到的注意力权重

乍看之下模型似乎学到了不同的注意力权重。为了对注意力机制有一个全局观念,我们衡量了注意力分布的熵。对于节点 i,{α_ij }_(j∈N(i)) 构成了一个在 i 邻节点上的离散概率分布。它的熵被定义为:

直观地说,熵低代表了概率高度集中,反之亦然。熵为 0 则所有的注意力都被放在一个点上。均匀分布具有最高的熵(log N(i))。在理想情况下,我们想要模型习得一个熵较低的分布(即某一、两个节点比其它节点重要的多)。注意由于节点的入度不同,它们注意力权重的分布所能达到的最大熵也会不同。

基于图中所有节点的熵,我们画了所有头注意力的直方图。 

图 3:Cora 数据集上学到的注意力权重直方图。

作为参考,下图是在所有节点的注意力权重都是均匀分布的情况下得到的直方图。 

出人意料的,模型学到的节点注意力权重非常接近均匀分布(换言之,所有的邻节点都获得了同等重视)。这在一定程度上解释了为什么在 Cora 上 GAT 的表现和 GCN 非常接近(在上面表格里我们可以看到两者的差距平均下来不到 2%)。由于没有显著区分节点,注意力并没有那么重要。

这是否说明了注意力机制没什么用?不!在接下来的数据集上我们观察到了完全不同的现象。

蛋白质交互网络 (PPI)

PPI(蛋白质间相互作用)数据集包含了 24 张图,对应了不同的人体组织。节点最多可以有 121 种标签(比如蛋白质的一些性质、所处位置等)。因此节点标签被表示为有 121 个元素的二元张量。数据集的任务是预测节点标签。

我们使用了 20 张图进行训练,2 张图进行验证,2 张图进行测试。平均下来每张图有 2372 个节点。每个节点有 50 个特征,包含定位基因集合、特征基因集合以及免疫特征。至关重要的是,测试用图在训练过程中对模型完全不可见。这一设定被称为归纳学习

我们比较了 dgl 实现的 GAT 和 GCN 在 10 次随机训练中的表现。模型的超参数验证集上进行了优化。在实验中我们使用了 micro f1 score 来衡量模型的表现。

在训练过程中,我们使用了 BCEWithLogitsLoss 作为损失函数。下图绘制了 GAT 和 GCN 的学习曲线;显然 GAT 的表现远优于 GCN。 

图 4:PPI 数据集上 GCN 和 GAT 学习曲线比较。

像之前一样,我们可以通过绘制节点注意力分布之熵的直方图来有一个统计意义上的直观了解。以下我们基于一个 3 层 GAT 模型中不同模型层不同注意力头绘制了直方图。

第一层学到的注意力

 第二层学到的注意力

最后一层学到的注意力 

作为参考,下图是在所有节点的注意力权重都是均匀分布的情况下得到的直方图。

可以很明显地看到,GAT 在 PPI 上确实学到了一个尖锐的注意力权重分布。与此同时,GAT 层与层之间的注意力也呈现出一个清晰的模式:在中间层随着层数的增加注意力权重变得愈发集中;最后的输出层由于我们对不同头结果做了平均,注意力分布再次趋近均匀分布。

不同于在 Cora 数据集上非常有限的收益,GAT 在 PPI 数据集上较 GCN 和其它图模型的变种取得了明显的优势(根据原论文的结果在测试集上的表现提升了至少 20%)。我们的实验揭示了 GAT 学到的注意力显著区别于均匀分布。虽然这值得进一步的深入研究,一个由此而生的假设是 GAT 的优势在于处理更复杂领域结构的能力。

拓展阅读

到目前为止我们演示了如何用 DGL 实现 GAT。简介起见,我们忽略了 dropout, skip connection 等一些细节。这些细节很常见且独立于 DGL 相关的概念。有兴趣的读者欢迎参阅完整的代码实现。

关于 DGL 专栏: DGL 是一款全新的面向图神经网络的开源框架。通过该专栏,我们 DGL 团队希望和大家一起学习图神经网络的最新进展。同时展示 DGL 的灵活性和高效性。通过系统学习算法,通过算法理解系统。

DGL专栏
DGL专栏

DGL是一款全新的面向图神经网络的开源框架。通过该专栏,我们DGL团队希望和大家一起学习图神经网络的最新进展。同时展示DGL的灵活性和高效性。通过系统学习算法,通过算法理解系统。

入门注意力机制图神经网络DGL
111
相关数据
激活函数技术

在 计算网络中, 一个节点的激活函数定义了该节点在给定的输入或输入的集合下的输出。标准的计算机芯片电路可以看作是根据输入得到"开"(1)或"关"(0)输出的数字网络激活函数。这与神经网络中的线性感知机的行为类似。 一种函数(例如 ReLU 或 S 型函数),用于对上一层的所有输入求加权和,然后生成一个输出值(通常为非线性值),并将其传递给下一层。

权重技术

线性模型中特征的系数,或深度网络中的边。训练线性模型的目标是确定每个特征的理想权重。如果权重为 0,则相应的特征对模型来说没有任何贡献。

机器学习技术

机器学习是人工智能的一个分支,是一门多领域交叉学科,涉及概率论、统计学、逼近论、凸分析、计算复杂性理论等多门学科。机器学习理论主要是设计和分析一些让计算机可以自动“学习”的算法。因为学习算法中涉及了大量的统计学理论,机器学习与推断统计学联系尤为密切,也被称为统计学习理论。算法设计方面,机器学习理论关注可以实现的,行之有效的学习算法。

参数技术

在数学和统计学裡,参数(英语:parameter)是使用通用变量来建立函数和变量之间关系(当这种关系很难用方程来阐述时)的一个数量。

学习曲线技术

在机器学习领域,学习曲线通常是表现学习准确率随着训练次数/时长/数据量的增长而变化的曲线

概率分布技术

概率分布(probability distribution)或简称分布,是概率论的一个概念。广义地,它指称随机变量的概率性质--当我们说概率空间中的两个随机变量具有同样的分布(或同分布)时,我们是无法用概率来区别它们的。

归纳式学习技术

归纳法或归纳推理(Inductive reasoning),有时叫做归纳逻辑,是论证的前提支持结论但不确保结论的推理过程。 它基于对特殊的代表(token)的有限观察,把性质或关系归结到类型;或基于对反复再现的现象的模式(pattern)的有限观察,公式表达规律。归纳学习则是基于归纳法的机器学习方法

损失函数技术

在数学优化,统计学,计量经济学,决策理论,机器学习和计算神经科学等领域,损失函数或成本函数是将一或多个变量的一个事件或值映射为可以直观地表示某种与之相关“成本”的实数的函数。

超参数技术

在机器学习中,超参数是在学习过程开始之前设置其值的参数。 相反,其他参数的值是通过训练得出的。 不同的模型训练算法需要不同的超参数,一些简单的算法(如普通最小二乘回归)不需要。 给定这些超参数,训练算法从数据中学习参数。相同种类的机器学习模型可能需要不同的超参数来适应不同的数据模式,并且必须对其进行调整以便模型能够最优地解决机器学习问题。 在实际应用中一般需要对超参数进行优化,以找到一个超参数元组(tuple),由这些超参数元组形成一个最优化模型,该模型可以将在给定的独立数据上预定义的损失函数最小化。

注意力机制技术

我们可以粗略地把神经注意机制类比成一个可以专注于输入内容的某一子集(或特征)的神经网络. 注意力机制最早是由 DeepMind 为图像分类提出的,这让「神经网络在执行预测任务时可以更多关注输入中的相关部分,更少关注不相关的部分」。当解码器生成一个用于构成目标句子的词时,源句子中仅有少部分是相关的;因此,可以应用一个基于内容的注意力机制来根据源句子动态地生成一个(加权的)语境向量(context vector), 然后网络会根据这个语境向量而不是某个固定长度的向量来预测词。

张量技术

张量是一个可用来表示在一些矢量、标量和其他张量之间的线性关系的多线性函数,这些线性关系的基本例子有内积、外积、线性映射以及笛卡儿积。其坐标在 维空间内,有 个分量的一种量,其中每个分量都是坐标的函数,而在坐标变换时,这些分量也依照某些规则作线性变换。称为该张量的秩或阶(与矩阵的秩和阶均无关系)。 在数学里,张量是一种几何实体,或者说广义上的“数量”。张量概念包括标量、矢量和线性算子。张量可以用坐标系统来表达,记作标量的数组,但它是定义为“不依赖于参照系的选择的”。张量在物理和工程学中很重要。例如在扩散张量成像中,表达器官对于水的在各个方向的微分透性的张量可以用来产生大脑的扫描图。工程上最重要的例子可能就是应力张量和应变张量了,它们都是二阶张量,对于一般线性材料他们之间的关系由一个四阶弹性张量来决定。

验证集技术

验证数据集是用于调整分类器超参数(即模型结构)的一组数据集,它有时也被称为开发集(dev set)。

卷积神经网络技术

卷积神经网路(Convolutional Neural Network, CNN)是一种前馈神经网络,它的人工神经元可以响应一部分覆盖范围内的周围单元,对于大型图像处理有出色表现。卷积神经网路由一个或多个卷积层和顶端的全连通层(对应经典的神经网路)组成,同时也包括关联权重和池化层(pooling layer)。这一结构使得卷积神经网路能够利用输入数据的二维结构。与其他深度学习结构相比,卷积神经网路在图像和语音识别方面能够给出更好的结果。这一模型也可以使用反向传播算法进行训练。相比较其他深度、前馈神经网路,卷积神经网路需要考量的参数更少,使之成为一种颇具吸引力的深度学习结构。 卷积网络是一种专门用于处理具有已知的、网格状拓扑的数据的神经网络。例如时间序列数据,它可以被认为是以一定时间间隔采样的一维网格,又如图像数据,其可以被认为是二维像素网格。

图神经网络技术

图网络即可以在社交网络或其它基于图形数据上运行的一般深度学习架构,它是一种基于图结构的广义神经网络。图网络一般是将底层图形作为计算图,并通过在整张图上传递、转换和聚合节点特征信息,从而学习神经网络基元以生成单节点嵌入向量。生成的节点嵌入向量可作为任何可微预测层的输入,并用于节点分类或预测节点之间的连接,完整的模型可以通过端到端的方式训练。

LeakyReLU技术

ReLU是将所有的负值都设为零,保留正值;相反,Leaky ReLU是给所有负值赋予一个非零斜率,即x<0时,y=α·x。

遗传算法技术

遗传算法是计算数学中用于解决最优化的搜索算法,是进化算法的一种。进化算法最初是借鉴了进化生物学中的一些现象而发展起来的,这些现象包括遗传、突变、自然选择以及杂交等。 遗传算法通常实现方式为一种计算机模拟。对于一个最优化问题,一定数量的候选解可抽象表示为染色体,使种群向更好的解进化。

优化器技术

优化器基类提供了计算梯度loss的方法,并可以将梯度应用于变量。优化器里包含了实现了经典的优化算法,如梯度下降和Adagrad。 优化器是提供了一个可以使用各种优化算法的接口,可以让用户直接调用一些经典的优化算法,如梯度下降法等等。优化器(optimizers)类的基类。这个类定义了在训练模型的时候添加一个操作的API。用户基本上不会直接使用这个类,但是你会用到他的子类比如GradientDescentOptimizer, AdagradOptimizer, MomentumOptimizer(tensorflow下的优化器包)等等这些算法。

推荐文章
期待下一篇,赶紧更新啊