[논문 리뷰] Self-supervised edge features for improved Graph Neural Network training

Arijit Sehanobish, Neal G. Ravindra외 1명|arXiv (Cornell University)|2020. 06. 23.

Single-cell and spatial transcriptomics참고 문헌 66인용 수 7

한 줄 요약

이 논문은 그래프 신경망(GNNs)를 위한 도메인에 관계없이 적용 가능한 엣지 특징을 생성하기 위해 그래프 어텐션 네트워크(GAT)의 어텐션 계수와 로컬 그래프 기하학을 위한 Forman-Ricci 곡률을 사용하는 자기지도 학습 프레임워크를 제안한다. 세트 트랜스포머를 통해 이러한 특징을 통합함으로써, 신경질환, SARS-CoV-2 감염, 코로나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 및 코로나19 환자 데이터에 프레임워크를 적용하여 잠재적 질병 바이오마커를 규명하기 위해.

제안 방법

• 노드 특징에 Louvain 클러스터링을 적용하여 커뮤니티 레이블을 생성하고, 이를 자기지도 학습 방식으로 GAT 모델을 훈련하는 데 사용한다.
• 훈련된 GAT에서 추출한 엣지 어텐션 계수를 자기지도 엣지 특징으로 사용하여 노드 간의 관계적 중요도를 캡처한다.
• 각 엣지에 대해 Forman-Ricci 곡률을 계산하여 국소 그래프 기하학과 위상적 구조를 인코딩한다.
• 세트 트랜스포머를 사용해 자기지도 엣지 특징과 GNN의 노드 특징을 엔드 투 엔드 학습 체계에서 집계하고 통합한다.
• 프레임워크를 GCN 및 GAT 아키텍처와 통합하여 세 가지 생물학적 데이터셋에서의 노드 분류에 활용한다.
• GAT과 세트 트랜스포머의 어텐션 메커니즘을 활용해 모델 예측을 해석하고 영향력 있는 유전자 및 세포 유형을 식별한다.

실험 결과

연구 질문

• RQ1 자기지도 학습을 통한 엣지 특징 학습이 수작업 또는 레이블링된 엣지 특징이 없는 노드 분류 작업에서 GNN 성능을 향상시킬 수 있는가?
• RQ2 엣지 특징으로서 Forman-Ricci 곡률을 통합할 경우 생물학적 그래프에서 국소 그래프 구조 표현이 어떻게 향상되는가?
• RQ3 GAT과 세트 트랜스포머의 어텐션 메커니즘이 SARS-CoV-2 감염 및 코로나19 심각도의 생물학적 결정 요인에 대해 얼마나 해석 가능한 통찰을 제공하는가?
• RQ4 제안된 프레임워크는 신경질환, in vitro SARS-CoV-2 감염, 인간 환자 샘플을 포함한 다양한 생물학적 그래프 데이터셋에 일반화되는가?
• RQ5 모델이 질병 진행과 심각도와 관련된 생물학적으로 의미 있는 유전자와 세포 유형을 어떻게 식별하여 치료 타겟에 대한 가설 수립을 지원할 수 있는가?

주요 결과

• 제안된 방법은 신경질환, SARS-CoV-2 감염, 인간 코로나19 환자와 관련된 세 가지 단세포 RNA-seq 데이터셋에서 기준 GCN 및 GAT 모델보다 향상된 노드 분류 성능을 달성한다.
• GAT 어텐션 계수에서 유도된 자기지도 엣지 특징은 수동으로 레이블링된 엣지 애너테이션 없이도 모델 훈련과 일반화 능력을 크게 향상시킨다.
• Forman-Ricci 곡률은 국소 그래프 기하학을 효과적으로 캡처하며, 특히 구조적으로 중요한 엣지를 식별하는 데 기여하여 표현 학습을 향상시킨다.
• 세트 트랜스포머를 통한 엣지 특징 통합은 관계적 및 구조적 정보를 노드 수준 표현과 효과적으로 융합시켜 예측 성능을 향상시킨다.
• GAT 및 세트 트랜스포머 레이어의 어텐션 메커니즘이 간섭성 신호 경로와 관련된 특정 유전자와 세포 유형을 성공적으로 강조하며, 이는 SARS-CoV-2 감염 및 질병 심각도 예측에 기여한다.
• 고곡률과 강한 어텐션을 보이는 세포 클러스터를 식별함으로써 면역병리학적 역할을 가질 가능성이 있으며, 향후 실험적 검증을 위한 잠재적 바이오마커와 치료 타겟을 제시한다.