Neural Execution Engines: Learning to Execute Subroutines
Yujun Yan The University of Michigan、Kevin Swersky Google Research、Danai Koutra The University of Michigan、Parthasarathy Ranganathan, Milad Hashemi Google Research {parthas, Work completed during an internship at Google.
Abstract
A significant effort has been made to train neural networks that replicate algorithmic reasoning, but they often fail to learn the abstract concepts underlying these algorithms. This is evidenced by their inability to generalize to data distributions that are outside of their restricted training sets, namely larger inputs and unseen data. We study these generalization issues at the level of numerical subroutines that comprise common algorithms like sorting, shortest paths, and minimum spanning trees. First, we observe that transformer-based sequence-to-sequence models can learn subroutines like sorting a list of numbers, but their performance rapidly degrades as the length of lists grows beyond those found in the training set. We demonstrate that this is due to attention weights that lose fidelity with longer sequences, particularly when the input numbers are numerically similar. To address the issue, we propose a learned conditional masking mechanism, which enables the model to strongly generalize far outside of its training range with near-perfect accuracy on a variety of algorithms. Second, to generalize to unseen data, we show that encoding numbers with a binary representation
中文速览
让神经网络真正"学会"算法而不只是记住训练数据里的模式，一直是个难题——模型在训练集范围内表现不错，一旦遇到更长的序列或没见过的数字就迅速崩溃。研究者聚焦于排序、最短路径、最小生成树等经典算法的基本子程序，发现基于Transformer的序列到序列模型之所以失败，根本原因是注意力权重（attention weights）在面对更长序列时变得"散焦"，尤其当输入数字彼此接近时更为严重。为此，他们提出了"神经执行引擎"（Neural Execution Engine，NEE），核心创新是一个可学习的条件掩码机制（learned conditional masking）：模型每步输出一个指针，该指针通过小型一维卷积网络动态更新编码器的注意力掩码，让模型始终精准聚焦在当前需要处理的那部分输入上；与此同时，用二进制表示数字并训练对应的嵌入向量，使模型能够自然地处理训练时从未出现过的数值。实验结果表明，NEE在远超训练范围的序列长度和数值域上都能达到接近完美的准确率，同一个比较子程序无需重新训练即可直接复用于Dijkstra和Prim等不同算法，这为构建真正具备强泛化能力的算法推理神经网络提供了切实可行的路径。
原文 arXiv:2006.08084;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2006.08084v3