[GNN] Do Transformers Really Perform Bad for Graph Representation?, Graphormer

Do Transformers Really Perform Bad for Graph Representation?

Link : https://arxiv.org/pdf/2106.05234.pdf

 

 

Abstract

 

트랜스포머 구조는 많은 도메인에서 지배적인 초이스가 되고 있지만, graph-level prediction task 리더보드에서는 아직 Popular하지 않으며, GNN variants 모델이 주로 사용되고 있다. 그래서 해당 논문에서는 graph representation learning에 있어 트랜스포머가 잘 적용될 수 있는지 미스테리를 Graphormer라는 모델을 통해 해소해보고자 한다. 키 아이디어는 모델에 그래프의 structural information을 효과적으로 인코딩하기 위해 트랜스포머 구조를 이용한다. 게다가 Graphormer의 expressive power를 수학적으로 특징짓고, structural information of graph를 인코딩 하는 방법과, 많은 popular GNN variants들은 Graphormer의 special case로 커버할 수 있음을 보인다.

 

 

Introduction

 

트랜스포머는 graph representation leaderboards에서 de-facto standard이지는 않으며 key module(e.g. feature aggregation)을 softmax attention로 대체한 GNN variants 정도만 쓰였다. 트랜스포머를 Graph representation에 잘 적용하기 위해서는 self-attention을 각 노드와 다른 노드의 semantic similarity만을 계산하는데 그치면 안되고 node pair들 사이의 관계를 반영하기 위한 graph structural information을 고려해야 한다고 저자들은 말하고 있다. 

 

첫째로, 해당 논문에서는 "Centrality Encoding"을 제안한다. Centrality Encoding은 그래프에서 노드의 중요도를 나타내며 이는 Self-attention module로 반영할 수 없는 정보이다. 특히 Centrality Encoding을 위해서 'Degree centrality'를 이용하는데, 이는 각 노드들의 degree를 따르는 learable vector로 assign되며 이 learnable vector는 input layer에서 node의 feature로 더해진다.

 

두번째로, "Spatial Encoding"을 제안한다. 이는 노드간 structural relation을 나타낸다. 다른 structured data로 부터 graph-structural data를 구별하는 notable geometrical 속성은 graph를 embed하기 위한 표준 grid에는 존재하지 않는다. node의 structure는 no-Euclidean space이며 edge로 연결되어있다. 이러한 structural information을 모델링하기 위해서, 그들의 spatial relation을 base로 한 learnable embedding을 assign한다. 이를 위해 두 노드의 shortest path를 사용하고, 이는 softmax attention의 bias term으로 encoding되며 그래프의 spatial dependency를 정확하게 모델이 알 수 있도록 도움을 준다. 게다가 때로는 분자 그래프의 두개의 원자 사이에 edge feature가 포함된 추가적인 spatial information이 존재하기도 한다.

 

정확하게 말하면, 각 노드 페어의 shortest path의 learnable emgedding과 edge feature의 dot-products의 평균을 계산하고, 이를 attention module에서 사용한다. 이러한 인코딩을 Graphormer가 갖추게 되면, 노드 페어의 관계를 더 잘 모델링할 수 있게 되고 graph를 표현할 수 있게 된다.

 

 

Preliminary

Graph Neural Network(GNN).

그래프 G를 Node set V={v_1, v_2, .. v_n}, Edge E로 표현하고 G=(V, E), node v_i의 feature vector는 x_i가 된다. GNN은 graph와 node의 presentation을 learning하는 것이 목적이다. 보통 modern GNN에서는 higher-order neighbor이나 첫번째 node aggregating representation 을 반복적으로 업데이트하는 learning 스키마를 따른다. node v_i의 representation을 h_i^(l)로 표기하고 l은 몇 번째 레이어인지 나타낸다. 그리고 h_i^(0)을 x_i로 정의한다. l 번째 iteration에서의 aggregation은 아래와 같이 AGGREGATE_COMBINE step으로 표현할 수 있다.

N(v_i)는 노드 v_i의 higher-order neighbors이거나 첫번째의 set이다. AGGREGATION 함수는 노드의 이웃들의 정보를 모으는데 사용되고, 일반적인 방법으로는 MEAN, MAX, SUM이 있다. COMBINE 함수는 각 이웃들의 정보를 node representation에 융합하기 위한 함수이다. 게다가 graph representation task를 위해 READOUT 함수는 전체 그래프의 representation h_G를 나타내기위해 최종 노드 feature h_i^(L)을 aggregate한다. 

READOUT은 summation이나 좀 더 정교한 graph-level pooling function같은 simple permuation invariant function에 의해 구현된다.

 

Transformer.

Transformer 아키텍쳐는 Transformer layer들로 이루어져 있고 각 레이어는 self-attention module과 position-wise feed-forward network(FNN) 이렇게 2 파트가 있다. position i의 hidden representation h_i은 self attention module의 input이고 H는 h_i set이며 input H는 W_Q, W_K, W_V 3개의 매트릭스에 의해 prjection된다. self-attention은 아래와 같이 계산된다.

A는 Query와 Key들 사이의 similarity를 capture한 매트릭스가 된다.

 

 

Graphormer

Graphormer는 graph representation을 배우기 위한 뉴럴넷의 inductive bias를 위해 디자인 되었다. 

 

 

Structural Encodings in Graphormer

Introduction에서 언급했듯이, 그래프의 structural information을 Transformer model에 활용하는 것이 무엇보다도 중요하며 이를 위해 간단하지만 효과적인 Encoding design에 대해 소개한다.

 

1. Centrality Encoding

QK를 내적한 A matrix를 Softmax에 넣어 나온 Attention Distribution Attn(H)는 노드 사이의 semantic correlation에 근거하여 계산된 값이다. 하지만 node centrality는 graph understanding에 있어 강한 signal이며 이를 측정하는 것은 중요하다. 예를 들어 social 망에 대한 그래프가 있다고 할 때 팔로워가 많은 celeb은 network의 트렌드를 predict하는데 있어 무엇보다 중요한 요소일 것이다. centrality는 현재 attention 만으로는 계산되지 않는다.

Graphormer에서 centrality를 위해서 standard centrality measurement 중 하나인 'degree centrality'를 network의 additional signal로 이용한다. 자세하게는, 각 노드의 indegree, outdegree를 따르는 두개의 real-valued embedding vector인 'centrality encoding'을 만들고 이를 간단하게 node feature에 더해서 최종 input을 만들게 된다.

z^-와 z^+는 learnable embedding vector이며 z^-는 v_i의 indegree, z^+는 v_i의 outdegree에 해당하고, undirected graph의 경우 indegree와 outdegree는 하나로 통합된다. 이 centrality encoding을 적용하면 softmax attention은 쿼리와 키의 node importance signal을 알 수 있게 된다. 그러므로 모델은 sematic corrleation과 node importance를 어텐션 메커니즘안에서 capture할 수 있게 되는 것이다.

 

2. Spatial Encoding

Transformer의 강점은 global receptive field를 잘 포착할 수 있다는 점이다. 각 트랜스포머 레이어에서 각 토큰은 어떤 포지션의 정보든 attend할 수 있고 이를 representation할 수 있기 때문이다. 그러나 이 operation은 모델이 explicitly 다른 포지션을 specify하거나 positional dependency(such as locality)를 encode하는 byproduct problem이 있다. sequence data에서 Transformer layer는 각 포지션의 embedding(i.e. absolute positional encoding)을 input으로써 줄 수 있고 또는 어떤 두 포지션의 상대적인 거리를 encoding(i.e. relative positional encoding)할 수도 있다.

그러나 그래프를 위해서는 노드가 sequence로 arrange되면 안된다. 노드들은 multi-dimensional spatial space를 가지고 있고 엣지로 연결되어 있다. 그래프의 structural information을 인코딩하기 위해서 'Spatial Encoding'을 제안한다. 어떤 function phi가 있다고 할 때, phi는 \phi(v_i, v_j): V x V -> R 가 되는 즉 v_i와 v_j의 spatial relation을 measure하는 함수이며 두 노드 사이의 conectivity를 정의한다. 여기서는 phi를 두 노드간 shortest path(SPD)를 사용했고 두 노드가 연결되어 있지 않으면 -1을 할당했다. phi로 나오는 두 노드의 conectivity는 phi(v_i, v_j)에 의해 인덱싱된 learnable scalar이며, self-attention module의 bias텀으로 들어간다. 아까 위 수식에 있는 A matrix는 아래와 같이 수정된다.

b_\phi(v_i, v_j)는 spatial Encoding이고 위처럼 A matrix의 바이어스 텀으로 더해지는 것을 볼 수 있고 이는 모든 레이어에 거쳐 공유된다. 이 learnable b_phi가 phi(v_i, v_j)에 대해 decreasing function으로 learning된다면, 모델은 가까이 있는 노드들에 더 attend하게 되고 멀리 있는 노드들을 덜 attend하게 된다.

 

3. Edge Encoding

분자 구조 그래프 데이터에서 두 원자 노드 쌍에 대한 결합 종류를 edge로 나타내면 이 edge에 대한 feature를 함께 encoding하는 것은 필수적이다. 그러므로 이전 연구들에서는 주로 2가지 edge encoding 방법이 쓰였는데,  첫번째는 연관 노드(associated node) 피처에 더하는 것이고 다른 하나는 연관 엣지(associated edge) 피처가 노드 피처 aggregation에 함께 사용되는 것이다. 연관 노드에 엣지 피처를 오직 엣지 정보에 propagate하는 이러한 방법은 전체 그래프의 representation에 edge information을 활용하기에 효과적인 방법은 아니다.

그래서 edge feature를 attention layer에 잘 적용하기 위해 새로운 edge encoding 방법을 제안한다. 어텐션 매커니즘은 각 노드 페어의 correlation을 estimate하도록 요구된다. 그리고, 우리는 연결된 edge는 correlation이 고려 됬을 것이라고 믿는다. 각 노드페어 (v_i, v_j)에서 엣지에 해당하는 shortest path SP_ij를 찾고, 엣지 피처의 dot-product의 평균값을 계산하고 이를 위와 마찬가지로 attention module의 bias 텀으로 준다. 이는 아래 식에서 c_ij에 해당한다.

 

x_{e_n}은 SP_ij에서 n번째 엣지 e_n의 피처이고, w_n^E는 n번째 weight embedding이고 d_E는 엣지 피처의 dimension이다.

 

 

Implementation Details of Graphormer

Transformer는 classic transformer를 그대로 사용했으나 효과적인 optimzation을 준다는 이유로 Layer Normalization을 FFN, MHA 이전에 적용한 것만 다르다.

 

그리고 모든 노드와 연결된 Virtual Node를 하나 주었고, Vnode라고 한다. 이는 BERT에서 [CLS] 토큰과 유사하며 sequence-level feature를 표현하게 된다. READOUT을 통해 graph feature를 나타내면 정보가 손실되는 영향도 있는 듯 하다. AGGREGATION-COMBINE step에서 normal node들과 똑같이 update되지만 spatial encoding에서는 pysical node와 virtual node의 연결을 구별하기 위해, Vnode의 b_phi는 reset하였다.

 

 

How Powerful is Graphormer?

Fact 1.

적절한 weight와 distance function phi를 선택함으로써, Graphormer laeyr는 GIN, GCN GraphSAGE와 같은 popular GNn 모델의 AGGREGATE, COMBINE step을 나타낼 수 있다.

이를 증명하기 위한 결과는,

 

1) Spatial Encoding은 node v_i의 이웃 set N(v_i)을 구별하기 위한 self-attention module을 가능하게 한다.

2) node의 degree를 알면, 이웃에 대한 평균을 이웃에 대한 합으로 변환할 수 있다.

3) MHA와 FFN에서 v_i와 N(v_i)의 reperesentation은 따로 processed 될 수 있고, 나중에 합쳐진다. 

 

와 같으며 자세한 proof는 Appendix A에 있고 이 포스팅에서는 생략했다.

더 나아가, Spatial Encoding을 사용하면 Graphormer는 표현력이 1-Weisfeiler-Lehman(WL) test에 불과한 기존 classic GNN의 message passing을 뛰어 넘을 수 있다는 것과 자세한 예시, 그리고 Graphormer는 1-WL test에 실패한 그래프를 구별하는데 어떻게 도움이 되는지에 대한 예시도 Appendix A에 있다.

 

Connection between Self-attention and Virtual Node

Graphormer가 popular GNN에 비해 뛰어난 표현력을 가지는 것 말고도 virtual node와 self-attention을 사용하는 것 사이에 흥미로운 connection이 있다는 것을 저자들은 발견했다. virtual node trick은 일반 GNN을 뛰어넘는 성능을 보여주고 개념적으로는, 모든 그래프의 정보를 GNN의 READOUT처럼, aggregate하고 그다음 각 노드에 propogate할 수 있다. 이 supernode를 전체 그래프에 대해 사용하면, 정보를 전파하는데 잇어 over-smoothing할 수 있다. 대신, 이러한 graph-level aggregation, propogate operation은 추가 인코딩 없이 Vanilla self-attention을 통해 자연스럽게 수행될 수 있다.

Fact 2.

적절한 weight를 선택함으로써 추가적인 인코딩 없는 그래포머의 결과 node reperesentation MEAN READOUT function으로 볼 수 있다.

이러한 사실은 각 노드가 모든 다른 노드를 attend하는 self attention의 이점을 취한 것이라고 볼 수 있다. 또한 이는 graph-level READOUT를 simuate하는 것이고 게다가 이 방법은 over-smoothing이 생기지 않는 것을 empirically 발견했다.

 

 

Experiments

OGB Large-Scale challenge

quantum chemistry regression challenge인 OGB-LSC(i.e. PCQM4M-LSC) task 결과이며 Graphormer 모델은 12개 레이어, 768 차원의 Transformer로 Graphormer안의 모든 Transformer는 같은 dimension을 가진다. Graphormer small은 6개 레이어, 512 차원이다. Experment setting에 대한 자세한 내용은 Appendix B에 있다.

이전 GNN variants model인 GCN, GIN, GINE, DeeperGNC, GT 등을 Baseline으로 삼았고, 뒤에 -VN 붙은 것은 Virtual Node를 적용한 결과이다. 대부분의 기존 SOTA model의 성능을 능가했고, Virtual Node를 사용하면 over-smoothing 현상이 적다는 것도 확인했다.(train, validate error가 지속적으로 감소하는 것을 통해서)

 

Graph Representation

그리고 graph-level prediction task OGB(OGBG-MolPCBA, OGBG-MolHIV)와 GNN benchmarking(GIN)에 대한 실험 결과이다. pre-training은 OGB에 의해 이루어졌고, OGB-LSC에 transferable capability가 좋은지 확인함을 통해 graph-representation이 잘 되었는지 확인한다.  

실험 결과를 보면 다른 GNN 모델들에 비해 outperform하였고, 가장 최근 제안된 Transformer-based model인 GT, SAN에 비해서도 좋은 성능을 보여줌을 확인할 수 있었다.

 

Ablation studies

Node Relation Encoding. 

먼저 트랜스포머에서 기존에 쓰이던 Positional Encoding(PE)와 Graphormer에서 제시한 Spatial Encoding을 비교하였을 때, 둘다 node간 관계를 구분짓는 정보를 인코딩 하기 위함이지만, Spatial Encoding이 더 좋은 성능을 보였다. 여기서 비교한 PE는 Laplacian PE이고, Weisfeiler-Lehman-PE(LW-PE)보다 이전 연구들에서 해당 PE가 더 좋았기 때문에 Laplacian PE와 비교했다고 한다. 

 

Centrality Encoding.

degree-based centrality encoding이 있을 때와 없을 때는 아주 큰 격차로 높은 성능을 보였으며 이는 Graph data에서 Centrality는 매우 중요한 정보임을 알 수 있다.

 

Edge Encoding.

3절에서 얘기한 2가지의 edge encoding과 해당 논문에서 제안한 attention bias로 edge encoding을 넣는 것을 비교해보았을 때, 기존 2 edge encoding 방식의 성능은 거의 차이가 없었지만 해당 논문에서 제시한 edge encoding 방식은 더 뛰어난 성능을 보였다. 이는 이 방법이 즉, edge encoding을 attention bias 텀으로 넣는 것이 더 효과적이고 엣지의 spatial information이 모델이 잘 인식한다는 것을 의미한다.

 

 

Conclusion

Graphormer는 기존에 Graph Transformer와 같이 Transformer를 Graph data에도 잘 적용해보려고 했던 다른 모델과 기존의 GNN variants 보다도 뛰어나고 효과적인 성능을 보였다. 해당 성능을 보인데 주요 키 attribution은, 3개의 noevel graph structural encoding(Centrality, Spatial, Edge Encoding) 방식에 있다고 볼 수 있다. 하지만 Self-attention의 quadratic complexity는 아직 아주 큰 스케일의 graph에는 적용하기엔 무리가 있다. 따라서 저자들은 효율적인 graphormer를 develop하는 것과, 특정 데이터에서 도메인 지식 encoding을 활용하는 것, node representation extraction에 적용할 수 있는 sampling strategy를 모색하는 것은 future work로 남겨두고 있다.