[논문 리뷰] Weisfeiler and Lehman Go Cellular: CW Networks

Cristian Bodnar, Fabrizio Frasca외 5명|arXiv (Cornell University)|2021. 06. 23.

Machine Learning in Materials Science참고 문헌 70인용 수 6

한 줄 요약

이 논문은 정규 세포 복합체로 메시지 전파를 일반화하여 계층적이고 증명 가능하게 더 표현력 있는 학습을 가능하게 하는 새로운 그래프 신경망 프레임워크인 CW 네트워크(CWNs)를 소개한다. 고리 구조를 2차 세포로 사용하여 분자 그래프를 세포 복합체로 올리는 방식으로, ZINC에서 0.079 MAE를 기록하며 분자 그래프 벤치마크에서 최고 성능을 달성하였으며, 복잡한 구조에 대한 강건성도 향상되었다.

ABSTRACT

Graph Neural Networks (GNNs) are limited in their expressive power, struggle with long-range interactions and lack a principled way to model higher-order structures. These problems can be attributed to the strong coupling between the computational graph and the input graph structure. The recently proposed Message Passing Simplicial Networks naturally decouple these elements by performing message passing on the clique complex of the graph. Nevertheless, these models can be severely constrained by the rigid combinatorial structure of Simplicial Complexes (SCs). In this work, we extend recent theoretical results on SCs to regular Cell Complexes, topological objects that flexibly subsume SCs and graphs. We show that this generalisation provides a powerful set of graph "lifting" transformations, each leading to a unique hierarchical message passing procedure. The resulting methods, which we collectively call CW Networks (CWNs), are strictly more powerful than the WL test and not less powerful than the 3-WL test. In particular, we demonstrate the effectiveness of one such scheme, based on rings, when applied to molecular graph problems. The proposed architecture benefits from provably larger expressivity than commonly used GNNs, principled modelling of higher-order signals and from compressing the distances between nodes. We demonstrate that our model achieves state-of-the-art results on a variety of molecular datasets.

연구 동기 및 목표

• 표준 그래프 신경망(GNNs)의 제한된 표현력 문제를 해결할 것. 이는 Weisfeiler-Lehman(WL) 테스트와 동일하거나 그 이하의 표현력이다.
• 단단한 심플리셜 복합체에 의존하는 기존 모델들(예: 메시지 전파 심플리셜 네트워크(MPSNs))의 한계를 극복할 것.
• 세포 복합체를 사용하여 계산 구조와 입력 그래프 구조를 분리하는 유연하고 계층적인 메시지 전파 프레임워크를 개발할 것.
• 분자 그래프 내 화학적 고리와 같은 고차원 구조를 체계적으로 모델링할 수 있도록 할 것.
• 유도 순환 기반의 새로운 올림 변환을 통해 표준 분자 그래프 벤치마크에서 최고 성능을 입증할 것.

제안 방법

• 정규 세포 복합체(CW 복합체) 기반의 메시지 전파 프레임워크를 제안하여, 심플리셜 복합체를 일반화하여 더 유연한 고차원 구조를 허용한다.
• 입력 그래프를 세포 복합체로 변환하기 위한 '올림 변환'을 도입하며, 이는 분자 그래프 내 유도 순환(고리)에 해당하는 2차 세포를 추가함으로써 이루어진다.
• 노드(0차 세포), 엣지(1차 세포), 고리(2차 세포)의 서로 다른 차원의 세포 간에 계층적 메시지 전파를 수행하며, 면 순서집합(face poset)의 이웃 관계에 기반한 업데이트를 수행한다.
• 결과로 생성된 CW 네트워크의 표현력을 이론적으로 분석하기 위해 세포 기반의 Weisfeiler-Lehman 테스트를 사용한다.
• 각 차원에서 세포 표현을 합산(ZINC) 또는 평균화(Mol-HIV)하는 방식으로 집계하는 리드아웃 메커니즘을 구현한다.
• 표준 최적화 프rotocol을 사용하여 모델을 훈련한다: Adam 최적화, 학습률 스케줄링, 조기 정지, 10회 반복한 가중치 초기화 평균화

실험 결과

연구 질문

• RQ1 세포 복합체 기반의 메시지 전파 프레임워크는 표준 GNN보다 더 높은 표현력을 가지는가?
• RQ2 고리와 같은 고차원 구조의 포함 여부가 분자 그래프 학습 과제에서 성능에 어떤 영향을 미치는가?
• RQ3 제안된 CW 네트워크 프레임워크는 WL 테스트보다 엄밀히 강력하며, 3-WL 테스트 이하의 표현력이 아닌가?
• RQ4 유도 순환 기반의 올림 변환은 분자 그래프에서 더 효과적이고 직관적인 계층적 메시지 전파 절차를 가능하게 하는가?
• RQ5 CW 네트워크의 계층적 다차원 메시지 전파 방식은 기존 최고 성능의 GNN보다 표준 분자 그래프 벤치마크에서 뛰어난 성능을 내는가?

주요 결과

• 세포 기반 WL 테스트를 통해 증명된 lin: CW 네트워크는 표준 Weisfeiler-Lehman(WL) 테스트보다 엄밀히 더 표현력이 뛰어나며, 3-WL 테스트 이하의 표현력이 아니며, 이는 이론적으로 보장된다.
• 고리 기반 올림을 사용한 CIN 모델은 ZINC 데이터셋에서 테스트 MAE 0.079 ± 0.006을 기록하며, GSN 및 HIMP를 포함한 모든 이전 베이스라인을 초월한다.
• 제거 실험 결과, 2차 세포(고리)를 제거하면 성능이 저하되며, 'CIN No-Rings'는 0.159 MAE를 기록한다. 이는 고리 간섭 없는 모델인 GIN-E Custom(0.196 MAE)보다도 우수하다.
• 용량이 감소한 모델(48개의 은닉 유닛, 2층)조차도 ZINC에서 0.094 MAE를 달성하여, 파rameter 제약 조건 하에서도 뛰어난 성능을 보인다.
• 엣지 특징과 고리 기반 2차 세포의 포함은 더 rich한 엣지 표현 학습을 가능하게 하며, 이는 'CIN No-Rings'가 'GIN-E Custom'보다 뛰어난 성능을 내는 것으로 확인된다.
• Mol-HIV에서 CIN 모델은 최고 성능을 기록하였으며, 보고된 테스트 AUC가 있으나, 제공된 텍스트에는 정확한 수치가 기재되어 있지 않다.