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 different inputs for which an algorithm will perform certain intermediate computations 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 exactly the same next trajectory step. We ensure invariance in the next-step prediction across such inputs, by employing a self-supervised objective derived by our observation, formalised in a causal graph. We prove that the resulting method, which we call Hint-ReLIC, 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）测试数据上性能显著下降。在本工作中，我们提出一个重要观察：对于许多不同输入，算法在执行过程中某些中间计算步骤会产生完全相同的结果。这一洞察使我们能够开发数据增强流程，针对算法的中间轨迹生成与目标算法下一步轨迹完全相同的输入。我们通过自监督目标函数（基于观察结果推导并形式化为因果图）确保跨此类输入的下一跳预测具有不变性。我们证明所提出的方法（称为Hint-ReLIC）能提升推理器的OOD泛化能力。在CLRS算法推理基准上评估显示，该方法在OOD测试数据上可实现高达3倍的性能提升。