Learning Graph Algorithms With Recurrent Graph Neural Networks

TL;DR

This paper introduces a recurrent GNN architecture with specific techniques that enable learning graph algorithms on small graphs and successfully extrapolating to larger graphs, addressing generalization challenges.

Contribution

The paper presents three key techniques—skip connections, state regularization, and edge convolutions—that improve the scalability and extrapolation ability of recurrent GNNs for graph algorithms.

Findings

Recurrent GNNs can be trained on small graphs and applied to larger ones.

The identified techniques significantly enhance GNN extrapolation capabilities.

Empirical validation shows successful generalization on algorithmic datasets.

Abstract

Classical graph algorithms work well for combinatorial problems that can be thoroughly formalized and abstracted. Once the algorithm is derived, it generalizes to instances of any size. However, developing an algorithm that handles complex structures and interactions in the real world can be challenging. Rather than specifying the algorithm, we can try to learn it from the graph-structured data. Graph Neural Networks (GNNs) are inherently capable of working on graph structures; however, they struggle to generalize well, and learning on larger instances is challenging. In order to scale, we focus on a recurrent architecture design that can learn simple graph problems end to end on smaller graphs and then extrapolate to larger instances. As our main contribution, we identify three essential techniques for recurrent GNNs to scale. By using (i) skip connections, (ii) state regularization, and (iii) edge convolutions, we can guide GNNs toward extrapolation. This allows us to train on small graphs and apply the same model to much larger graphs during inference. Moreover, we empirically validate the extrapolation capabilities of our GNNs on algorithmic datasets.