Discrete Neural Algorithmic Reasoning

TL;DR

This paper introduces a neural reasoning approach that maintains discrete execution states, enabling neural models to perfectly replicate classical algorithms and generalize reliably across different data distributions.

Contribution

It proposes a novel neural architecture that separates discrete states from continuous data, allowing for perfect alignment with classical algorithms and proven correctness.

Findings

Achieved perfect test scores on multiple algorithmic tasks.

Enabled proof of correctness for learned algorithms.

Demonstrated robust generalization across out-of-distribution data.

Abstract

Neural algorithmic reasoning aims to capture computations with neural networks by training models to imitate the execution of classical algorithms. While common architectures are expressive enough to contain the correct model in the weight space, current neural reasoners struggle to generalize well on out-of-distribution data. On the other hand, classical 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.