[논문 리뷰] MAG-GNN: Reinforcement Learning Boosted Graph Neural Network

Lecheng Kong, Jiarui Feng외 4명|arXiv (Cornell University)|2023. 10. 29.

Advanced Graph Neural Networks인용 수 5

한 줄 요약

MAG-GNN는 그래프 신경망(GNNs)의 표현력을 높이기 위해 모든 하위그래프를 순열로 나열하는 데 드는 계산 비용을 줄이기 위해 강화학습(RL) 기반 방법을 제안한다. 하위그래프의 최소 표현 집합을 식별함으로써, MAG-GNN는 여러 그래프 벤치마크에서 최신 기술 수준의 성능을 달성하면서도 하위그래프 GNN보다 인퍼런스 시간을 크게 단축시켰다.

ABSTRACT

While Graph Neural Networks (GNNs) recently became powerful tools in graph learning tasks, considerable efforts have been spent on improving GNNs' structural encoding ability. A particular line of work proposed subgraph GNNs that use subgraph information to improve GNNs' expressivity and achieved great success. However, such effectivity sacrifices the efficiency of GNNs by enumerating all possible subgraphs. In this paper, we analyze the necessity of complete subgraph enumeration and show that a model can achieve a comparable level of expressivity by considering a small subset of the subgraphs. We then formulate the identification of the optimal subset as a combinatorial optimization problem and propose Magnetic Graph Neural Network (MAG-GNN), a reinforcement learning (RL) boosted GNN, to solve the problem. Starting with a candidate subgraph set, MAG-GNN employs an RL agent to iteratively update the subgraphs to locate the most expressive set for prediction. This reduces the exponential complexity of subgraph enumeration to the constant complexity of a subgraph search algorithm while keeping good expressivity. We conduct extensive experiments on many datasets, showing that MAG-GNN achieves competitive performance to state-of-the-art methods and even outperforms many subgraph GNNs. We also demonstrate that MAG-GNN effectively reduces the running time of subgraph GNNs.

연구 동기 및 목표

• 표현력을 높이기 위해 모든 가능한 하위그래프를 순열로 나열하는 데 드는 높은 계산 비용을 해결하기 위해.
• 완전한 하위그래프 순열과 유사한 표현력을 가지는 소규모 최적의 하위그래프 집합이 존재하는지 조사하기 위해.
• 완전한 탐색 없이도 높은 GNN 표현력을 유지하면서도 효율적이고 확장 가능한 방법을 개발하기 위해.
• 하위그래프 GNN의 인퍼런스 시간을 단축시키면서도 예측 성능을 유지하거나 향상시키기 위해.

제안 방법

• MAG-GNN는 입력 그래프에서 루트를 가진 하위그래프의 무작위로 샘플된 후보 집합으로 시작한다.
• RL 에이전트는 모델 성능 기반의 보상 신호를 최대화하기 위해 후보 집합 내 하위그래프를 반복적으로 교체하도록 훈련된다.
• 에이전트는 목표 하위그래프를 새로운 후보 하위그래프로 교체할 때의 기대 보상을 추정하기 위해 Q-테이블을 사용한다.
• 하위그래프 교체 과정은 각 단계에서 가장 높은 기대 보상을 갖는 하위그래프 업데이트를 선택하는 정책에 의해 이끌린다.
• 최종적으로 선택된 하위그래프 집합은 예측 작업을 위한 후행 GNN의 입력으로 사용된다.
• 이 방법은 지수적 하위그래프 순열 복잡도를 상수 단계의 RL 탐색으로 줄여, 효율적이고 표현력 있는 GNN를 가능하게 한다.

실험 결과

연구 질문

• RQ1 소규모 최적의 하위그래프 집합이 GNN에서 전체 하위그래프 순열과 동일한 표현력을 달성할 수 있는가?
• RQ2 강화학습이 GNN 성능 향상에 기여하는 가장 분류 능력이 높은 하위그래프를 효과적으로 식별할 수 있는가?
• RQ3 하위그래프 순열을 상수 단계의 RL 과정으로 줄였을 때, GNN 성능은 유지되거나 향상되며 런타임은 감소하는가?
• RQ4 정확도와 효율성 측면에서 MAG-GNN는 최신 기술 수준의 하위그래프 GNN과 비교해 어떻게 성과를 내는가?

주요 결과

• MAG-GNN는 최신 기술 수준의 방법들과 경쟁 가능한 성능을 달성했으며, ZINC 및 OGBG-MOLHIV를 포함한 여러 데이터셋에서 몇몇 하위그래프 GNN을 능가했다.
• ZINC 데이터셋에서 MAG-GNN는 테스트 손실 0.106 ± 0.014를 기록했으며, GIN(0.163 ± 0.004)을 크게 능가했고, 더 복잡한 모델인 CIN(0.079 ± 0.006)의 성능에 가까이 다가섰다.
• 사전학습(MAG-GNN-PRE)을 적용한 경우 성능이 더욱 향상되어 ZINC에서 테스트 손실 0.096 ± 0.009, OGBG-MOLHIV에서 78.30 ± 1.08의 성능을 기록했다.
• MAG-GNN는 인퍼런스 시간을 크게 단축시켰으며, I²GNN와 같은 엣지 루트 하위그래프 GNN보다 빠르며, 유사한 효율성에도 불구하고 NGNN를 정확도에서 능가했다.
• 모델은 다양한 기본 GNN 아키텍처에 적응 가능성을 보이며, 아키텍처 통합을 통해 추가 성능 향상 가능성을 시사한다.
• 아블레이션 연구는 MAG-GNN의 RL 기반 하위그래프 선택이 표현력과 효율성 사이의 균형을 효과적으로 유지하며, 고정된 아키텍처를 가진 하위그래프 GNN에서 관찰되는 과적합을 피한다는 것을 확인했다.