[论文解读] A Generalist Neural Algorithmic Learner

Borja Ibarz, Vitaly Kurin+13 more|arXiv (Cornell University)|Sep 22, 2022

Machine Learning and Algorithms被引用 7

一句话总结

本文提出了一种通用神经算法学习器——一个具有共享参数的单一图神经网络，能够同时学习并执行多种经典算法，如排序、最短路径、动态规划和几何计算。通过改进输入表示、训练制度和GNN架构，该模型在CLRS-30基准上实现了平均性能20%以上的绝对提升，并在分布外泛化能力上达到或超过专用模型的表现，展示了在异构控制流之间有效整合多任务知识的能力。

ABSTRACT

The cornerstone of neural algorithmic reasoning is the ability to solve algorithmic tasks, especially in a way that generalises out of distribution. While recent years have seen a surge in methodological improvements in this area, they mostly focused on building specialist models. Specialist models are capable of learning to neurally execute either only one algorithm or a collection of algorithms with identical control-flow backbone. Here, instead, we focus on constructing a generalist neural algorithmic learner -- a single graph neural network processor capable of learning to execute a wide range of algorithms, such as sorting, searching, dynamic programming, path-finding and geometry. We leverage the CLRS benchmark to empirically show that, much like recent successes in the domain of perception, generalist algorithmic learners can be built by "incorporating" knowledge. That is, it is possible to effectively learn algorithms in a multi-task manner, so long as we can learn to execute them well in a single-task regime. Motivated by this, we present a series of improvements to the input representation, training regime and processor architecture over CLRS, improving average single-task performance by over 20% from prior art. We then conduct a thorough ablation of multi-task learners leveraging these improvements. Our results demonstrate a generalist learner that effectively incorporates knowledge captured by specialist models.

研究动机与目标

• 开发一个能够通过共享参数学习并执行多种经典算法任务的单一神经网络。
• 克服专用模型在控制流结构相似性或单一算法上的局限性。
• 通过解决联合训练中放大的数值不稳定性问题，提升神经算法推理中多任务学习的稳定性和性能。
• 证明通过改进的训练和架构组件，专用模型的知识可被有效整合到通用学习器中。
• 建立一个可扩展、通用的框架，适用于多种算法领域中的神经算法推理。

提出的方法

• 模型采用具有共享权重的统一图神经网络处理器，通过特定任务的编码器-解码器头处理多种算法任务。
• 通过引入结构化、与任务无关的特征，增强输入表示，以提升模型泛化能力和训练稳定性。
• 引入分块训练制度，以减少多任务学习期间的梯度不稳定性，提升收敛性和性能。
• 避免使用带噪声的教师强制策略，以防止训练与推理之间的分布偏移，增强分布外泛化能力。
• 架构改进包括三元组推理模块、门控消息传递机制，以及基于Sinkhorn的可微路由与注意力操作。
• 应用梯度裁剪和效率优化，以稳定训练并加速在CLRS-30基准上的收敛速度。

实验结果

研究问题

• RQ1 一个单一神经网络是否能有效学习并执行具有根本不同控制流的多种算法任务？
• RQ2 与单任务专用模型相比，多任务训练在多大程度上能提升在多样化算法任务上的泛化能力？
• RQ3 输入表示、训练制度和架构的改进如何共同提升多任务算法推理性能？
• RQ4 在训练期间避免教师强制噪声是否能带来更好的分布外泛化表现？
• RQ5 通用模型是否能匹配或超越专用模型的性能，尤其是在更大、未见过的实例上？

主要发现

• 所提出的通用模型在CLRS-30基准上的平均性能相比先前最先进方法实现了20.3%的绝对提升。
• 该模型在多个任务（包括排序和最短路径）上的分布外泛化性能与对应单任务专用模型相当或更优。
• 消融实验表明，教师强制噪声会显著损害泛化能力，当训练期间使用真实标签提示时性能明显下降。
• 分块训练制度相比原始多任务训练带来了显著的性能提升，表明训练稳定性得到改善。
• 架构改进的组合（如三元组推理模块和Sinkhorn操作符）对模型的鲁棒性和性能有显著贡献。
• 结果表明，当数值稳定性和表示质量得到仔细处理时，神经算法推理中的有效多任务学习是可行的。