博客详情

主页 / 博客详情
blog

理解图注意力网络

从 图卷积网络 (GCN) 中，我们了解到结合局部图结构和节点级特征可以在节点分类任务上取得很好的性能。然而，GCN的聚合方式依赖于结构，这可能会损害其泛化能力。

一种替代方法是简单地对所有邻居节点的特征进行平均，就像 GraphSAGE 中那样。 图注意力网络 提出了一种替代方法，通过依赖于特征且与结构无关的归一化方式，对邻居特征进行加权，这借鉴了注意力的思想。

本教程的目标

  • 解释什么是图注意力网络。
  • 演示如何在DGL中实现它。
  • 理解学到的注意力。
  • 介绍归纳学习。

将注意力机制引入GCN

GAT和GCN之间的主要区别在于如何聚合来自一跳邻居的信息。

对于GCN，图卷积操作生成邻居节点特征的归一化总和

其中是其一跳邻居的集合（若要包含在集合中，只需为每个节点添加自环），是基于图结构的归一化常数，是激活函数（GCN使用），并且是用于节点特征转换的共享权重矩阵。在 GraphSAGE 中提出的另一个模型采用了相同的更新规则，只是它们将.

GAT引入了注意力机制，作为静态归一化卷积操作的替代。下面是计算节点嵌入的公式的层来自层嵌入:

解释

  • 公式 (1) 是对下层嵌入的线性变换并且是其可学习的权重矩阵。
  • 公式 (2) 计算两个邻居之间的成对的*未归一化*注意力分数。在这里，它首先连接两个节点的嵌入，其中表示拼接，然后将其与可学习的权重向量进行点积，最后应用LeakyReLU。这种形式的注意力通常称为*加性注意力*，与Transformer模型中的点积注意力形成对比。
  • 公式 (3) 应用softmax对每个节点的传入边的注意力分数进行归一化。
  • 公式 (4) 类似于GCN。来自邻居的嵌入被聚合在一起，并根据注意力分数进行缩放。

论文中还有其他细节，例如dropout和跳跃连接。为了简化起见，本教程中省略了这些内容，并在末尾为感兴趣的读者提供了完整示例的链接。

本质上，GAT只是一个不同的聚合函数，它使用对邻居特征的注意力进行加权，而不是简单的平均聚合。

DGL中的GAT

我们先来大致了解一下 GATLayer 模块在DGL中是如何实现的。不用担心，我们将逐一分解上面的四个公式。

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
        # equation (1)
        self.fc = nn.Linear(in_dim, out_dim, bias=False)
        # equation (2)
        self.attn_fc = nn.Linear(2 * out_dim, 1, bias=False)
    
    def edge_attention(self, edges):
        # edge UDF for equation (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):
        # message UDF for equation (3) & (4)
        return {'z' : edges.src['z'], 'e' : edges.data['e']}
    
    def reduce_func(self, nodes):
        # reduce UDF for equation (3) & (4)
        # equation (3)
        alpha = F.softmax(nodes.mailbox['e'], dim=1)
        # equation (4)
        h = torch.sum(alpha * nodes.mailbox['z'], dim=1)
        return {'h' : h}
    
    def forward(self, h):
        # equation (1)
        z = self.fc(h)
        self.g.ndata['z'] = z
        # equation (2)
        self.g.apply_edges(self.edge_attention)
        # equation (3) & (4)
        self.g.update_all(self.message_func, self.reduce_func)
        return self.g.ndata.pop('h')

公式 (1)

第一个很简单。线性变换非常常见，可以使用Pytorch的 torch.nn.Linear 轻松实现。

公式 (2)

未归一化注意力分数是使用相邻节点的嵌入计算的并且。这表明注意力分数可以看作是边数据，可以使用 apply_edges API计算。 apply_edges 的参数是一个 边用户自定义函数 (Edge UDF)，定义如下

    def edge_attention(self, edges):
        # edge UDF for equation (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 会将所有边数据 批量 处理到一个张量中，因此这里的 cat 和 attn_fc 是并行应用于所有边的。

公式 (3) 和 (4)

类似于GCN，使用 update_all API触发所有节点上的消息传递。消息函数发送两个张量：源节点的变换后的 z 嵌入以及每条边上的未归一化注意力分数 e。Reduce函数随后执行两个任务

  1. 使用softmax对注意力分数进行归一化（公式 (3)）。
  2. 聚合按注意力分数加权的邻居嵌入（公式 (4)）。

两个任务都首先从mailbox中获取数据，然后在第二个维度 (dim=1) 上对其进行操作，消息在该维度上进行批量处理。

    def reduce_func(self, nodes):
        # reduce UDF for equation (3) & (4)
        # equation (3)
        alpha = F.softmax(nodes.mailbox['e'], dim=1)
        # equation (4)
        h = torch.sum(alpha * nodes.mailbox['z'], dim=1)
        return {'h' : h}

多头注意力

类似于ConvNet中的多通道，GAT引入了 多头注意力 来丰富模型容量并稳定学习过程。每个注意力头都有自己的参数，它们的输出可以通过两种方式合并

或者

其中是头部的数量。作者建议在中间层使用拼接，在最终层使用平均。

我们可以使用上面定义的单头 GATLayer 作为下方 MultiHeadGATLayer 的构建块

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':
            # concat on the output feature dimension (dim=1)
            return torch.cat(head_outs, dim=1)
        else:
            # merge using average
            return torch.mean(torch.stack(head_outs))

整合所有部分

现在，我们可以定义一个两层的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)
        # Be aware that the input dimension is hidden_dim*num_heads since
        #   multiple head outputs are concatenated together. Also, only
        #   one attention head in the output layer.
        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()

# create the model
net = GAT(g, 
          in_dim=features.size()[1], 
          hidden_dim=8, 
          out_dim=7, 
          num_heads=8)
print(net)

# create optimizer
optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)

# main loop
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数据集上的模型性能。

模型 准确率
GCN (论文) %
GCN (dgl) %
GAT (论文) %
GAT (dgl) %

我们的模型学到了什么样的注意力分布？

由于注意力权重与边相关联，我们可以通过边的颜色来可视化它。下面我们选取Cora数据集的一个子图，并绘制最后一个 GATLayer 的注意力权重。节点根据其标签着色，而边根据注意力权重的大小着色，这可以参考右侧的颜色条。

可以看到，模型似乎学到了不同的注意力权重。为了更彻底地理解其分布，我们测量了注意力分布的 熵。对于任何节点, 与其所有邻居一起形成一个离散概率分布，其熵由下式给出

直观地说，低熵意味着高度集中，反之亦然；熵为0意味着所有注意力都集中在一个源节点上。均匀分布的熵最高，为。理想情况下，我们希望看到模型学习到熵较低的分布（即，一两个邻居比其他邻居重要得多）。

注意，由于节点可以具有不同的度，最大熵也会不同。因此，我们绘制了整个图中所有节点熵值的聚合直方图。下面是每个注意力头学习到的注意力直方图。

作为参考，这是所有节点具有均匀注意力权重分布时的直方图。

可以看出，学习到的注意力值与均匀分布非常相似（即所有邻居同等重要）。这部分解释了为什么GAT在Cora上的性能与GCN接近（根据 作者报告的结果，平均精度差异在次运行的平均精度差异小于％）；注意力机制无关紧要，因为它没有区分出太多东西。

这是否意味着注意力机制没有用？ 不是的！不同的数据集呈现出截然不同的模式，我们接下来会展示。

蛋白质-蛋白质相互作用 (PPI) 网络

此处使用的PPI数据集包含个图，对应于不同的人体组织。节点最多可以有种标签，因此节点的标签表示为一个大小为的二进制张量。任务是预测节点标签。

我们使用个图用于训练，个用于验证，个用于测试。每个图的平均节点数为。每个节点有

个特征，这些特征由位置基因集、motif基因集和免疫学特征组成。关键的是，测试图在训练期间完全未被观察到，这种设置称为“归纳学习”。我们比较了GAT和GCN在此任务上

模型 F1 分数(micro)
GAT
GCN
论文

上表是这项实验的结果，我们使用 micro F1 分数 来评估模型性能。

训练期间，我们使用 BCEWithLogitsLoss 作为损失函数。下面展示了GAT和GCN的学习曲线；显而易见的是GAT相对于GCN的显著性能优势。

ppi-curve

和之前一样，我们可以通过展示节点级注意力熵的直方图来对学习到的注意力进行统计分析。下面是不同注意力层学习到的注意力直方图。

第1层学习到的注意力

第2层学习到的注意力

最终层学习到的注意力

再次与均匀分布进行比较

显然，GAT确实学习到了尖锐的注意力权重！层之间也存在明显的模式：注意力随着层数的增加变得更加尖锐。

与Cora数据集上GAT的提升微不足道不同，在PPI数据集上，GAT与 GAT论文 中比较的其他GNN变体之间存在显著的性能差距（至少％），并且两者之间的注意力分布明显不同。虽然这值得进一步研究，但一个直接的结论是，GAT的优势可能更多地体现在处理邻居结构更复杂的图的能力上。

下一步？

至此，我们演示了如何使用DGL实现GAT。还有一些遗漏的细节，例如dropout、跳跃连接和超参数调整，这些都是常见做法，不涉及DGL相关概念。我们推荐感兴趣的读者参考完整示例。

  • 在此处查看优化后的完整示例 此处。
  • 敬请关注我们的下一个教程，内容是如何通过并行化多个注意力头和SPMV优化来加速GAT模型。

2月 17日

博客 新闻 发布

作者： Hao Zhang, Mufei Li, Minjie Wang, Zheng Zhang，分类： 博客