[論文レビュー] 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）の注目係数と局所的グラフ幾何を表すフォーマン＝リッチィ曲率を用いて、ドメインに依存しないエッジ特徴を自己教師付きで生成するフレームワークを提案する。これらの特徴をセット変換器を介して統合することで、神経疾患、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の意思決定の解釈可能性を高めること。
• 自己教師付きエッジ特徴をノード表現と組み合わせ、エンドツーエンドの学習が可能なセット変換器を用いて統合すること。
• SARS-CoV-2およびCOVID-19患者データを含む実世界の生物学的データセットにこのフレームワークを適用し、疾患バイオマーカーの同定を行うこと。

提案手法

• ノード特徴に対してLouvainコミュニティ検出を適用し、コミュニティラベルを生成。これを自己教師付きでGATモデルの学習に用いる。
• 学習済みGATから抽出したエッジ注目係数を、ノード間の関係的重要性を表す自己教師付きエッジ特徴として用いる。
• 各エッジに対してフォーマン＝リッチィ曲率を計算し、局所的グラフ幾何およびトポロジー構造を符号化する。
• セット変換器を用いて、自己教師付きエッジ特徴とGNNからのノード特徴を統合・組み合わせ、エンドツーエンドの学習スキームを実現する。
• GCNおよびGATアーキテクチャとフレームワークを統合し、3つの生物学的データセットにおけるノード分類に適用する。
• GATおよびセット変換器内の注目メカニズムを活用し、予測の解釈と、影響力のある遺伝子および細胞型の同定を実現する。

実験結果

リサーチクエスチョン

• RQ1 手作業で作成したか、ラベル付きのエッジ特徴を必要とせずに、エッジ特徴の自己教師付き学習がGNNのノード分類性能を向上させられるか。
• RQ2 エッジ特徴としてフォーマン＝リッチィ曲率を組み込むことで、生物学的グラフにおける局所的グラフ構造の表現がどの程度向上するか。
• RQ3 GATおよびセット変換器内の注目メカニズムは、SARS-CoV-2感染およびCOVID-19重症度の生物学的要因をどの程度解釈可能なインサイトを提供できるか。
• RQ4 本フレームワークは、神経疾患、in vitroでのSARS-CoV-2感染、ヒト患者由来の単細胞RNA-seqデータを含む多様な生物学的グラフデータセットに一般化可能か。
• RQ5 モデルは、疾患の進行および重症度に関連する生物学的に意味のある遺伝子および細胞型を、治療標的の仮説生成を支援する形で同定できるか。

主な発見

• 提案手法は、神経疾患、SARS-CoV-2感染、ヒトCOVID-19患者由来の単細胞RNA-seqデータセットの3つすべてにおいて、ベースラインのGCNおよびGATモデルよりも優れたノード分類性能を達成した。
• GATの注目係数から導出された自己教師付きエッジ特徴は、ラベル付きエッジアノテーションがなくても、モデルの学習および一般化性能を顕著に向上させた。
• フォーマン＝リッチィ曲率は、局所的グラフ幾何を効果的に捉え、特に構造的に重要なエッジを同定するうえで表現学習の向上に寄与した。
• セット変換器によるエッジ特徴の統合により、関係性と構造的情報をノードレベルの表現と効果的に統合でき、予測性能の向上が達成された。
• GATおよびセット変換器の注目メカニズムは、インターフェロンシグナルに関与する遺伝子など、SARS-CoV-2感染および病態の重症度を予測する要因となる特定の遺伝子および細胞型を的確に強調した。
• 高い曲率と強い注目度を示す細胞クラスタが同定されたことから、免疫病態の関与が示唆され、今後の実験的検証に向けた候補バイオマーカーおよび治療標的が得られた。