Utilizing a transformer for link prediction

Abstract

Transformers have been widely employed and demonstrated success in many domains such as language and vision. While recent approaches have extended and applied them to the graph domain, e.g., graph property prediction, their success has been limited to predicting the properties of small-scale graphs. However, in practice, it is crucial to efficiently predict graph properties (e.g., node and edge) for large graph-structured data such as those from the Internet and biology. To this end, we propose a novel Link Prediction Transformer, called LiT, that exploits specialized tokens from the extracted subgraphs as embedding inputs to the Transformer for link prediction in a large-scale graph. With our proposed type, node, and position identifiers that effectively encode the graph properties, we verify that local subgraphs embed sufficient information for predicting links in a graph, and demonstrate at least comparable or superior performance to those of Graph Neural Networks (GNN)-based state-of-the-art models. We further identify important nodes and edges from the attention score matrices and examine the effect of each identifier and each encoding in the position identifier in subgraph token encoder.