《Pooling Regularized Graph Neural Network for fMRI Biomarker Analysis》阅读笔记

前言

2020年MICCAI的论文，Individual-level的。

论文地址：https://arxiv.org/abs/2007.14589
代码地址：https://github.com/xxlya/PRGNN_fMRI

---

一、模型

模型由convolutional layer、pooling layer以及readout layer组成。

1.定义

样本： $G=(V,\epsilon )$。

节点集： V = { v 1 , . . . , v N } V=\{v_1,...,v_N\} V={v1​,...,vN​}，$N$为ROI的个数。

邻接矩阵： $E=[e_{ij}]\in R^{N\times N}$，$e_{ij}=0,if(v_i,v_j)\notin \epsilon$。

属性矩阵： $H=[h_i,...,h_N{]}^T$，$h_i$表示节点$i$的属性。

2.Convolutional Layer

在GAT的基础上，加入了边的属性。

1. 计算节点对其邻居节点的注意力系数。注意与GAT的区别，GAT模型是没有$e_{i,j}$的。
   
   

2. 对注意力系数归一化。
   
   

3. 更新节点$i$的嵌入表示。
   
   

3.Pooling Layer

(1).现存方法及其弊端

在graph上做池化的方法可以分为两种：

第一种，基于聚类的池化方法，根据图的拓扑结构把几个节点聚合成一个超级节点。在脑网络中，这种池化方法的可解释性不强。

第二种，基于重要性的池化方法，计算每个节点的重要程度并保存排名靠前的几个节点。然而，把gPool或SAGPool等现有的池化模型照搬到脑网络中，有以下几点缺陷：1. 被保留节点与被丢弃节点之间的区分度不大，不利于寻找有代表性的区域生物标志；2.同一类被试（患者或正常人）中，各节点的重要程度可能完全不同，不利于找到组水平的生物标志。

(2).本文池化方法及其创新

虽然基于节点重要程度的池化方法有上述这些缺点，但只要对其加以限制，还是能利用这类方法去找到生物标志的。因此，本文还是采用的基于重要性的方法：

SAGE pooling中函数与卷积层的计算方法相似，只不过${\varphi }_i^{\theta }\in R^{1\times d^{(l)}}$，得到的结果是一个标量。

在得到各个节点的重要程度后，只保留前k个（或一定比例）节点。

论文中只描述了池化后的邻接矩阵如何表示，并没有描述如何设计属性矩阵，等我看完代码再把这个部分补充进来。

上述两个池化公式，其实就是gPool和SAGPool中的公式。如果只是这样的话，是没有创新可言的。本文的创新点在于为池化部分设计了两个损失函数：Distance Loss用来加强被保留/丢弃的节点之间的差异；Group-level Consistency Loss则迫使同类graph的节点重要性排序相似。具体的函数表达式放到后面部分讲。

4.Readout Layer

尽管用了池化的方法，但最后并不是只剩下一个节点。因此，仍然需要从剩余的节点中“总结”出全局表示向量。

class NNGAT_Net(torch.nn.Module):
    def __init__(self, ratio, indim, poolmethod = 'topk'):
        super(NNGAT_Net, self).__init__()
        self.dim1 = 32
        self.dim2 = 32
        self.dim3 = 8
        self.indim = indim
        self.poolmethod = poolmethod
		# 这里直接调用torch_geometric中的GAT，应该是没用到边的属性的。
		# 即没有用Convoluton Layer中的第一个公式，而是用了GAT原本的更新公式
        self.conv1 = GATConv( self.indim, self.dim1)
        self.bn1 = torch.nn.BatchNorm1d(self.dim1)
        if self.poolmethod == 'topk':
            self.pool1 = TopKPooling(self.dim1, ratio=ratio, multiplier=1, nonlinearity=torch.sigmoid)
        elif self.poolmethod == 'sag':
            self.pool1 = SAGPooling(self.dim1, ratio=ratio, GNN=GATConv,nonlinearity=torch.sigmoid) #0.4 data1 10 fold

        self.conv2 = GATConv(self.dim1, self.dim2)
        self.bn2 = torch.nn.BatchNorm1d(self.dim2)
        if self.poolmethod == 'topk':
            self.pool2 = TopKPooling(self.dim2, ratio=ratio, multiplier=1, nonlinearity=torch.sigmoid)
        elif self.poolmethod == 'sag':
            self.pool2 = SAGPooling(self.dim2, ratio=ratio, GNN=GATConv,nonlinearity=torch.sigmoid)

		# 原文中是在最后一层池化后用mean来获得全局向量
		# 这里用了mean||max，并且用skip-connection拼接了所有池化层后的向量表示，与原文有出入。
        self.fc1 = torch.nn.Linear((self.dim1+self.dim2)*2, self.dim2)
        self.bn4 = torch.nn.BatchNorm1d(self.dim2)
        self.fc2 = torch.nn.Linear(self.dim2, self.dim3)
        self.bn5 = torch.nn.BatchNorm1d(self.dim3)
        self.fc3 = torch.nn.Linear(self.dim3, 2)

    def forward(self, x, edge_index, batch, edge_attr):
        # edge_attr = edge_attr.squeeze()  
        # edge_index, edge_attr = self.augment_adj(edge_index, edge_attr, x.size(0))
        x = self.conv1(x, edge_index)
        if x.norm(p=2, dim=-1).min() == 0:
            print('x is zeros')
        x, edge_index, edge_attr, batch, perm, score1 = self.pool1(x, edge_index, edge_attr, batch)
        x1 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)

        edge_attr = edge_attr.squeeze()
        edge_index, edge_attr = self.augment_adj(edge_index, edge_attr, x.size(0))

        x = self.conv2(x, edge_index)
        x, edge_index, edge_attr, batch, perm, score2  = self.pool2(x, edge_index, edge_attr, batch)
        x2 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)

        x = torch.cat([x1,x2], dim=1) #concate

        x = self.bn4(F.relu(self.fc1(x)))
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.bn5(F.relu(self.fc2(x)))
        x = F.dropout(x, p=0.5, training=self.training)
        x = F.log_softmax(self.fc3(x), dim=-1)

        return x, score1, score2

    def augment_adj(self, edge_index, edge_weight, num_nodes):
        edge_index, edge_weight = add_self_loops(edge_index, edge_weight,
                                                 num_nodes=num_nodes)
        edge_index, edge_weight = sort_edge_index(edge_index, edge_weight,
                                                  num_nodes)
        edge_index, edge_weight = spspmm(edge_index, edge_weight, edge_index,
                                         edge_weight, num_nodes, num_nodes,
                                         num_nodes)
        edge_index, edge_weight = remove_self_loops(edge_index, edge_weight)

        return edge_index, edge_weight

5.Loss Function

(1).Distance Loss

引入这个损失的目的是为了增强被保留节点与被丢弃节点间的区分度。

${\hat{s}}_m^{(l)}=[{\hat{s}}_{m,1}^{(l)},...,{\hat{s}}_{m,N}^{(l)}]$ 代表第$m$个被试的各节点在第$l$层的重要程度（降序排列）。

记前$k^{(l)}$个节点为$a_{m,i}^{(l)}={\hat{s}}_{m,i}^{(l)},i=1,...,k^{(l)}$；

其余节点则记为$b_{m,j}^{(l)}={\hat{s}}_{m,j+k^{(l)}}^{(l)},j=1,...,N^{(l)}-k^{(l)}$

1). MDD Loss

引用了GAN中的损失函数。

if opt.distL == 'mmd':
	mmd = MDD_loss()
	def dist_loss(s, ratio):
		s = s.sort(dim=1).values
		# 这里和原文有出入。如果ratio > 0.5，那么source和target就会有重复的节点。
		# 按论文的意思应该是没有的重复的节点，而且source和target第二维的维度可以不相同。
		source = s[:,-int(s.size(1)*ratio):]
		target = s[:,:int(s.size(1)*ratio)]
		res = mmd(source, target)
		return -res

class MDD_loss(nn.Module):
   def __init__(self, kernel_mul=2.0, kernel_num=5):
       super(MMD_loss, self).__init__()
       self.kernel_num = kernel_num
       self.kernel_mul = kernel_mul
       self.fix_sigma = None

    def guassian_kernel(self, source, target, kernel_mul=2.0, kernel_num=1, fix_sigma=None):
        n_samples = int(source.size()[0]) + int(target.size()[0])
		# 假设batch_size=1，s=[s_1, s_2, ... ,s_n]
		# source = [s_1, s_2, ..., s_k]; target = [s_n-k, s_n-k+1, ..., s_n]
		# 记source = [a_1, a_2, ..., a_k]; target = [b_1, b_2, ..., b_k]

		# total = [[a_1, a_2, ..., a_k], [b_1, b_2, ..., b_k]]
        total = torch.cat([source, target], dim=0)
        # toal0 = [[[a_1, a_2, ..., a_k], [b_1, b_2, ..., b_k]], [[a_1, a_2, ..., a_k], [b_1, b_2, ..., b_k]]]
        total0 = total.unsqueeze(0).expand(int(total.size(0)), int(total.size(0)), int(total.size(1)))
        # total1 = [[[a_1, a_2, ..., a_k],[a_1, a_2, ..., a_k]],[[b_1, b_2, ..., b_k],[b_1, b_2, ..., b_k]]]
        total1 = total.unsqueeze(1).expand(int(total.size(0)), int(total.size(0)), int(total.size(1)))
        # L2_distance = [[aa, ab], [ba, bb]]
        L2_distance = ((total0 - total1) ** 2).sum(2)
        if fix_sigma:
            bandwidth = fix_sigma
        else:
            bandwidth = torch.sum(L2_distance.data) / (n_samples ** 2 - n_samples)
        bandwidth /= kernel_mul ** (kernel_num // 2)
        bandwidth_list = [bandwidth * (kernel_mul ** i) for i in range(kernel_num)]
        # 论文里公式写错了
        kernel_val = [torch.exp(-L2_distance / bandwidth_temp) for bandwidth_temp in bandwidth_list]
        return sum(kernel_val)

	def forward(self, source, target):
		batch_size = int(source.size()[0])
		kernels = self.guassian_kernel(source, target, kernel_mul=self.kernel.mul, kernel_num=self.kernel_num, fix_sigma=self.fix_sigma)
		
		XX = kernels[:batch_size, :batch_size]   # 被保留节点之间的MDD
		YY = kernels[batch_size:, batch_size:]   # 被丢弃节点之间的MDD
		XY = kernels[:batch_size, batch_size:]   # 被保留/丢弃节点之间的MDD
		YX = kernels[batch_size:, :batch_size]   # 被丢弃/保留节点之间的MDD
		loss = torch.mean(XX + YY - XY - YX)
		# 和原文的公式有出入
		return loss

2).BCE Loss

理想状态下，被保留节点的得分应该尽可能接近于1，而被舍弃节点的得分接近于0。因此也可以用二值交叉熵作为损失函数。

if opt.distL == 'bce':
	def dist_loss(s, ratio):
		# s.shape = batch_size x num_nodes
		if ratio > 0.5:
			ratio = 1 - ratio
		s = s.sort(dim=1).values
		# 这里的计算与原文有出入。这里是取相同数量的排名靠前及靠后的节点计算损失。
		res = -torch.log(s[:, -int(s.size(1)*ratio):]+EPS).mean() - torch.log(1-s[:,:int(s.size(1)*ratio)]+EPS).mean()
		return res

(2).Group-level Consistency Loss

前面提到了，不同被试的节点重要性排序可能截然不同，这是不利于寻找组水平的生物标记的。该损失函数能迫使不同被试的$s^{(1)}$尽可能相似（注意，这是只要求第一层的排序相似。因为经过一次池化后，不同被试所保留的节点不一定相同）。

上述公式代表每一类的损失。其中$L_c=D_c-W_c$，$D_c$为对角阵，对角线元素的值均为$M_c$；$W_c$则为全1矩阵。

def consist_loss(s):
	if len(s) == 0:
		ruturn 0
	else:
		s = torch.sigmoid(s)
		# W: 全1矩阵
		W = torch.ones(s.shape[0], s.shape[0])
		# D: 对角阵，对角元素为一个batch中属于该类的被试个数
		D = torch.eye(s.shape[0]) * torch.sum(W, dim=1)
		L = D - W
		L = L.to(device)
		res = torch.trace(torch.transpose(s,0,1) @ L @ s) / (s.shape[0] * s.shape[0])
		return res

再加上分类损失$L_{ce}$，最终的损失函数为：

二、实验结果

1.数据预处理

这篇论文的数据预处理方式和其他论文的处理方法有点不同。通常，会将ROI内所有体素的时间序列的均值作为该区域的时间序列。但这篇文章仅取了1/3的体素，并且用这种方法做了10倍的数据扩充（即对每个ROI随机取10次1/3的体素）。

同时，本篇文章构建graph的方法也有所不同。节点间的连接强度用的是偏相关系数（一般用的是皮尔逊系数），并且只保留前10%的边。如果graph不连通，则按照偏相关系数的降序，继续添加边直到连通。

节点的属性则是用节点间的皮尔逊相关系数表示。

def read_sigle_data(data_dir, filename):
	temp = h5py.File(osp.join(data_dir, filename), 'r')
	
	# read edge and edge attribute
	pcorr = np.abs(temp['pcorr'].value)
	# 这里应该是2019年MICCAI文章的代码，和这篇文章保留前10%不符。
	th = np.percentile(pcorr.reshape(-1), 95)
	pcorr[pcorr < th] = 0
	num_nodes = pcorr.shape[0]
	
	G = from_numpy_matrix(pcorr)
	A = nx.to_scipy_sparse_matrix(G)
	adj = A.tocoo()
	# len(adj.row): 边的数量
	edge_att = np.zeros((len(adj.row)))
	for i in range(len(adj.row)):     
		edge_att[i] = pcorr[adj.row[i], adj.col[i]]
	edge_index = np.stack([adj.row, adj.col])
	edge_index, edge_att = remove_self_loops(torch.from_numpy(edge_index).long(), torch.from_numpy(edge_att).float())
	edge_index, edge_att = coalesce(edge_index, edge_att, num_nodes, num_nodes)

	# 节点属性，皮尔逊系数矩阵
	att = temp['corr'].value

	return edge_att.data.numpy(), edge_index.data.numpy(), att, temp['indicator'].value, num_nodes

2.消融实验

${\lambda }_2=0$时，模型会在训练集上过拟合；而${\lambda }_2$太大时（> 0.1），则会在欠拟合。

随着迭代的进行，节点的重要程度有了较为明显的分化（除了用SAGE pooling的）。被保留节点的得分逐渐趋于1，而被舍弃节点的得分则趋近于0。

3.对比实验

机器学习方法的输入为皮尔逊相关系数矩阵的上三角阵；BrainNetCNN的输入则是皮尔逊相关系数矩阵；GNN方法的输入则与PR-GNN相同。

4.可解释性

当${\lambda }_2$设置的较小时（图a），模型更倾向于寻找个人的生物标志，被试间重叠区域不多；而${\lambda }_2$较大时（图b-c），则更倾向于寻找组水平的生物标志，并且随着${\lambda }_2$的增大，重叠区域增多。

---