An Empirical Study of Retrieval-enhanced Graph Neural Networks

TL;DR

This paper introduces GRAPHRETRIEVAL, a retrieval-enhanced scheme for GNNs that improves graph property prediction by retrieving similar graphs, demonstrating significant benefits across multiple datasets and addressing label distribution issues.

Contribution

The paper proposes a retrieval-augmented GNN framework called GRAPHRETRIEVAL that is model-agnostic and enhances performance on graph tasks through example retrieval.

Findings

GRAPHRETRIEVAL outperforms existing GNNs on 13 datasets.

Retrieval enhancement alleviates long-tailed label distribution problems.

The scheme is effective across various GNN architectures.

Abstract

Graph Neural Networks (GNNs) are effective tools for graph representation learning. Most GNNs rely on a recursive neighborhood aggregation scheme, named message passing, thereby their theoretical expressive power is limited to the first-order Weisfeiler-Lehman test (1-WL). An effective approach to this challenge is to explicitly retrieve some annotated examples used to enhance GNN models. While retrieval-enhanced models have been proved to be effective in many language and vision domains, it remains an open question how effective retrieval-enhanced GNNs are when applied to graph datasets. Motivated by this, we want to explore how the retrieval idea can help augment the useful information learned in the graph neural networks, and we design a retrieval-enhanced scheme called GRAPHRETRIEVAL, which is agnostic to the choice of graph neural network models. In GRAPHRETRIEVAL, for each input graph, similar graphs together with their ground-true labels are retrieved from an existing database. Thus they can act as a potential enhancement to complete various graph property predictive tasks. We conduct comprehensive experiments over 13 datasets, and we observe that GRAPHRETRIEVAL is able to reach substantial improvements over existing GNNs. Moreover, our empirical study also illustrates that retrieval enhancement is a promising remedy for alleviating the long-tailed label distribution problem.