BeamWitness: Rooted Subgraph Beam Search for Selective Graph Representation Learning

Standard graph neural networks (GNNs) are built on message passing and neighborhood aggregation, following the Weisfeiler--Lehman refinement paradigm. Despite their wide adoption, these models tend to compress and mix neighbor information from the very first layer. This early mixing obscures the critical subgraphs that drive the prediction. Attention-based and subgraph-based methods partially address this by reweighting nodes or enriching local representations, but both families remain passive as they usually rely on soft weighting or predefined subgraph templates. We propose BeamWitness, a selective graph representation learning framework that treats subgraph identification as an active search problem. Starting from a root node, a policy trained with task rewards expands a connected subgraph one node at a time or terminates early, and beam search retains a small set of candidate subgraphs as witnesses for prediction. For node-level tasks, the target node serves as the root, while for graph-level tasks, BeamWitness launches searches from multiple root nodes and aggregates the selected witnesses into a graph-level representation. On controlled motif benchmarks with known ground-truth subgraphs, BeamWitness achieves near-perfect accuracy and selects witnesses that align with the true motifs. On real-world node and graph classification benchmarks, BeamWitness is competitive with standard GNNs and related methods while providing an explicit, interpretable subgraph-based account of each prediction.