[논문 리뷰] DOTIN: Dropping Task-Irrelevant Nodes for GNNs

Shaofeng Zhang, Feng Zhu외 3명|arXiv (Cornell University)|2022. 04. 28.

Advanced Graph Neural Networks인용 수 4

한 줄 요약

DOTIN은 K개의 학습 가능한 가상 노드를 활용하여 작업에 특화된 전역 표현을 캡처함으로써, GNN에서 작업과 관련이 없는 노드까지 최대 90%까지 적응적으로 제거하는 새로운 작업 인식 그래프 프루닝 방법을 제안한다. 트랜스포머 기반의 어텐션 메커니즘을 사용하여 노드 중요도를 순위 매김함으로써, 성능 저하 없이 D&D에서 약 50% 빠른 훈련 속도와 약 60% 적은 메모리 사용량을 달성하며, 단일 및 다중 작업 설정에서 풀링 기반 기준선을 능가한다.

ABSTRACT

Scalability is an important consideration for deep graph neural networks. Inspired by the conventional pooling layers in CNNs, many recent graph learning approaches have introduced the pooling strategy to reduce the size of graphs for learning, such that the scalability and efficiency can be improved. However, these pooling-based methods are mainly tailored to a single graph-level task and pay more attention to local information, limiting their performance in multi-task settings which often require task-specific global information. In this paper, departure from these pooling-based efforts, we design a new approach called DOTIN (\underline{D}r\underline{o}pping \underline{T}ask-\underline{I}rrelevant \underline{N}odes) to reduce the size of graphs. Specifically, by introducing $K$ learnable virtual nodes to represent the graph embeddings targeted to $K$ different graph-level tasks, respectively, up to 90\% raw nodes with low attentiveness with an attention model -- a transformer in this paper, can be adaptively dropped without notable performance decreasing. Achieving almost the same accuracy, our method speeds up GAT by about 50\% on graph-level tasks including graph classification and graph edit distance (GED) with about 60\% less memory, on D\&D dataset. Code will be made publicly available in https://github.com/Sherrylone/DOTIN.

연구 동기 및 목표

• 그래프 수준 작업에서 GNN의 확장성 한계를 줄이기 위해 계산 및 메모리 비용을 감소시키는 것.
• 기존 풀링 방법의 단점을 해결하기 위해 국소 정보에 중점을 두고 작업에 무관한 방법을 개선하는 것.
• 가상 노드를 통해 전역 표현을 학습함으로써 적응적이고 작업에 특화된 노드 프루닝을 가능하게 하는 것.
• 성능 저하 없이 다중 작업 GNN 설정에서 효율성과 확장성을 향상시키는 것.
• GNN에 종속되지 않고 파라미터 효율적인 방법으로 계층적 그래프 크기 감소를 제공하는 것.

제안 방법

• K개의 학습 가능한 가상 노드를 도입하여, K개의 다른 그래프 수준 작업에 대해 각각 작업에 특화된 전역 임베딩을 표현한다.
• 모든 작업을 통해 노드의 주의도 점수를 계산하기 위해 트랜스포머 기반 어텐션 메커니즘을 사용한다.
• 모든 작업에 걸쳐 평균 주의도가 낮은 노드를 식별하고, 이는 작업과 관련이 없는 노드로 간주한다.
• 각 GNN 레이어 이후에 계층적으로 노드를 제거하여, 추가적인 GNN 또는 클러스터링 레이어 없이도 그래프 크기를 감소시킨다.
• 특히 다중 작업 환경에서 높은 작업 특화 주의도를 가지는 노드를 유지함으로써 모델 성능을 유지를 한다.
• 비용이 많이 드는 클러스터링 또는 추가 GNN 레이어를 피하기 위해 가상 노드의 주의도 점수를 정렬함으로써 O(N log N) 복잡도를 달성한다.

실험 결과

연구 질문

• RQ1 그래프 수준 작업에서 성능 저하 없이 작업과 관련이 없는 노드를 식별하고 제거할 수 있는가?
• RQ2 분류 및 그래프 편집 거리와 같은 다양한 그래프 수준 작업 간에 작업에 특화된 노드 중요도는 어떻게 달라지는가?
• RQ3 GNN에 종속되지 않는 방법이 풀링 기반 접근법보다 더 높은 효율성과 확장성을 달성할 수 있는가?
• RQ4 다중 작업 어휘 기반의 적응적 노드 제거가 성능 향상과 메모리 사용량 감소에 기여하는가?
• RQ5 이 방법은 새로운 작업과 데이터셋에 일반화되어 효율성과 정확도를 유지할 수 있는가?

주요 결과

• D&D 데이터셋에서 DOTIN은 기준선 대비 성능 저하 없이 노드 수를 최대 90%까지 감소시킨다.
• D&D 데이터셋에서 DOTIN은 GAT 백본을 사용할 경우 약 50% 빠른 훈련 속도와 약 60% 적은 메모리 사용량을 달성한다.
• 은닉 차원이 512와 4096일 경우, DOTIN은 90% 노드 제거 비율에서 각각 기준선 대비 42%와 57%의 메모리만 사용한다.
• 시각화 결과는 노드 중요도가 작업 간에 크게 다름을 보여주며, 예를 들어 그래프 분류에 중요한 노드는 GED에서는 종종 관련이 없고, 그 반대의 경우도 마찬가지다.
• gPool과 같은 풀링 기반 방법은 작업에 무관한 설계로 인해 작업과 관련 있는 노드를 제거할 수 있어, DOTIN이 그에 비해 뛰어난 성능을 보인다.
• 이 방법은 D&D, PROTEINS, NCI1, NCI109, FRANKENSEIN 등 여러 데이터셋에서 잘 일반화되며, 다양한 백본과 제거 비율에서도 높은 정확도를 유지한다.