Open-Book Neural Algorithmic Reasoning

TL;DR

The paper addresses the limitation of standard neural algorithmic reasoning that relies on single-instance inputs by introducing an open-book framework that allows a model to consult training instances during reasoning. It adds a Dataset Encoder to embed training examples and an Open-Book Processor to integrate these representations via cross-attention with the current reasoning state, compatible with existing architectures and supporting both single-task and multi-task settings. Empirical results on the CLRS Algorithmic Reasoning Benchmark show significant performance gains across most tasks and the ability to replicate multi-task training effects with similar training effort, while also enabling interpretable insights through learned task-attention patterns. The framework advances practical neural algorithmic reasoning by combining memory-augmented reasoning with interpretable multi-task transfer, and future work will focus on ensuring robust gains across all tasks and exploring richer memory mechanisms.

Abstract

Neural algorithmic reasoning is an emerging area of machine learning that focuses on building neural networks capable of solving complex algorithmic tasks. Recent advancements predominantly follow the standard supervised learning paradigm -- feeding an individual problem instance into the network each time and training it to approximate the execution steps of a classical algorithm. We challenge this mode and propose a novel open-book learning framework. In this framework, whether during training or testing, the network can access and utilize all instances in the training dataset when reasoning for a given instance. Empirical evaluation is conducted on the challenging CLRS Algorithmic Reasoning Benchmark, which consists of 30 diverse algorithmic tasks. Our open-book learning framework exhibits a significant enhancement in neural reasoning capabilities. Further, we notice that there is recent literature suggesting that multi-task training on CLRS can improve the reasoning accuracy of certain tasks, implying intrinsic connections between different algorithmic tasks. We delve into this direction via the open-book framework. When the network reasons for a specific task, we enable it to aggregate information from training instances of other tasks in an attention-based manner. We show that this open-book attention mechanism offers insights into the inherent relationships among various tasks in the benchmark and provides a robust tool for interpretable multi-task training.

Figures (9)

• Figure 1: An illustration of the open-book framework. At each reasoning step $t$, we simultaneously input $(\mathbf{x},\mathbf{y}^{(t-1)})$ and instances from the training set $\mathbf{T}$, yielding $\mathbf{y}^{(t)}$.
• Figure 2: Comparison of the MPNN architecture's performance before and after augmentation with the open-book framework. The 30 tasks are arranged in descending order of improvement magnitude.
• Figure 3: Comparisons between our multi-task augmented approach and Triplet-GMPNN. The 30 tasks are arranged in descending order of improvement magnitude.
• Figure 4: Comparison of the PGN architecture's performance before and after augmentation with the open-book framework. The 30 tasks are arranged in descending order of improvement magnitude.
• Figure 5: Comparison of the Triplet-GMPNN architecture's performance before and after augmentation with the open-book framework. The 30 tasks are arranged in descending order of improvement magnitude.