Richer Representations for Neural Algorithmic Reasoning via Auxiliary Reconstruction

2026-05-30

AI Summary

Existing neural algorithmic reasoning approaches are limited by insufficient representational capacity of their encoders, hindering effective simulation of classical algorithm execution. This work proposes an encoder enhancement mechanism based on an auxiliary reconstruction task, which leverages self-supervised learning to model dependencies among internal features of the input state, thereby strengthening the preservation and expressiveness of state information. Moving beyond prior methods that solely optimize the processor module, the proposed architecture achieves substantial performance gains on standard neural algorithmic reasoning benchmarks, demonstrating its ability to learn richer and more structured state representations.

Abstract

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.

Problem

• neural algorithmic reasoning
• encoder representation
• state representation
• feature dependencies
• representation learning

Innovation

• auxiliary reconstruction
• neural algorithmic reasoning
• representation learning
• self-supervised learning
• encoder enhancement