Neural Execution Engines: Learning to Execute Subroutines

中文速览

神经网络虽然能学会排序、最短路等算法的表面操作，却常常一遇到更长的输入或训练中没见过的数字就失灵，原因是注意力在长序列、数值相近时难以准确找到该处理的位置。研究者在Transformer基础上设计了神经执行引擎（Neural Execution Engine, NEE），让模型同时输出计算结果和指针，再用条件掩码动态更新下一步关注的区域，并用二进制表示数字以支持更大范围的数值。实验表明，NEE在选择排序、归并排序、Dijkstra最短路和Prim最小生成树等任务上都能远超训练长度范围、接近完美地泛化，而且同一个已学会的比较子程序还能直接复用于多种算法；二进制嵌入也能对训练中缺失的数字进行可靠插值。这个结果说明，给模型加入类似程序指针和动态记忆访问的结构，可能是让神经网络真正学会可迁移算法规则、而不只是记住训练样本的关键。

摘要

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 leads to embeddings with rich structure once trained on downstream tasks like addition or multiplication. This allows the embedding to handle missing data by faithfully interpolating numbers not seen during training.

术语表

algorithmic reasoning

算法推理

strong generalization

强泛化

subroutine imitation

子程序模仿

numerical subroutine

数值子程序

sequence-to-sequence model

序列到序列模型

Transformer

Transformer

self-attention

自注意力

attention weights

注意力权重

attention entropy

注意力熵

conditional masking mechanism

条件掩码机制

learned conditional masking

学习式条件掩码

Neural Execution Engine (NEE)

Neural Execution Engine（NEE）

pointer

指针

encoder mask

编码器掩码

self-attention mask

自注意力掩码

dynamic graph

动态图

Graph Attention Networks

Graph Attention Networks

binary representation

二进制表示

binary embedding

二进制嵌入

interpolation

插值

selection sort

选择排序

merge sort

归并排序

Dijkstra’s algorithm

Dijkstra 算法

Prim’s algorithm

Prim 算法

shortest path

最短路径

minimum spanning tree

最小生成树

comparison subroutine

比较子程序

arithmetic subroutine

算术子程序

pointer manipulation

指针操作