Link Prediction(LP) is a fundamental problem in graph machine learning that aims to predict the existence of links between nodes. Most current research on LP adopts Graph Neural Networks (GNNs) to learn the representation of subgraphs, but fails to efficiently capture global topological information in large graphs. In response to this issue, we focus on the global attention mechanism of Transformers. Nevertheless, original Transformers are not inherently suitable for learning graph-structured data, and their deep multi-head attention architecture has been limited by prohibitive compute and memory costs. In this work, we propose a lightweight model: a single-layer, single-head Transformer, which constructs subgraph structural features based on hash estimation. It provides a new perspective for applying Transformers in graph-structured data processing. Firstly, we utilize MinHash and HyperLogLog techniques to estimate the structural information of subgraphs, then fuse subgraph structural features with node features to achieve efficient message passing. In this case, our model does not need to extract or manipulate enclosed subgraphs, and structural features can be preprocessed. Additionally, we design a single-layer, single-head Transformer as the encoder for graph-structured data, utilizing the attention mechanism to capture global effects. Meanwhile, our model does not require extra positional encoding, which significantly reduce computational complexity. Extensive experiments demonstrate that our model achieves optimal prediction accuracy on large-scale datasets with high efficiency. The code is available at https://github.com/sunteng6/HeFormer .

错误:搜索内容不能为空,请输入英文关键词
错误:关键词超出字数限制,请精简
高级检索

HeFormer: A Lightweight Transformer Combining Hash Estimation for Link Prediction

  • Teng Sun,
  • Xiaoqiang Xiao,
  • Xu Zhang,
  • Weixun Ning

摘要

Link Prediction(LP) is a fundamental problem in graph machine learning that aims to predict the existence of links between nodes. Most current research on LP adopts Graph Neural Networks (GNNs) to learn the representation of subgraphs, but fails to efficiently capture global topological information in large graphs. In response to this issue, we focus on the global attention mechanism of Transformers. Nevertheless, original Transformers are not inherently suitable for learning graph-structured data, and their deep multi-head attention architecture has been limited by prohibitive compute and memory costs. In this work, we propose a lightweight model: a single-layer, single-head Transformer, which constructs subgraph structural features based on hash estimation. It provides a new perspective for applying Transformers in graph-structured data processing. Firstly, we utilize MinHash and HyperLogLog techniques to estimate the structural information of subgraphs, then fuse subgraph structural features with node features to achieve efficient message passing. In this case, our model does not need to extract or manipulate enclosed subgraphs, and structural features can be preprocessed. Additionally, we design a single-layer, single-head Transformer as the encoder for graph-structured data, utilizing the attention mechanism to capture global effects. Meanwhile, our model does not require extra positional encoding, which significantly reduce computational complexity. Extensive experiments demonstrate that our model achieves optimal prediction accuracy on large-scale datasets with high efficiency. The code is available at https://github.com/sunteng6/HeFormer .