[论文解读] Towards Better Out-of-Distribution Generalization of Neural Algorithmic Reasoning Tasks

Sadegh Mahdavi, Kevin Swersky+4 more|arXiv (Cornell University)|Nov 1, 2022

Advanced Neural Network Applications被引用 4

一句话总结

本文通过识别CLRS基准中的独特挑战，提出一种针对输入分布偏移的修复方法，并引入一种基于注意力机制的2-WL图神经网络处理器，在30项算法任务上将SOTA模型的OOD性能提升了3%。

ABSTRACT

In this paper, we study the OOD generalization of neural algorithmic reasoning tasks, where the goal is to learn an algorithm (e.g., sorting, breadth-first search, and depth-first search) from input-output pairs using deep neural networks. First, we argue that OOD generalization in this setting is significantly different than common OOD settings. For example, some phenomena in OOD generalization of image classifications such as \emph{accuracy on the line} are not observed here, and techniques such as data augmentation methods do not help as assumptions underlying many augmentation techniques are often violated. Second, we analyze the main challenges (e.g., input distribution shift, non-representative data generation, and uninformative validation metrics) of the current leading benchmark, i.e., CLRS \citep{deepmind2021clrs}, which contains 30 algorithmic reasoning tasks. We propose several solutions, including a simple-yet-effective fix to the input distribution shift and improved data generation. Finally, we propose an attention-based 2WL-graph neural network (GNN) processor which complements message-passing GNNs so their combination outperforms the state-of-the-art model by a 3% margin averaged over all algorithms. Our code is available at: \url{https://github.com/smahdavi4/clrs}.

研究动机与目标

• 研究神经算法推理中OOD泛化所面临的独特挑战，这些挑战与标准视觉OOD设置不同。
• 识别CLRS基准中的关键问题，包括输入分布偏移、代表性不足的数据生成方式，以及具有误导性的分布内验证指标。
• 通过改进数据分布、输入表示修复和新型GNN架构，提升OOD泛化能力。
• 设计一种结合消息传递GNN与基于注意力机制的2-WL处理器的混合模型，以提升对未见图大小的性能表现。
• 在30项算法推理任务上实现SOTA的OOD泛化性能，平均准确率达到74.33%。

提出的方法

• 通过修改输入表示以更好地反映算法不变性，提出一种简单但有效的输入分布偏移修复方法。
• 扩大分布内训练数据集，以推动OOD泛化极限并提升模型鲁棒性。
• 设计一种基于注意力机制的2-WL图神经网络处理器，以增强其表达能力，超越标准消息传递GNN。
• 将基于注意力机制的2-WL处理器与消息传递GNN结合，构建一种混合架构，其性能优于SOTA。
• 利用中间算法步骤的提示作为额外监督信号，引导模型学习逐步推理过程。
• 在包含节点、边和图级别特征的图结构输入上，端到端训练模型，并为全部30项任务使用统一框架。

实验结果

研究问题

• RQ1 神经算法推理中的OOD泛化动态与标准图像分类OOD设置有何不同？
• RQ2 CLRS基准中OOD泛化表现不佳的主要原因是什么，特别是关于数据分布和验证指标方面？
• RQ3 通过架构或表示层面的修改，能否有效缓解输入分布偏移问题？
• RQ4 当与消息传递GNN结合时，基于注意力机制的2-WL GNN处理器在多大程度上能提升OOD泛化能力？
• RQ5 改进的数据生成与模型选择策略能否带来更可靠的OOD性能估计并实现更好的泛化？

主要发现

• CLRS基准表现出标准视觉基准中未见的独特OOD挑战，如输入分布偏移和代表性不足的数据生成方式。
• 分布内验证指标接近完美，导致其无法可靠预测OOD性能。
• 对输入分布偏移的简单修复显著提升了OOD泛化能力，尤其在更大图尺寸下效果更明显。
• 所提出的基于注意力机制的2-WL GNN处理器在全部30项算法任务上平均优于SOTA模型3%。
• 混合模型在所有算法上的平均OOD测试准确率达到74.33%，表明其对未见图大小具有强大的泛化能力。
• 尽管OOD准确率较高，模型在某一特定基准任务上仍表现不佳，表明在算法不变性学习方面仍有改进空间。