TGSBM: Transformer-Guided Stochastic Block Model for Link Prediction
Zhejian Yang, Songwei Zhao, Zilin Zhao, Hechang Chen
Abstract
Link prediction is a cornerstone of the Web ecosystem, powering applications from recommendation and search to knowledge graph completion and collaboration forecasting. However, largescale networks present unique challenges: they contain hundreds of thousands of nodes and edges with heterogeneous and overlapping community structures that evolve over time. Existing approaches face notable limitations: traditional graph neural networks struggle to capture global structural dependencies, while recent graph transformers achieve strong performance but incur quadratic complexity and lack interpretable latent structure. We propose TGSBM (Transformer-Guided Stochastic Block Model), a framework that integrates the principled generative structure of Overlapping Stochastic Block Models with the representational power of sparse Graph Transformers. TGSBM comprises three main components: (i) expander-augmented sparse attention that enables near-linear complexity and efficient global mixing, (ii) a neural variational encoder that infers structured posteriors over community memberships and strengths, and (iii) a neural edge decoder that reconstructs links via OSBM's generative process, preserving interpretability. Experiments across diverse benchmarks demonstrate competitive performance (mean rank 1.6 under HeaRT protocol), superior scalability (up to 6× faster training), and interpretable community structures. These results position TGSBM as a practical approach that strikes a balance between accuracy, efficiency, and transparency for largescale link prediction. CCS Concepts • Computing methodologies → Latent variable models; • Information systems → Data mining.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Builds on17
- Open Graph Benchmark: Datasets for Machine Learning on GraphsWeihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong et al.NeurIPS 2020 · 3,935 citations
- Do Transformers Really Perform Badly for Graph Representation?Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng et al.NeurIPS 2021 · 1,632 citations
- Rethinking Graph Transformers with Spectral AttentionDevin Kreuzer, Dominique Beaini, William L. Hamilton, Vincent Létourneau et al.NeurIPS 2021 · 854 citations
- Neural Bellman-Ford Networks: A General Graph Neural Network Framework for Link PredictionZhaocheng Zhu, Zuobai Zhang, Louis-Pascal A. C. Xhonneux, Jian TangNeurIPS 2021 · 546 citations
- NodeFormer: A Scalable Graph Structure Learning Transformer for Node ClassificationQitian Wu, Wentao Zhao, Zenan Li, David P. Wipf et al.NeurIPS 2022 · 472 citations
Related papers
- Neo-GNNs: Neighborhood Overlap-aware Graph Neural Networks for Link PredictionSeongjun Yun, Seoyoon Kim, Junhyun Lee, Jaewoo Kang et al.NeurIPS 2021 · 183 citations
- Plain Transformers are Surprisingly Powerful Link PredictorsQuang Truong, Yu Song, Donald Loveland, Mingxuan Ju et al.ICML 2026
- Exphormer: Sparse Transformers for GraphsHamed Shirzad, Ameya Velingker, Balaji Venkatachalam, Danica J. Sutherland et al.ICML 2023 · 219 citations
- Simplifying and Empowering Transformers for Large-Graph RepresentationsQitian Wu, Wentao Zhao, Chenxiao Yang, Hengrui Zhang et al.NeurIPS 2023 · 318 citations
- Efficient Learning on Large Graphs using a Densifying Regularity LemmaJonathan Kouchly, Ben Finkelshtein, Michael M. Bronstein, Ron LevieICLR 2026 · 2 citations
