Cooperative Graph Neural Network
마지막으로 over-squashing 현상을 해결하는 방법으로 cooperative graph neural network에 대해서 알아보자. 이 방법은 기존의 graph rewiring을 manual하게 하는 것이 아니라 이 과정을 학습하고자 한다. 먼저 각 vertex의 state를 부여해준다. 여기서 state는 각 vertex가 이웃과 연결되어 정보를 주고 받을 수 있는지 아니면 고립되어 있는지를 나타낸다. 각 vertex를 하나의 player로 설정해서 'listen', 'broadcast', 'listen and broadcast', 'isolate'와 같은 상태를 할당받게 되는 것이다. 기존의 message propagation의 경우 모든 node가 이웃들에게 'listen and broadcast'를 하는 상태를 부여받은 특수한 경우에 해당하게 된다. 이와 다르게 저자들은 좀더 flexible하고 dynamic한 message passing을 위해서 모든 node가 주어진 state에 맞게 propagation 전략을 세울 수 있도록 만들었다. 이렇게해서 network의 각 layer마다 message passing 구조를 조절할 수 있도록 만들었고, 이를 통해 over-squashing 현상을 해결하고자 하였다.
제목이 cooperative graph neural network인 이유가 각 vertex가 다른 vertex와 서로 협력하여 squashed된 정보를 피하고자 하기 위해서이다.
각 layer마다 node들은 discrete한 결정을 취하게 되며, 어떻게보면 이러한 방식이 reinforcement learning과 유사해보일 수 있다. 각 vertex마다 확률적으로 특정 행동을 취해야하는 점은 reinforcement learning의 방식을 따르고 있다. STANDARD는 동시에 message를 처리하여 정보를 만들어 전달하는 것을 의미한다. LISTEN은 오로지 input으로 listen만 받게 된다. 반대로 BROADCAST는 listen을 밖으로 보내는 것을 말한다. 마지막으로 ISOLATE는 이웃들과 어떠한 행동도 취하지 않는 것을 의미한다.
Input으로 node feature가 주어졌을 때 policy network가 각 state에 따른 확률들을 각 node에 할당하게 된다. 각 state의 확률로부터 해당 node가 어떠한 state를 가지게 되는지 결정할 수 있다. Message passing이 BROADCAST나 ISOLATE의 경우에는 끊어지게 되며, 반대로 STANDARD나 LISTEN일 때에는 연결되어 정보를 aggregation하게 된다. 만약 여기서 network를 학습시키고 싶으면 discrete optimization을 수행해주면 된다. 이는 위와 같이 straight-through Gumbel-softmax estimator를 사용해주면 되는데, 이는 discrete variable들을 학습해준다. Cooperative graph neural network는 실험 결과도 굉장히 흥미롭지만 아이디어 자체만으로도 충분히 재미있는 방법이다.
이 논문에서는 universal approximation theorem을 위와같이 이야기하고 있다. 이는 expressiveness power를 정량화하는데, 만약 충분하게 전체 graph를 평균화시키면 어떠한 graph function도 approximation할 수 있음을 이야기한다. 이러한 아이디어는 GNN에서 random augmentation과 비슷하다. 만약 무작위로 graph를 augment한다면 GNN의 expressiveness power가 올라가게 된다.
