Recent work on neural algorithmic reasoning has investigated the reasoning capabilities of neural networks, effectively demonstrating they can learn to execute classical algorithms on unseen data coming from the train distribution. However, the performance of existing neural reasoners significantly degrades on out-of-distribution (OOD) test data, where inputs have larger sizes. In this work, we make an important observation: there are many \emph{different} inputs for which an algorithm will perform certain intermediate computations \emph{identically}. This insight allows us to develop data augmentation procedures that, given an algorithm's intermediate trajectory, produce inputs for which the target algorithm would have \emph{exactly} the same next trajectory step. Then, we employ a causal framework to design a corresponding self-supervised objective, and we prove that it improves the OOD generalisation capabilities of the reasoner. We evaluate our method on the CLRS algorithmic reasoning benchmark, where we show up to 3$\times$ improvements on the OOD test data.

翻译：近期关于神经算法推理的研究探讨了神经网络的推理能力，有效证明了它们能够学习执行训练数据分布中未见过的经典算法。然而，现有神经推理器在输入规模更大的分布外（OOD）测试数据上的性能显著下降。在本工作中，我们提出一个重要观察：对于许多不同的输入，算法在执行过程中会进行完全相同的中间计算。这一洞见使我们能够开发数据增强流程——给定算法的中间轨迹，生成与目标算法下一步轨迹完全一致的输入。进而，我们采用因果框架设计相应的自监督目标函数，并证明该目标能提升推理器的OOD泛化能力。我们在CLRS算法推理基准上评估了该方法，结果显示在OOD测试数据上实现了高达3倍的性能提升。