Unified Graph Transformer (UGT) 연구에 대한 회고

Disassortative 및 Heterophilic 그래프 핸들링에 대한 고민들


read the paper read Thuy's Post get the code


이 글에서는 우리 연구실(Network Science Lab @ CUK)의 Van Thuy Hoang 박사과정과 함께 진행하여 AAAI 2024에서 발표한 Unified Graph Transformer (UGT) 연구Van Thuy Hoang, O-Joun Lee: Transitivity-Preserving Graph Representation Learning for Bridging Local Connectivity and Role-based Similarity. The 38th AAAI Conference on Artificial Intelligence (AAAI 2024), Vancouver, Canada; 02/2024. DOI:10.1609/aaai.v38i11.29138의 배경과 설계 원칙, 그리고 논문에는 포함되지 못했던 근본적인 개념들에 대해 회고하고자 한다.

Motivation

우리 연구는 기존 Graph Neural Networks (GNNs)와 Graph Transformers (GTs)의 주된 작동 원리인 Node Homophily 기반의 Message Passing Mechanism이 실제로 Disassortative하거나 Heterophilic한 현실 그래프 환경에서도 효과적으로 작동할 수 있는지에 대한 의문에서 시작되었다. 특히, 우리가 중점적으로 다룬 문제는 비슷한 구조적 특성을 가진 노드들이 유사한 벡터 표현(Vector Representation)을 갖도록 하는 Structural Distinctiveness 문제였다. 기존의 연구들은 주로 구조적으로 유사한 노드 사이에 Pseudo Edge를 추가하는 방법을 제안해왔지만, 우리는 이러한 방식이 오히려 원본 그래프의 구조를 왜곡할 위험이 있다고 판단했다. 즉, 추가된 Edge는 구조적 유사성을 반영하기 위한 것이지만, 실제로는 Graph Encoder의 구조적 식별력(structural discriminability)을 오히려 저하시킬 가능성이 있었다.

Bridging Local Connectivity and Structural Similarity

우리는 이 문제의 본질적인 해법을 Local Connectivity 정보와 Structural Similarity 정보를 동시에 잘 반영하는 접근법에서 찾아야 한다고 생각했다. 이를 위해 제안한 개념이 바로 k-step Transition Probability이다. k-step Transition Probability는 두 노드 사이에 k-hop 이내의 모든 경로 정보를 반영하는 개념으로, 이는 경로의 개수뿐 아니라 노드들의 Degree에 따른 정보 전파의 확률과 감쇄 효과를 포함하여 주변의 구조적 특성을 자연스럽게 표현한다. 따라서 우리는 이 k-step Transition Probability를 보존하는 것을 목표로 한 Pre-Training Task를 설계하여, Local Connectivity와 Structural Similarity 사이의 근본적인 격차를 효과적으로 줄일 수 있었다.

Injecting Structural Features

구조적 식별력 측면에서는 기존의 Graph Isomorphism Network (GIN)이 Sum Aggregation과 Multi-layer Perceptron (MLP)을 이용해 1-dimensional Weisfeiler-Lehman (1d-WL) 테스트 수준의 구조 식별력을 제공하는 것으로 알려져 있다. 그러나 분자 구조 분석과 같은 보다 정교한 구조 식별력이 요구되는 도메인에서는 더 높은 수준의 식별력이 필요하다. 물론 Subgraph Sampling 기반의 GNN과 GT 모델들이 2d-WL 수준의 식별력을 달성할 수 있지만, 이러한 모델들은 k-hop Subgraph Sampling과 Self-Attention의 결합 과정에서 일부 구조 정보를 잃어버릴 수 있는 단점이 존재한다. 기존 연구에서는 이를 해결하기 위해 Laplacian Eigenvector나 Multi-hop Ordered Degree Sequence와 같은 추가적인 Structural Feature를 Initial Node Representation이나 Attention Score에 주입하는 방법을 제안한 바 있다. UGT는 이러한 접근을 한층 발전시켜, Laplacian Eigenvector와 Multi-hop Ordered Degree Sequence, 그리고 앞서 언급한 k-step Transition Probability 정보를 Initial Node Representation과 Attention Score, 그리고 Initial Residual에 결합하여 구조 식별력을 한층 더 향상시켰다.

Achieving Expressive Power as 3d-WL

이러한 설계를 통해 UGT는 기존 Structural Feature Injection 기법 및 Structural Similarity를 고려한 Graph Augmentation 기법과 k-step Transition Probability 보존 Task 간에 뛰어난 시너지를 창출하였다. 실제로 Node Classification, Node Clustering, Graph Classification 등 다양한 Downstream Task에서 State-of-the-Art (SOTA) 성능을 기록하였으며, 특히 Heterophilic 그래프와 Homophilic 그래프 모두에서 뛰어난 성능을 보였다는 점에서 우리의 설계 목표를 충분히 달성했다고 판단한다. 또한 Graph Isomorphism Testing에서 3d-WL에 준하는 수준의 구조 식별력을 달성했다는 점은 UGT가 기존 모델 대비 구조 식별력 측면에서 분명한 진보를 이루었음을 입증한 결과였다.