[论文解读] Neural Execution Engines: Learning to Execute Subroutines

Yujun Yan, Kevin Swersky+3 more|arXiv (Cornell University)|Jun 15, 2020

Topic Modeling参考文献 6被引用 7

一句话总结

本文提出神经执行引擎（NEEs），一种基于Transformer的架构，能够以强大的泛化能力学习执行排序和最短路径等算法子程序。通过引入可学习的条件掩码机制并使用二进制数表示，NEEs在远超训练分布范围的输入上实现近乎完美的准确率，包括更长的序列和未见过的数据范围。

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 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.

研究动机与目标

• 解决神经网络在算法推理中缺乏强泛化能力的问题，特别是在输入超出训练分布范围时。
• 克服Transformer在长序列上处理数值相近值时的注意力退化问题，该问题会导致误预测和误差累积。
• 通过二进制表示学习结构化、插值型的数字符号表示，实现对未见数据分布的泛化能力。
• 设计一种模块化框架，使单个训练好的子程序（如比较操作）可被重用于Dijkstra算法和Prim算法等多种算法，而无需重新训练。
• 在多种图类型和输入分布下展现鲁棒性，证明在分布偏移下性能下降极小。

提出的方法

• 使用序列到序列的Transformer模型，为每个子程序步骤预测一个值和一个指针，实现动态状态追踪。
• 引入可学习的条件掩码机制，根据预测的指针动态更新编码器掩码，提升长序列上的注意力保真度。
• 以二进制形式表示数字，构建分层的、指数可扩展的嵌入表示，支持对未见数值的插值。
• 在加法和乘法等下游任务上进行训练，以学习结构化、可解释的数字符号表示。
• 通过在子程序而非完整算法上进行训练，将算法逻辑的学习与数据分布解耦，实现模块化复用。
• 对不同算法（如Dijkstra算法、Prim算法）应用相同的NEE，使用相同的已学习子程序，无需重新训练，证明其可迁移性。

实验结果

研究问题

• RQ1 基于Transformer的模型能否在训练分布之外实现对排序和最短路径等子程序的强泛化执行？
• RQ2 标准Transformer为何在长序列或数值相近的输入上无法泛化？如何保持注意力的保真度？
• RQ3 二进制数字符号表示能否生成结构化、可插值的嵌入，从而支持对未见数值的泛化？
• RQ4 单个已学习的子程序（如比较操作）在不重新训练的情况下，能在多大程度上被复用于多种不同算法？
• RQ5 神经执行引擎在输入数据分布发生偏移（如不同图类型或输入范围）时，其鲁棒性如何？

主要发现

• 在测试图规模达100的情况下，NEEs在Dijkstra最短路径和Prim最小生成树任务上均达到100.00%的准确率，即使训练时仅使用大小为8的图。
• 即使测试图的类型与训练图不同，图算法性能仍保持近乎完美（最短路径为99.91%，MST为93.00%），证明对分布偏移具有强鲁棒性。
• 通过二进制数字符号表示，模型学习到结构化、可解释的嵌入，支持插值——例如，训练时未见的65%的数值被准确插值。
• 可学习的条件掩码机制实现了强泛化：在短序列（如长度为8）上训练的模型，能准确泛化至长达100个元素的序列。
• 在单一子程序（如比较操作）上训练的NEEs可直接应用于多种算法（如Dijkstra、Prim）而无需重新训练，展现出高度的模块化与可迁移性。
• 采用二进制嵌入的模型在参数效率和可扩展性方面优于独热输出，同时在所有任务和输入范围内保持高性能。