From Shortcut Learning to Discrete Neural Insertion Sort

The recent field of neural algorithmic reasoning (NAR) studies the ability of graph neural networks (GNNs) to emulate classical algorithms like Bellman-Ford, a phenomenon known as algorithmic alignment. At the same time, recent advances in large language models (LLMs) have spawned the study of mechanistic interpretability, which aims to identify granular model components like circuits that perform specific computations. In this work, we introduce Mechanistic Interpretability for Neural Algorithmic Reasoning (MINAR), an efficient circuit discovery toolbox that adapts attribution patching methods from mechanistic interpretability to the GNN setting. We show through two case studies that MINAR recovers faithful neuron-level circuits from GNNs trained on algorithmic tasks. Our study sheds new light on the process of circuit formation and pruning during training, as well as giving new insight into how GNNs trained to perform multiple tasks in parallel reuse circuit components for related tasks. Our code is available at https://github.com/pnnl/MINAR.

Algorithmic reasoning -- the ability to perform step-by-step logical inference -- is a synthetic benchmark for evaluating multi-step reasoning abilities, designed for graph neural networks and also for transformer models. Prior work has evaluated reasoning for executing a single algorithmic task, whereas a more desirable objective is to perform multiple algorithmic reasoning tasks simultaneously. We start by noting that this is inherently difficult due to differences arising from the execution traces of the algorithms (such as depth- vs. breadth-first search), which cause interference when they are trained together. In this paper, we introduce {branching neural networks}, a new architecture for multitask algorithmic reasoning. The main idea is to search for a recursive tree-structured partition of $n$ algorithmic tasks into a $k$-ary tree (divided into $L$ layers). Naive search requires $O(k^{nL})$ complexity; we develop an algorithm that reduces this to $O(nL)$ by solving a convex relaxation at each layer to approximate an optimal partition. Our approach clusters these tasks using gradient-based affinity and can be used on top of any base model. We validate our approach on algorithmic reasoning benchmarks and their extensions with text descriptions. We show that gradient-based affinity scores help estimate true performance with less than 5% error, measured across eight different architectures with up to 34 billion parameters. On the CLRS benchmark, our approach outperforms existing graph neural networks by 3.7% and baselines by 1.2%, while reducing runtime by 48% and memory usage by 26%. The learned branching structure shows a hierarchical clustering of related algorithms. On three text-based graph reasoning benchmarks, our approach improves over baseline methods by 3.2%. Finally, we validate our approach for overlapping community detection.

Neural algorithmic reasoning (NAR) is a paradigm that trains neural networks to execute classic algorithms by supervised learning. Despite its successes, important limitations remain: inability to construct valid solutions without post-processing and to reason about multiple correct ones, poor performance on combinatorial NP-hard problems, and inapplicability to problems for which strong algorithms are not yet known. To address these limitations, we reframe the problem of learning algorithm trajectories as a Markov decision process, which imposes structure on the solution construction procedure and unlocks the powerful tools of imitation and reinforcement learning (RL). We propose the GNARL framework, encompassing the methodology to translate problem formulations from NAR to RL and a learning architecture suitable for a wide range of graph-based problems. We achieve high rates of reaching correct states on several CLRS-30 problems and performance matching or exceeding much narrower NAR approaches for NP-hard problems. Remarkably, GNARL remains applicable when no expert algorithm is available, though reward-driven learning can exhibit greater variability than direct supervision.