[论文解读] Unveiling Transformers with LEGO: a synthetic reasoning task

Yi Zhang, Artūrs Bačkurs+4 more|arXiv (Cornell University)|Jun 9, 2022

Explainable Artificial Intelligence (XAI)被引用 4

一句话总结

本文提出 LEGO，一种合成推理任务，用于评估 Transformer 模型通过结构化注意力机制进行逻辑推理链追踪的能力。作者识别出一种新型的‘关联’注意力模式，该模式专门针对相同标记，展示了预训练如何通过结构化注意力提升性能，并提出一种 LEGO 注意力模块，在降低 FLOPs 的同时保持或提升大规模预训练中的准确率。

ABSTRACT

We propose a synthetic reasoning task, LEGO (Learning Equality and Group Operations), that encapsulates the problem of following a chain of reasoning, and we study how the Transformer architectures learn this task. We pay special attention to data effects such as pretraining (on seemingly unrelated NLP tasks) and dataset composition (e.g., differing chain length at training and test time), as well as architectural variants such as weight-tied layers or adding convolutional components. We study how the trained models eventually succeed at the task, and in particular, we manage to understand some of the attention heads as well as how the information flows in the network. In particular, we have identified a novel \emph{association} pattern that globally attends only to identical tokens. Based on these observations we propose a hypothesis that here pretraining helps for LEGO tasks due to certain structured attention patterns, and we experimentally verify this hypothesis. We also observe that in some data regime the trained transformer finds ``shortcut" solutions to follow the chain of reasoning, which impedes the model's robustness, and moreover we propose ways to prevent it. Motivated by our findings on structured attention patterns, we propose the LEGO attention module, a drop-in replacement for vanilla attention heads. This architectural change significantly reduces Flops and maintains or even \emph{improves} the model's performance at large-scale pretraining.

研究动机与目标

• 设计一种合成推理任务，以隔离并评估 Transformer 在序列推理方面的能力。
• 研究在无关 NLP 任务上进行预训练以及数据集构成对推理任务泛化能力的影响。
• 识别并分析注意力头行为，特别是新型结构化注意力模式（如‘关联’）。
• 理解为何预训练能提升推理任务上的表现，并通过实验验证该假设。
• 提出一种更高效的注意力机制——LEGO 注意力，以在不牺牲性能的前提下降低计算成本。

提出的方法

• LEGO 任务被设计为一种合成推理基准，模型需学习追踪并推理由相同标记构成的链式结构。
• 作者在 LEGO 数据上训练 Transformer 模型，使用不同链长，并通过激活分析和注意力模式分析来研究注意力头行为。
• 提出一种新型的‘关联’注意力模式，该模式仅全局关注相同标记，从而实现对序列中重复元素的精确追踪。
• 提出假设：预训练通过诱导结构化注意力模式来提升性能，该假设通过受控的消融实验进行验证。
• 提出 LEGO 注意力模块作为标准多头注意力的即插即用替代品，通过限制注意力仅作用于相同标记来减少 FLOPs。
• 在多种数据设置下进行实验，包括训练和测试时的链长变化，以评估模型的鲁棒性及对捷径学习的敏感性。

实验结果

研究问题

• RQ1 在无关 NLP 任务上进行预训练以及数据集构成如何影响 Transformer 在未见链长的推理任务中的泛化能力？
• RQ2 在 LEGO 任务训练过程中，哪些特定注意力模式出现？它们如何支持对标记链的推理？
• RQ3 为何预训练能提升在合成推理任务上的表现？这是否可归因于结构化注意力模式的出现？
• RQ4 在哪些数据设置下，模型会学习到依赖表面统计线索而非真实推理的‘捷径’解决方案？此类行为如何缓解？
• RQ5 与标准注意力相比，改进的注意力机制——LEGO 注意力——是否能在大规模预训练中实现更高的效率和性能？

主要发现

• 识别出一种新型的‘关联’注意力模式，特定注意力头会全局仅关注相同标记，从而实现对序列中重复元素的精确追踪。
• 在无关 NLP 任务上进行预训练可提升 LEGO 任务表现，且该提升与结构化注意力模式的出现高度相关。
• 在某些数据设置下，模型会学习到依赖表面统计线索的捷径解决方案，从而降低对分布偏移的鲁棒性。
• 所提出的 LEGO 注意力模块显著降低了 FLOPs，同时在大规模预训练基准上保持或提升了性能。
• 研究结果表明，注意力模式并非随机，而是具有结构化的；架构选择（如权重共享或卷积组件）会影响推理行为。
• 研究结果表明，预训练对推理任务的益处可能源于诱导出特定的、与任务无关的注意力归纳偏置，而不仅仅是数据多样性。