Richer Representations for Neural Algorithmic Reasoning via Auxiliary Reconstruction
Authors: Jiafu Huang, Chao Peng, Chenyang Xu, Zhengfeng Yang, Kecheng Cai, Chenhao Zhang, Yi Wang, Yiwei Gong, +2 more
Organizations: Shanghai Key Laboratory of Trustworthy Computing, East China Normal University, China · Overseas Education College, Xiamen University, China · Dept. of Electrical and Electronic Engineering, The University of Melbourne, Australia
Neural algorithmic reasoning has emerged as a popular research direction. It aims to train neural networks to mimic the step-by-step behavior of classical rule-based algorithms. More specifically, the execution of such algorithms can be abstracted as a sequence of states, where each state represents the intermediate outcome after an execution step. The training objective is to generate state sequences that replicate the underlying algorithmic process. A common framework for this task adopts an encoder-processor-decoder architecture, where the encoder learns representations of states, the processor simulates algorithmic steps, and the decoder reconstructs output states. While prior work has focused on improving the processor, the role of the encoder in representation learning has received little attention. Most methods rely on simple MLP encoders, raising the question of whether such representations are sufficiently informative for supporting algorithmic reasoning. This paper investigates how to improve encoder representations for neural algorithmic reasoning. We propose a reconstruction module that aims to recover the input state from its encoded representation. This auxiliary reconstruction task encourages the encoder to retain critical information about the input. We demonstrate that incorporating this task during training improves the performance of existing neural architectures on standard benchmarks. Furthermore, we observe that current encoders often underutilize the correlations among features within a state. To address this, we draw inspiration from self-supervised learning and design an enhanced variant of the auxiliary task that encourages the encoder to capture intra-state feature dependencies. Experimental results show that our method enables the encoder to learn richer representations, thereby enhancing the performance of existing processors on algorithmic reasoning tasks.
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(knL) 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.
Alex Schutz, Victor-Alexandru Darvariu, Efimia Panagiotaki +2
How should future neural reasoning systems implement extended computation? Recursive Reasoning Models (RRMs) offer a promising alternative to autoregressive sequence extension by performing iterative latent-state refinement with shared transition functions. Yet existing RRMs are largely deterministic, following a single latent trajectory and converging to a single prediction. We introduce Generative Recursive reAsoning Models (GRAM), a framework that turns recursive latent reasoning into probabilistic multi-trajectory computation. GRAM models reasoning as a stochastic latent trajectory, enabling multiple hypotheses, alternative solution strategies, and inference-time scaling through both recursive depth and parallel trajectory sampling. This yields a latent-variable generative model supporting conditional reasoning via pθ(y∣x) and, with fixed or absent inputs, unconditional generation via pθ(x). Trained with amortized variational inference, GRAM improves over deterministic recurrent and recursive baselines on structured reasoning and multi-solution constraint satisfaction tasks, while demonstrating an unconditional generation capability. https://ahn-ml.github.io/gram-website