[논문 리뷰] A Unified Lottery Ticket Hypothesis for Graph Neural Networks

Tianlong Chen, Yongduo Sui외 3명|arXiv (Cornell University)|2021. 02. 12.

Advanced Graph Neural Networks참고 문헌 71인용 수 28

한 줄 요약

이 논문은 Unified GNN Sparsification (UGS)을 도입하여 그래프와 GNN 가중치를 함께 가지치고, Lottery Ticket Hypothesis를 그래프 신경망에 일반화하여 Graph Lottery Tickets (GLTs)를 제공하며 추론 비용을 크게 절감하고 성능을 유지한다.

ABSTRACT

With graphs rapidly growing in size and deeper graph neural networks (GNNs) emerging, the training and inference of GNNs become increasingly expensive. Existing network weight pruning algorithms cannot address the main space and computational bottleneck in GNNs, caused by the size and connectivity of the graph. To this end, this paper first presents a unified GNN sparsification (UGS) framework that simultaneously prunes the graph adjacency matrix and the model weights, for effectively accelerating GNN inference on large-scale graphs. Leveraging this new tool, we further generalize the recently popular lottery ticket hypothesis to GNNs for the first time, by defining a graph lottery ticket (GLT) as a pair of core sub-dataset and sparse sub-network, which can be jointly identified from the original GNN and the full dense graph by iteratively applying UGS. Like its counterpart in convolutional neural networks, GLT can be trained in isolation to match the performance of training with the full model and graph, and can be drawn from both randomly initialized and self-supervised pre-trained GNNs. Our proposal has been experimentally verified across various GNN architectures and diverse tasks, on both small-scale graph datasets (Cora, Citeseer and PubMed), and large-scale datasets from the challenging Open Graph Benchmark (OGB). Specifically, for node classification, our found GLTs achieve the same accuracies with 20%~98% MACs saving on small graphs and 25%~85% MACs saving on large ones. For link prediction, GLTs lead to 48%~97% and 70% MACs saving on small and large graph datasets, respectively, without compromising predictive performance. Codes available at https://github.com/VITA-Group/Unified-LTH-GNN.

연구 동기 및 목표

• 대규모 그래프에서 GNN 추론의 계산 비용을 그래프 구조와 모델 매개변수를 공동으로 가지치기함으로써 감소시키는 것을 목표로 한다.
• 그래프 Lottery Ticket Hypothesis를 GNN에 일반화하여 희소 부분 그래프와 희소 부분 네트워크를 결합한 Graph Lottery Tickets (GLTs)를 정의한다.
• GLTs가 무작위 초기화 및 자기지도 사전학습을 통해 다양한 GNN 아키텍처와 작업에서 발견될 수 있음을 보여준다.
• UGS/GLTs의 대형 데이터셋에서의 확장성 및 희소성 수준과 하드웨어 영향에 대한 통찰을 제공한다.

제안 방법

• 그래프(m_g)와 모델 가중치(m_theta)에 대한 미분 가능 마스크를 사용한 Unified GNN Sparsification (UGS)을 제안하여 A와 Θ를 공동으로 가지치기하고 손실 L_UGS = L({m_g ⊙ A, X}, m_theta ⊙ Θ) + γ1||m_g||1 + γ2||m_theta||1를 최적화한다.
• 반복 가지치기(Algorithm 2)를 사용하여 UGS를 반복 적용하고 원래 초기화로 가중치를 되돌려서 GLT를 찾고, 그래프 희소성 s_g와 모델 희소성 s_theta를 가진 GLT를 얻는다.
• GLT를 (희소 그래프 G_s, 희소 마스크 m_theta, 초기화 Θ0)의 삼중항으로 정의하여 전체 그래프와 밀집 모델로 학습했을 때의 성능에 상응하는 것을 목표로 한다.
• 작은 그래프(Cora, Citeseer, PubMed) 및 큰 그래프(OGB)에서의 GCN/GIN/GAT 아키텍처와 노드 분류 및 링크 예측 작업에 대해 GLT를 평가한다.

실험 결과

연구 질문

• RQ1 통합 희소화 프레임워크가 그래프 구조와 GNN 가중치를 성능 저하 없이 모두 가지치기할 수 있는가?
• RQ2 Graph Lottery Tickets가 서로 다른 GNN 아키텍처와 작업에서 존재하며 무작위 초기화나 자기지도 사전학습에서 발견될 수 있는가?
• RQ3 그래프와 모델의 희소도 수준이 소그래프와 대형 그래프의 노드 분류 및 링크 예측에 어떤 영향을 미치는가?
• RQ4 UGS가 대규모 데이터셋의 깊은 GNN으로 확장 가능하며 그래프 밀도가 가지치기에 대한 견고성에 어떤 영향을 미치는가?

주요 결과

• GLTs는 Cora, Citeseer, PubMed에서 GCN, GIN, GAT에 대해 대부분의 MACs 절감과 성능 저하 거의 없이 존재한다.
• GLTs는 노드에 대해 20%–98% MACs 절감, 대형 그래프에 대해 25%–85% MACs 절감, 최대 58.19% 그래프 희소성과 97.75% 모델 희소성을 달성한다.
• 링크 예측의 경우 소형 그래프에서 48%–97% MACs 절감, 대형 그래프에서 70% 절감을 달성하며 정확도 손실이 없다.
• UGS는 임의 가지치기보다 우수하고 ADMM과 비교해도 경쟁력이 있으며, 대형 그래프 및 최신 GNN 변형으로의 확장성에 있어 이점을 가진다.
• GLTs는 자기지도 사전학습(GraphCL)에서 도출될 수 있으며 많은 설정에서 임의 초기화에 비해 견고성 및 성능 향상을 제공할 수 있다.
• UGS는 대형 그래프의 28층 ResGCN에도 확장되어 상당한 MACs 절감을 달성(예: 일부 데이터세트에서 최대 85%)하며 성능을 유지한다.