[论文解读] Self-supervised edge features for improved Graph Neural Network training

Arijit Sehanobish, Neal G. Ravindra+1 more|arXiv (Cornell University)|Jun 23, 2020

Single-cell and spatial transcriptomics参考文献 66被引用 7

一句话总结

该论文提出了一种自监督框架，通过图注意力网络（GAT）的注意力系数和Forman-Ricci曲率来生成与领域无关的图神经网络（GNN）边特征，用于捕捉局部图几何结构。通过使用集合变换器（Set Transformer）整合这些特征，该方法在神经疾病、SARS-CoV-2感染和COVID-19的单细胞RNA-seq数据集中显著提升了节点分类性能，同时实现了对与疾病严重程度相关的关键基因和细胞类型的可解释性识别。

ABSTRACT

Graph Neural Networks (GNN) have been extensively used to extract meaningful representations from graph structured data and to perform predictive tasks such as node classification and link prediction. In recent years, there has been a lot of work incorporating edge features along with node features for prediction tasks. One of the main difficulties in using edge features is that they are often handcrafted, hard to get, specific to a particular domain, and may contain redundant information. In this work, we present a framework for creating new edge features, applicable to any domain, via a combination of self-supervised and unsupervised learning. In addition to this, we use Forman-Ricci curvature as an additional edge feature to encapsulate the local geometry of the graph. We then encode our edge features via a Set Transformer and combine them with node features extracted from popular GNN architectures for node classification in an end-to-end training scheme. We validate our work on three biological datasets comprising of single-cell RNA sequencing data of neurological disease, extit{in vitro} SARS-CoV-2 infection, and human COVID-19 patients. We demonstrate that our method achieves better performance on node classification tasks over baseline Graph Attention Network (GAT) and Graph Convolutional Network (GCN) models. Furthermore, given the attention mechanism on edge and node features, we are able to interpret the cell types and genes that determine the course and severity of COVID-19, contributing to a growing list of potential disease biomarkers and therapeutic targets.

研究动机与目标

• 解决图神经网络（GNN）在图结构化生物数据中缺乏可泛化、手工设计的边特征的问题。
• 在无需标注边特征的情况下，提升GNN在节点分类任务中的性能。
• 通过识别单细胞转录组中与生物相关的基因和细胞类型，实现对GNN决策的可解释性。
• 利用集合变换器将自监督边特征与节点表征进行整合，实现端到端训练。
• 将该框架应用于真实世界生物数据集，包括SARS-CoV-2和COVID-19患者数据，以发现潜在疾病生物标志物。

提出的方法

• 对节点特征应用Louvain聚类以生成社区标签，并利用这些标签自监督地训练GAT模型。
• 提取训练后GAT的边注意力系数，作为自监督边特征，以捕捉节点间的关系重要性。
• 为每条边计算Forman-Ricci曲率，以编码局部图几何结构和拓扑特征。
• 使用集合变换器在端到端训练方案中聚合并融合自监督边特征与GNN的节点特征。
• 将该框架与GCN和GAT架构集成，用于在三个生物数据集上的节点分类任务。
• 利用GAT和集合变换器中的注意力机制解释模型预测结果，并识别具有影响力的基因和细胞类型。

实验结果

研究问题

• RQ1 自监督学习边特征是否能在无需手工设计或标注边特征的情况下，提升GNN在节点分类任务中的性能？
• RQ2 将Forman-Ricci曲率作为边特征引入，如何增强生物图中局部图结构的表征能力？
• RQ3 GAT和集合变换器中的注意力机制在多大程度上能提供关于SARS-CoV-2感染和COVID-19严重程度生物决定因素的可解释性见解？
• RQ4 所提出的框架是否能在包括神经疾病、体外SARS-CoV-2感染和人类患者样本在内的多样化生物图数据集中实现泛化？
• RQ5 该模型能否以支持治疗靶点假设生成的方式，识别出与疾病进展和严重程度相关的生物上具有意义的基因和细胞类型？

主要发现

• 所提出的方法在三个单细胞RNA-seq数据集（包括神经疾病、SARS-CoV-2感染和人类COVID-19患者数据）上，相较于基线GCN和GAT模型，实现了更优的节点分类性能。
• 从GAT注意力系数中提取的自监督边特征显著提升了模型训练效率和泛化能力，即使在无标注边注释的情况下亦然。
• Forman-Ricci曲率有效捕捉了局部图几何结构，并在识别结构上显著的边方面对表征学习有积极贡献。
• 通过集合变换器集成边特征，实现了关系与结构信息与节点表征的有效融合，从而提升了预测性能。
• GAT和集合变换器层中的注意力机制成功突出了特定的基因和细胞类型（如干扰素信号通路相关基因），这些基因和细胞类型对SARS-CoV-2感染和疾病严重程度具有预测价值。
• 模型识别出高曲率和强注意力的细胞簇，提示其在免疫病理中的潜在作用，为后续实验验证提供了潜在生物标志物和治疗靶点。