[논문 리뷰] Probabilistically Rewired Message-Passing Neural Networks

Chendi Qian, Andrei Manolache외 5명|arXiv (Cornell University)|2023. 10. 03.

Advanced Graph Neural Networks인용 수 4

한 줄 요약

이 논문은 정확하고 미분 가능한 k-하위집합 샘플링을 사용하여 작업에 관련된 사전 지식에 기반해 그래프 내에서 간선을 추가하거나 제거하는 것을 학습하는 확률적 재연결 메시지 전파 신경망(PR-MPNNs)을 제안한다. 이 방법은 표현력 향상, 과도한 압축 및 부족한 정보 전달 완화를 가능하게 하며, 분자의 그래프 및 인용 그래프 벤치마크에서 최신 기술 수준 또는 경쟁 가능한 성능을 달성하여 표준 MPNNs와 그래프 트랜스포머를 뛰어넘는 성과를 보인다.

ABSTRACT

Message-passing graph neural networks (MPNNs) emerged as powerful tools for processing graph-structured input. However, they operate on a fixed input graph structure, ignoring potential noise and missing information. Furthermore, their local aggregation mechanism can lead to problems such as over-squashing and limited expressive power in capturing relevant graph structures. Existing solutions to these challenges have primarily relied on heuristic methods, often disregarding the underlying data distribution. Hence, devising principled approaches for learning to infer graph structures relevant to the given prediction task remains an open challenge. In this work, leveraging recent progress in exact and differentiable $k$-subset sampling, we devise probabilistically rewired MPNNs (PR-MPNNs), which learn to add relevant edges while omitting less beneficial ones. For the first time, our theoretical analysis explores how PR-MPNNs enhance expressive power, and we identify precise conditions under which they outperform purely randomized approaches. Empirically, we demonstrate that our approach effectively mitigates issues like over-squashing and under-reaching. In addition, on established real-world datasets, our method exhibits competitive or superior predictive performance compared to traditional MPNN models and recent graph transformer architectures.

연구 동기 및 목표

• 고정된 그래프 구조로 인해 발생하는 메시지 전파 GNN의 한계, 예를 들어 과도한 압축과 제한된 표현력 문제를 해결하기 위해.
• 예측 작업에 맞게 조정되는 원칙적이고 미분 가능한 그래프 구조 수정 방법을 개발하기 위해.
• 데이터 분포에 적응하지 않거나 정당성 없이 무작위 간선을 추가하는 히우리스틱 재연결 방법의 단점을 보완하기 위해.
• 확률적 재연결이 모델 표현력을 어떻게 향상시키는지 이론적으로 분석하고, 랜덤 재연결보다 성능이 뛰어나지 않는 조건을 규명하기 위해.
• 실제 데이터셋에서 PR-MPNN이 MPNN 및 그래프 트랜스포머와 비교해 우수하거나 경쟁 가능한 성능을 내는지 경험적으로 검증하기 위해.

제안 방법

• PR-MPNNs는 후행 모델을 통해 후보 간선에 대한 확률 분포를 학습하며, 정확히 k개의 간선으로 조건화된 k-하위집합 분포로 파arameter화한다.
• 이 프레임워크는 학습된 분포에서 여러 개의 k-간선 인접 행렬을 샘플링하고, 이를 집계(예: 뺄셈을 통해)하여 최종적인, 미분 가능한 그래프 구조를 형성한다.
• 역전파 중에는 k-하위집합 분포의 정확하고 미분 가능한 마진 분포를 사용하여 학습 중 편향과 분산을 감소시킨다.
• 최근의 정확하고 미분 가능한 k-하위집합 샘플링 기술을 활용하여 간선 선택 메커니즘의 엔드 투 엔드 학습을 가능하게 한다.
• 최종 그래프 구조는 예측을 위한 하류 MPNN의 입력으로 사용되며, 이로써 구조 최적화와 표현 학습이 동시에 이루어진다.
• 이 방법은 확장성과 미분 가능성에 기반하여 확률적 간선 선택 과정을 효과적으로 역전파할 수 있도록 설계되어 있다.

실험 결과

연구 질문

• RQ1 고정된 또는 무작위로 재연결된 그래프와 비교해, 확률적 재연결이 메시지 전파 GNN의 표현력을 향상시키는 조건은 무엇인가?
• RQ2 미분 가능한 k-하위집합 샘플링 메커니즘이 간선 선택 정책의 효과적이고 저분산 학습을 어떻게 가능하게 하는가?
• RQ3 PR-MPNN은 그래프 신경망에서 과도한 압축과 정보 부족 현상을 어느 정도 완화하는가?
• RQ4 실제 그래프 벤치마크에서 PR-MPNN은 그래프 트랜스포머 및 최신 기술 수준의 MPNN보다 성능이 뛰어나거나 경쟁 가능한가?
• RQ5 학습된 간선 선택 정책은 데이터 분포와 예측 작업에 어떻게 적응하는가?

주요 결과

• LRGB 벤치마크의 Peptides-struct 데이터셋에서, PR-MPNN는 0.2477 ± 0.0005의 최고 MAE 스코어를 기록하여 DRew-GIN(0.2799 ± 0.0016) 및 기타 베이스라인을 모두 앞섰다.
• 이질적이고 전이 가능한 WebKB 데이터셋에서, PR-MPNN는 Cornell(0.659 ± 0.040), Texas(0.827 ± 0.032), Wisconsin(0.750 ± 0.015)에서 가장 높은 정확도를 기록했으며, 모든 MPNN 기반 모델을 뛰어넘고 그래프 트랜스포머와도 경쟁했다.
• 장거리 간선을 추가하여 정보 흐름을 향상시키는 방식으로 과도한 압축과 정보 부족 현상을 크게 감소시켜, 전반적인 구조적 전역 정보가 필요한 작업에서 성능 향상을 입증했다.
• 간격 성질 예측 작업에서 기반 GIN 모델 대비 평균 절대 오차(MAE)에서 1.98배 향상된 성과를 기록했다(1.70 ± 0.01 vs. 3.37 ± 0.03).
• ZPVE 성질 예측에서 기반 GIN 모델 대비 14.13배 향상된 성능(4.73 ± 0.08 vs. 66.87 ± 1.45)을 기록하여 정보 압축 현상을 강력하게 완화함을 보였다.
• 이론적 분석을 통해 PR-MPNNs는 학습된 간선의 분포와 작업에 관련된 구조와의 일치 조건이 정확히 충족될 경우, 표준 MPNN보다 더 높은 표현력을 달성할 수 있음을 확인했다.