[논문 리뷰] A Generalization of ViT/MLP-Mixer to Graphs

Xiaoxin He, Bryan Hooi외 4명|arXiv (Cornell University)|2022. 12. 27.

Advanced Graph Neural Networks인용 수 7

한 줄 요약

이 논문은 그래프를 패치로 분할하고 토큰 믹싱 및 채널 믹싱 메커니즘을 적용함으로써 비전 트랜스포머와 MLP-Mixer를 그래프로 일반화한 새로운 GNN 아키텍처인 Graph ViT/MLP-Mixer를 제안한다. 이는 노드와 엣지의 선형 시간 복잡도를 달성하고 과도한 압축을 완화하며, 장거리 의존성 및 표현력 벤치마크에서 MP-GNNs와 그래프 트랜스포머를 능가한다.

ABSTRACT

Graph Neural Networks (GNNs) have shown great potential in the field of graph representation learning. Standard GNNs define a local message-passing mechanism which propagates information over the whole graph domain by stacking multiple layers. This paradigm suffers from two major limitations, over-squashing and poor long-range dependencies, that can be solved using global attention but significantly increases the computational cost to quadratic complexity. In this work, we propose an alternative approach to overcome these structural limitations by leveraging the ViT/MLP-Mixer architectures introduced in computer vision. We introduce a new class of GNNs, called Graph ViT/MLP-Mixer, that holds three key properties. First, they capture long-range dependency and mitigate the issue of over-squashing as demonstrated on Long Range Graph Benchmark and TreeNeighbourMatch datasets. Second, they offer better speed and memory efficiency with a complexity linear to the number of nodes and edges, surpassing the related Graph Transformer and expressive GNN models. Third, they show high expressivity in terms of graph isomorphism as they can distinguish at least 3-WL non-isomorphic graphs. We test our architecture on 4 simulated datasets and 7 real-world benchmarks, and show highly competitive results on all of them. The source code is available for reproducibility at: \url{https://github.com/XiaoxinHe/Graph-ViT-MLPMixer}.

연구 동기 및 목표

• 메시지 전파 GNN의 한계, 특히 과도한 압축과 열악한 장거리 의존성 문제를 해결한다.
• 그래프 트랜스포머의 제곱형 복잡도를 극복하면서도 표현력을 유지하거나 향상시킨다.
• 누적된 메시지 전파 레이어에 의존하지 않고도 효율적으로 장거리 의존성을 포착할 수 있는 GNN 아키텍처를 개발한다.
• 1-WL을 초월하는 그래프 이sovomorphism 표현력을 향상시켜 더 복잡한 비이sovomorphism 그래프를 구별할 수 있도록 한다.
• 다양한 그래프 벤치마크, 특히 대규모 및 희소 그래프에서 선형 시간 및 메모리 복잡도를 달성하면서도 높은 성능을 내는 것을 목표로 한다.

제안 방법

• 메시지 전파 GNN의 한계를 해결하기 위해 METIS 그래프 분할 기법을 사용하여 그래프를 P개의 패치로 분할하고, 연결성을 유지하기 위해 한 단계 이내로 겹침을 도입한다.
• 각 패치에 기본 GNN(예: GCN, GIN)을 적용하여 패치 임베딩을 계산하며, 표준 GNN 대비 일정한 오버헤드를 유발한다.
• 패치 간 장거리 의존성을 모델링하기 위해 크로스-패치 어텐션 또는 피드포워드 연산을 수행하는 토큰 믹서 레이어를 적용한다.
• 각 패치 내에서 노드 표현을 개선하기 위해 채널 믹서 레이어를 적용하여 기능 수준의 상호작용을 가능하게 한다.
• 하류 작업을 위해 패치 표현을 집계하기 위해 글로벌 분류 헤드를 사용한다.
• MLP-Mixer 스타일 아키텍처를 활용하여 공간적(패치 기반) 및 채널 기반 특징 학습을 분리함으로써 효율적인 글로벌 모델링을 가능하게 한다.

실험 결과

연구 질문

• RQ1 비전 트랜스포머/MLP-Mixer 기반 아키텍처가 그래프 구조 데이터에 효과적으로 일반화될 수 있는가?
• RQ2 제안된 Graph ViT/MLP-Mixer가 표준 MP-GNNs와 그래프 트랜스포머보다 더 나은 장거리 의존성 모델링 성능을 보일 수 있는가?
• RQ3 1-WL 및 2-WL GNN보다 높은 표현력을 확보하면서도 선형 시간 및 메모리 복잡도를 유지할 수 있는가?
• RQ4 TreeNeighbourMatch 및 Long Range Graph Benchmark와 같은 장거리 의존성과 구조적 복잡성이 높은 벤치마크에서 모델의 성능은 어떠한가?
• RQ5 모델 성능이 패치 수와 그래프 분할 전략에 얼마나 의존하는가?

주요 결과

• Peptides-func에서 평균 정밀도에 0.056의 평균 향상이 있었고, Peptides-struct에서는 MAE가 0.028 감소하여 장거리 상호작용 모델링에서 뛰어난 성능을 입증했다.
• TreeNeighbourMatch 데이터셋에서 모델은 깊이 r=7까지 일반화 가능했지만, 표준 MP-GNN은 r=4에서 실패하여 과도한 압축 완화 효과를 입증했다.
• SUN과 유사한 성능(0.6730 on Peptides-func, 0.2498 on Peptides-struct)을 달성했지만, 메모리 사용량은 44배 줄이고 학습 속도는 19배 빠르게 하여 효율성을 입증했다.
• Graph MLP-Mixer는 높은 표현력을 보이며 적어도 3-WL 비이sovomorphism 그래프를 구별할 수 있어 표준 MP-GNN을 능가한다.
• 모델은 시간 및 메모리 복잡도 모두 선형 복잡도 O(N + E)를 유지하며, 그래프 트랜스포머 및 SUN과 같은 표현력 있는 GNN보다 확장성 면에서 뚜렷한 우수성을 보였다.
• 4개의 시뮬레이션 및 7개의 실세계 벤치마크에서의 실험 결과는 모든 작업에서 경쟁력 있는 성능을 보였으며, 아키텍처의 강건성과 일반화 능력을 검증했다.