[논문 리뷰] Sketch-GNN: Scalable Graph Neural Networks with Sublinear Training Complexity

Mucong Ding, Tahseen Rabbani외 3명|arXiv (Cornell University)|2024. 06. 21.

Advanced Graph Neural Networks인용 수 7

한 줄 요약

Sketch-GNN은 인접 행렬과 노드 특징의 압축 스케치에서 GNN을 훈련시켜 그래프 크기에 비례하지 않는 차등적/부분선형 학습 시간과 메모리 사용을 달성하면서도 경쟁력 있는 정확도를 유지합니다. 비선형 활성화 함수와 메시지 전달을 근사하기 위해 다항 텐서 스케치와 학습 가능한 LSH를 사용합니다.

ABSTRACT

Graph Neural Networks (GNNs) are widely applied to graph learning problems such as node classification. When scaling up the underlying graphs of GNNs to a larger size, we are forced to either train on the complete graph and keep the full graph adjacency and node embeddings in memory (which is often infeasible) or mini-batch sample the graph (which results in exponentially growing computational complexities with respect to the number of GNN layers). Various sampling-based and historical-embedding-based methods are proposed to avoid this exponential growth of complexities. However, none of these solutions eliminates the linear dependence on graph size. This paper proposes a sketch-based algorithm whose training time and memory grow sublinearly with respect to graph size by training GNNs atop a few compact sketches of graph adjacency and node embeddings. Based on polynomial tensor-sketch (PTS) theory, our framework provides a novel protocol for sketching non-linear activations and graph convolution matrices in GNNs, as opposed to existing methods that sketch linear weights or gradients in neural networks. In addition, we develop a locality-sensitive hashing (LSH) technique that can be trained to improve the quality of sketches. Experiments on large-graph benchmarks demonstrate the scalability and competitive performance of our Sketch-GNNs versus their full-size GNN counterparts.

연구 동기 및 목표

• 매우 큰 그래프에서 전체 그래프 메모리 없이 또는 그래프 크기에 대한 선형 시간 의존성 없이 확장 가능한 GNN 학습을 촉진한다.
• 학습 전에 인접 행렬과 특징 행렬을 압축하는 스케치 기반 프레임워크를 도입한다.
• GNN에서 비선형 활성화를 근사하기 위해 다항 텐서 스케치를 확장한다.
• 학습 가능한 locality-sensitive hashing를 개발하여 학습 중 스케치를 조정하고 정확도를 향상시킨다.

제안 방법

• 메시지 전달을 CXW 위의 비선형 연산으로 표현하고 CXW를 다항 텐서 스케치를 통해 근사한다.
• C와 XW를 작은 행렬로 스케치하기 위해 count sketch와 tensor sketch를 사용하여 부분선형 메모리와 시간을 가능하게 한다.
• 비선형성을 멱급수로 확장하고 계수 ck를 학습하여 활성화를 근사한다.
• 진화하는 은닉 표현에 맞추어 스케치를 적응시키기 위해 학습 가능한 SimHash 기반의 LSH를 도입한다.
• 레이어 간에 스케치 표현의 효율적 계산을 가능하게 하는 양면 스케치 프레임워크를 제공한다.
• 레이어 간 오류 누적을 완화하기 위해 스킵 연결을 도입한다.

실험 결과

연구 질문

• RQ1 인접 행렬과 특징 행렬을 스케치하여 그래프 크기에 대해 부분선형 시간 및 메모리로 GNN 학습을 수행할 수 있는가?
• RQ2 전체 공간으로 언스케치하지 않고도 스케치 기반 프레임워크 내에서 GNN의 비선형 활성화를 어떻게 근사할 수 있는가?
• RQ3 학습 가능한 LSH가 학습 중 스케치의 품질을 향상시키고 경쟁력 있는 정확도를 유지할 수 있는가?
• RQ4 일반적인 GNN 아키텍처 전반에서 Sketch-GNN의 메모리, 시간, 정확도 간의 이론적 및 실증적trade-off는 무엇인가?

주요 결과

• Sketch-GNN은 대형 벤치마크에서 경쟁력 있는 정확도를 유지하면서 그래프 크기에 대해 부분선형 학습 복잡도를 달성한다.
• 프레임워크는 n×n 인접 행렬과 n×d 특징 행렬을 훨씬 더 작은 스케치로 스케치하여 각 레이어당 O(c) 메모리와 시간(단, c는 n과 무관)을 가능하게 한다.
• Polynomial tensor sketch (PTS)는 CX와 XW의 텐서/카운트 스케치와 결합하여 비선형 활성화를 근사하는 데 사용된다.
• 학습 가능한 LSH 접근 방식은 학습 중 해시 테이블을 조정하여 스케치 품질을 향상시키고 성능 손실을 줄인다.
• 스킵 연결은 스케치 기반 파이프라인에서 레이어 간의 오류 누적을 완화하는 데 도움을 준다.
• 실험 결과는 전체 그래프 GNN과의 경쟁력 있는 성능과 대형 그래프에서의 우수한 확장성을 보여준다.