[논문 리뷰] MLPInit: Embarrassingly Simple GNN Training Acceleration with MLP Initialization

Xiaotian Han, Tong Zhao외 3명|arXiv (Cornell University)|2022. 09. 30.

Advanced Graph Neural Networks인용 수 12

한 줄 요약

MLPInit는 동일한 아키텍처와 가중치 공간을 가진 사전 훈련된 PeerMLP의 가중치로 GNN을 초기화하는 간단하면서도 효과적인 방법을 제안한다. 이 방법은 훈련 시간을 최대 33배 빠르게 하고, 여러 벤치마크에서 노드 분류 작업에서 최대 7.97% 향상되고 링크 예측 작업에서 최대 17.81% 향상된다.

ABSTRACT

Training graph neural networks (GNNs) on large graphs is complex and extremely time consuming. This is attributed to overheads caused by sparse matrix multiplication, which are sidestepped when training multi-layer perceptrons (MLPs) with only node features. MLPs, by ignoring graph context, are simple and faster for graph data, however they usually sacrifice prediction accuracy, limiting their applications for graph data. We observe that for most message passing-based GNNs, we can trivially derive an analog MLP (we call this a PeerMLP) with an equivalent weight space, by setting the trainable parameters with the same shapes, making us curious about extbf{\emph{how do GNNs using weights from a fully trained PeerMLP perform?}} Surprisingly, we find that GNNs initialized with such weights significantly outperform their PeerMLPs, motivating us to use PeerMLP training as a precursor, initialization step to GNN training. To this end, we propose an embarrassingly simple, yet hugely effective initialization method for GNN training acceleration, called MLPInit. Our extensive experiments on multiple large-scale graph datasets with diverse GNN architectures validate that MLPInit can accelerate the training of GNNs (up to 33X speedup on OGB-Products) and often improve prediction performance (e.g., up to $7.97\%$ improvement for GraphSAGE across $7$ datasets for node classification, and up to $17.81\%$ improvement across $4$ datasets for link prediction on metric Hits@10). The code is available at \href{https://github.com/snap-research/MLPInit-for-GNNs}.

연구 동기 및 목표

• 대규모 그래프에서 GNN의 높은 계산 비용과 긴 훈련 시간을 해결한다.
• 메시지 전달 GNN과 MLP의 가중치 공간이 동일하게 만들 수 있는지 조사한다.
• 먼저 PeerMLP를 훈련하고, 수렴한 가중치를 사용해 GNN을 초기화하면 훈련을 가속화하고 성능을 향상시킬 수 있는지 탐색한다.
• 아키텍처 변경이나 하이퍼파라미터 튜닝 없이도 간단하고 즉시 사용 가능한 초기화 방법을 개발한다.
• PeerMLP 가중치로 초기화된 GNN이 무작위 초기화된 GNN과 그 자체 PeerMLP보다도 성능이 뛰어나다는 것을 입증한다.

제안 방법

• 주어진 GNN과 동일한 아키텍처와 가중치 차원을 가진 MLP를 PeerMLP로 정의하고, 동일한 학습 가능한 파rameter 공간을 공유한다.
• 희소 메시지 전달을 피하고 노드 특성만을 사용하여 동일한 그래프 데이터에서 PeerMLP를 훈련시켜 더 빠른 수렴을 이룬다.
• 수렴한 PeerMLP의 가중치를 해당 GNN의 초기 가중치로 사용한다.
• 이 초기화 상태에서 GNN을 미세 조정하여 사전에 학습된 특성 표현을 활용해 수렴 속도를 가속화한다.
• 다양한 GNN 아키텍처(예: GraphSAGE, GCN, ClusterGCN)와 대규모 데이터셋(예: OGB-products, Reddit, Flickr)에 이 방법을 적용한다.
• 무작위 초기화와의 공정한 비교를 위해 표준 훈련 프로토콜을 사용하고 조기 정지 및 검증 기반 모델 선택을 적용한다.

실험 결과

연구 질문

• RQ1 GNN과 해당 PeerMLP의 가중치 공간은 형태와 차원 측면에서 동일하게 만들 수 있는가?
• RQ2 노드 특성만을 사용해 PeerMLP를 훈련하면, 이 표현을 GNN으로 전이했을 때 무작위 초기화보다 하류 작업 성능이 향상되는가?
• RQ3 PeerMLP의 수렴한 가중치가 GNN의 효과적인 초기화로 기능할 수 있는가? 이는 수렴에 필요한 훈련 에포크 수를 줄일 수 있는가?
• RQ4 이 초기화 방법은 다양한 GNN 아키텍처와 대규모 그래프 데이터셋에서 일관되게 훈련을 가속화하는가?
• RQ5 PeerMLP 가중치로 초기화된 GNN의 성능이 노드 분류 및 링크 예측 작업에서 PeerMLP 자체의 성능을 초월하는가?

주요 결과

• MLPInit는 랜덤 초기화 대비 OGB-products 데이터셋에서 GNN 훈련 시간을 최대 33배 빠르게 한다.
• GraphSAGE는 MLPInit를 사용해 일곱 개의 데이터셋에서 노드 분류 정확도를 최대 7.97% 향상시킨다.
• 링크 예측 작업에서는 MLPInit가 네 개의 데이터셋에서 Hits@10 성능을 최대 17.81% 향상시킨다.
• PeerMLP 가중치로 초기화된 GNN은 랜덤 초기화된 대비 반 이하의 에포크 수로 수렴한다.
• PeerMLP 자체는 수렴한 가중치로 초기화되었을 때 GNN보다 성능이 열 劣하다. 이는 GNN의 메시지 전달 메커니즘이 전달된 표현에서 유의미한 이점을 얻기 때문이다.
• 이 방법은 다양한 GNN 아키텍처(예: GraphSAGE, GCN, ClusterGCN, GraphSAINT)와 데이터셋(예: Reddit, Flickr, OGB-arXiv, OGB-products)에 대해 강건하다.