TGSBM: Transformer-Guided Stochastic Block Model for Link Prediction

TL;DR

TGSBM addresses the challenge of scalable, interpretable link prediction in large-scale networks with overlapping communities. It couples the principled generative structure of Overlapping Stochastic Block Models with the expressive power of sparse Graph Transformers, via expander-augmented sparse attention, a neural variational encoder for structured OSBM posteriors, and a neural edge decoder that preserves interpretability. The approach delivers near-linear time complexity, competitive accuracy, and interpretable latent communities, achieving up to 6x faster training and strong performance under the HeaRT protocol. This work advances practical deployment of graph learning by balancing predictive power with transparency and efficiency, enabling robust link prediction on web-scale graphs.

Abstract

Link prediction is a cornerstone of the Web ecosystem, powering applications from recommendation and search to knowledge graph completion and collaboration forecasting. However, large-scale 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 \textbf{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) \emph{expander-augmented sparse attention} that enables near-linear complexity and efficient global mixing, (ii) a \emph{neural variational encoder} that infers structured posteriors over community memberships and strengths, and (iii) a \emph{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\times$ faster training), and interpretable community structures. These results position TGSBM as a practical approach that strikes a balance between accuracy, efficiency, and transparency for large-scale link prediction.

Figures (4)

• Figure 1: Link prediction with overlapping communities. Researchers belong to multiple overlapping communities (Deep Learning, Computer Vision, Bayesian Methods). The task is to predict missing links (dashed edge with "?") based on shared community memberships while preserving interpretability.
• Figure 2: The overview of TGSBM architecture. TGSBM consists of three main components: (a) Sparse Graph Transformer Encoder processes node representations over local edges augmented with expander connections. (b) Variational Overlapping Stochastic Block Model uses three specialized heads to infer stick-breaking variables, binary community memberships, and Gaussian membership strengths. (c) Neural Edge Decoder reconstructs links via the OSBM generative process, maintaining interpretability through probabilistic community interactions.
• Figure 3: Training time per 100 epochs on Cora, Citeseer, and Pubmed. TGSBM (orange) vs. LPFormer (blue). Bars represent median wall-clock time across 3 runs.
• Figure 4: Qualitative analysis on synthetic data with 100 nodes and 10 communities. (a) Ground-truth adjacency matrix sorted by community memberships, with white/black/gray denoting link/no-link/hidden edges. (b) Learned latent structure $\mathbf{Z} \in \mathbb{R}^{100 \times 10}$ from TGSBM, where color intensity indicates community membership strength. (c) Reconstructed adjacency probabilities using only the first 2 dimensions of $\mathbf{Z}$ (columns 1-2), demonstrating that dominant communities capture essential network structure.