Persistent Message Passing

Abstract

Graph neural networks (GNNs) are a powerful inductive bias for modelling algorithmic reasoning procedures and data structures. Their prowess was mainly demonstrated on tasks featuring Markovian dynamics, where querying any associated data structure depends only on its latest state. For many tasks of interest, however, it may be highly beneficial to support efficient data structure queries dependent on previous states. This requires tracking the data structure's evolution through time, placing significant pressure on the GNN's latent representations. We introduce Persistent Message Passing (PMP), a mechanism which endows GNNs with capability of querying past state by explicitly persisting it: rather than overwriting node representations, it creates new nodes whenever required. PMP generalises out-of-distribution to more than 2× larger test inputs on dynamic temporal range queries, significantly outperforming GNNs which overwrite states.

Graph Neural Networks are Dynamic Programmers

“…However, it quickly became apparent that it is not enough to just train any GNN-for many algorithmic tasks, careful attention is required. Several papers illustrated special cases of GNNs that align with sequential algorithms (Veličković et al, 2019), linearithmic sequence processing (Freivalds et al, 2019), physics simulations (Sanchez-Gonzalez et al, 2020), iterative algorihtms (Tang et al, 2020), data structures (Veličković et al, 2020) or auxiliary memory (Strathmann et al, 2021). Some explanations for this lack of easy generalisation have arisen-we now have both geometric (Xu et al, 2020) and causal (Bevilacqua et al, 2021) views into how better generalisation can be achieved.…”

Neural algorithmic reasoning

“…Lastly, there are representational issues associated with dynamically allocated memory-it may be unclear what is the best way to represent the internal memory storage and its usage in algorithm trajectories. One example of the ambiguity is in asking whether the algorithm executor should start with a "scratch space" defined by the space complexity of the problem that gets filled up, or dynamically generate such space 5 (Strathmann et al, 2021). As such, we for now exclude all algorithms that require allocating memory which cannot be directly attached to the set of objects provided at input time.…”