# Richer Representations for Neural Algorithmic Reasoning via Auxiliary Reconstruction

arXiv: https://arxiv.org/abs/2606.00559 · HTML: https://mlsanity.com/paper/2606.00559/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 20 of its 27 references that are in the map.
- First posted on arXiv 30 May 2026.

## Builds on (most cited first, 20 of 20)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Deep contextualized word representations](https://mlsanity.com/paper/1802.05365/) (2018)
- [Variational Graph Auto-Encoders](https://mlsanity.com/paper/1611.07308/) (2016)
- [Unsupervised Visual Representation Learning by Context Prediction](https://mlsanity.com/paper/1505.05192/) (2015)
- [Strategies for Pre-training Graph Neural Networks](https://mlsanity.com/paper/1905.12265/) (2019)
- [GraphMAE: Self-Supervised Masked Graph Autoencoders](https://mlsanity.com/paper/2205.10803/) (2022)
- [SimGRACE: A Simple Framework for Graph Contrastive Learning without Data Augmentation](https://mlsanity.com/paper/2202.03104/) (2022)
- [GraphMAE2: A Decoding-Enhanced Masked Self-Supervised Graph Learner](https://mlsanity.com/paper/2304.04779/) (2023)
- [A Generalist Neural Algorithmic Learner](https://mlsanity.com/paper/2209.11142/) (2022)
- [Pointer Graph Networks](https://mlsanity.com/paper/2006.06380/) (2020)
- [Neural Algorithmic Reasoners are Implicit Planners](https://mlsanity.com/paper/2110.05442/) (2021)
- [Dual Algorithmic Reasoning](https://mlsanity.com/paper/2302.04496/) (2023)
- [Neural Priority Queues for Graph Neural Networks](https://mlsanity.com/paper/2307.09660/) (2023)
- [Recurrent Aggregators in Neural Algorithmic Reasoning](https://mlsanity.com/paper/2409.07154/) (2024)
- [A Context-Enhanced Framework for Sequential Graph Reasoning](https://mlsanity.com/paper/2412.09056/) (2024)

## Abstract (arXiv)

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.