Discrete Neural Algorithmic Reasoning

Neural algorithmic reasoning aims to capture computations with neural networks by training models to imitate the execution of classic algorithms. While common architectures are expressive enough to contain the correct model in the weights space, current neural reasoners struggle to generalize well on out-of-distribution data. On the other hand, classic computations are not affected by distributional shifts as they can be described as transitions between discrete computational states. In this work, we propose to force neural reasoners to maintain the execution trajectory as a combination of finite predefined states. To achieve this, we separate discrete and continuous data flows and describe the interaction between them. Trained with supervision on the algorithm's state transitions, such models are able to perfectly align with the original algorithm. To show this, we evaluate our approach on multiple algorithmic problems and achieve perfect test scores both in single-task and multitask setups. Moreover, the proposed architectural choice allows us to prove the correctness of the learned algorithms for any test data.

Neural networks often struggle to reliably perform algorithmic tasks, especially when tested on data that differs from what they were trained on. Unlike traditional algorithms, which follow precise, predefined steps to guarantee correct results, neural networks can produce unpredictable errors when faced with unfamiliar inputs.In this work, we design neural networks that mimic classical algorithms more closely by enforcing a structured, step-by-step approach. Our key idea is to split the computation into two parts: one that handles discrete, algorithm-like states and another that processes continuous data. By combining these two flows and training the model with explicit guidance on how algorithms transition between states, we ensure the network behaves predictably and correctly—even on unseen or larger-scale inputs.We test our approach on a variety of algorithmic tasks and show that these models achieve perfect accuracy, matching the performance of classical algorithms. Importantly, our method provides guarantees that the learned neural networks will always produce correct results, no matter the input. This advance could lead to more reliable and interpretable AI systems for tasks requiring precise, algorithmic reasoning.