Recursive Algorithmic Reasoning

TL;DR

This paper introduces a stack-augmented graph neural network that enables recursive reasoning, significantly improving the generalization of neural models to larger graphs in algorithmic tasks like DFS.

Contribution

The paper proposes augmenting GNNs with a stack to enable recursive reasoning, addressing memory limitations and improving algorithmic generalization.

Findings

Enhanced generalization to larger graphs in DFS tasks

Stack augmentation enables recursive reasoning in GNNs

Significant improvement over prior methods in algorithmic tasks

Abstract

Learning models that execute algorithms can enable us to address a key problem in deep learning: generalizing to out-of-distribution data. However, neural networks are currently unable to execute recursive algorithms because they do not have arbitrarily large memory to store and recall state. To address this, we (1) propose a way to augment graph neural networks (GNNs) with a stack, and (2) develop an approach for capturing intermediate algorithm trajectories that improves algorithmic alignment with recursive algorithms over previous methods. The stack allows the network to learn to store and recall a portion of the state of the network at a particular time, analogous to the action of a call stack in a recursive algorithm. This augmentation permits the network to reason recursively. We empirically demonstrate that our proposals significantly improve generalization to larger input graphs over prior work on depth-first search (DFS).