Learning to Execute Graph Algorithms Exactly with Graph Neural Networks

TL;DR

This work establishes exact learnability guarantees for executing graph algorithms with graph neural networks under bounded-degree and finite-precision constraints, using Neural Tangent Kernel theory to connect local binary instructions to global graph execution. It introduces a graph template matching framework in which an ensemble of MLPs learns local update rules, which are then embedded into a GNN update to perform entire LOCAL-model algorithms exactly with high probability. The main contributions include formal learnability results for Message Flooding, BFS, DFS, and Bellman-Ford, explicit resource bounds (dataset size, embedding dimension, and ensemble size), and an experimental validation of ensemble complexity and NTK-based predictions. The approach offers a principled path to exact algorithmic execution on graphs with scalable memory and communication, with potential impact on distributed computing benchmarks and theory-guided GNN design.

Abstract

Understanding what graph neural networks can learn, especially their ability to learn to execute algorithms, remains a central theoretical challenge. In this work, we prove exact learnability results for graph algorithms under bounded-degree and finite-precision constraints. Our approach follows a two-step process. First, we train an ensemble of multi-layer perceptrons (MLPs) to execute the local instructions of a single node. Second, during inference, we use the trained MLP ensemble as the update function within a graph neural network (GNN). Leveraging Neural Tangent Kernel (NTK) theory, we show that local instructions can be learned from a small training set, enabling the complete graph algorithm to be executed during inference without error and with high probability. To illustrate the learning power of our setting, we establish a rigorous learnability result for the LOCAL model of distributed computation. We further demonstrate positive learnability results for widely studied algorithms such as message flooding, breadth-first and depth-first search, and Bellman-Ford.

Figures (9)

• Figure 1: Outline of our approach. During training (top), we train $K$ MLP instances $\Phi$ on binary, block-structured instructions that teach local operations. The bits are split into a computation section (yellow) and a message/communication section (green). Inputs are encoded by a fixed $\Psi_\text{Enc}$, processed by $\Phi$, and trained by minimizing MSE to ground-truth instruction outputs. We then form an ensemble $\hat{\mu}$ by averaging MLP predictions on each data point. At inference time (bottom), we apply $\hat{\mu}$ to an attributed graph. In Step 1, we compute local node outputs from features using $\Psi_\text{Enc}$ and $\hat{\mu}$, then binarize with the step function $\Psi_H$. In Step 2, communication is carried out by message passing using masking matrices $P_C$ and $P_M$, which keep computation blocks local and transmit only message blocks. This procedure yields a GNN that can learn and execute multiple graph algorithms, including Message Flooding, Breadth-First Search, Depth-First Search, Bellman–Ford, and more generally, algorithms representable by the LOCAL model of distributed computation.
• Figure 2: Illustration of the template-matching framework on a two-node line graph: the left (blue) node and the right (purple) node. The task computes the logical AND of the bits stored at both nodes in the Bit coordinate. Top-left shows the initial node states, followed by alternating steps of local template matching and communication. Each node begins with a local ID encoding its neighborhood position, followed by its bit value. In Round 1.a, local template matching uses the template set (right), which specifies each block’s input–output relation. The computation ($\boldsymbol{x}_{C}$) and message ($\boldsymbol{x}_{M}$) components (orange and green) indicate which entries are retained locally and which are shared. Each node sends its Bit value through a designated slot in the Msg. component, determined by its local ID. Here each slot is one bit, and the first Msg. coordinate corresponds to the blue node and the second to the purple node. In Round 1.b, one communication round fills these slots so each node receives its neighbor’s bit. In the next matching step (Round 2.a), the AND is applied, and the system converges to a steady state with the result stored in the AND coordinate.
• Figure 3: Numerical and theoretical lower bounds on the ensemble size required for Message Flooding. Left: ensemble size vs. message bits $l$ with $D=2$. Right: ensemble size vs. $D$ with $l=1$. When all other parameters are fixed, the bound grows as $\mathcal{O}(l^2)$ in $l$ and as $\mathcal{O}(D^4)$ in $D$, respectively.
• Figure 4: Impact of ensemble size (number of trained models) on Message Flooding accuracy as message size $l$ (left) and maximum node degree $D$ (right) increase. In the left plot, we fix $D=2$ and in the right, $l=1$. Accuracy is averaged over all trees with $n=7$ nodes. Accuracy improves with ensemble size, and larger $l$ or $D$ typically requires larger ensembles to match simpler settings.
• Figure 5: Accuracy on the ablation task for variants of \ref{['eq:architecture']}. We run Message Flooding with $l=3$-bit messages, and measure accuracy over all non-empty messages. Variants enable or disable $\Psi_{\mathrm{enc}}$, disable $\Psi_H$, or apply $\Psi_H$ before (baseline) or after aggregation. Omitted curves indicate zero accuracy throughout.

Theorems & Definitions (17)

• Theorem 5.1: Learnability and execution of any LOCAL model using the GNN in \ref{['eq:architecture']} (Informal)
• proof : Proof outline
• Theorem 6.1: Message Flooding
• Theorem 6.2: BFS
• Theorem 6.3: DFS
• Theorem 6.4: Bellman-Ford
• Lemma D.1: Turing Completeness of the Template Matching Framework
• proof
• Remark D.2
• Lemma D.3: Expressivity of the graph template-matching framework