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基础上设计了神经执行引擎（Neural Execution Engine, NEE），让模型同时输出计算结果和指针，再用条件掩码动态更新下一步关注的区域，并用二进制表示数字以支持更大范围的数值。实验表明，NEE在选择排序、归并排序、Dijkstra最短路和Prim最小生成树等任务上都能远超训练长度范围、接近完美地泛化，而且同一个已学会的比较子程序还能直接复用于多种算法；二进制嵌入也能对训练中缺失的数字进行可靠插值。这个结果说明，给模型加入类似程序指针和动态记忆访问的结构，可能是让神经网络真正学会可迁移算法规则、而不只是记住训练样本的关键。
原文 arXiv:2006.08084;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2006.08084v3