用 LEGO 揭示 Transformer:一个合成推理任务
机器学习
2023-02-21 v3 人工智能
计算与语言
摘要
我们提出了一个合成推理任务 LEGO(Learning Equality and Group Operations,学习等式与群操作),它概括了遵循推理链的问题,并研究 Transformer 架构如何学习该任务。我们特别关注数据效应,例如(在看似无关的自然语言处理任务上的)预训练以及数据集构成(例如训练与测试时推理链长度不同),以及架构变体例如权重绑定层或添加卷积组件。我们研究训练后的模型最终如何成功完成该任务,特别是我们得以理解部分注意力头以及信息在网络中的流动方式。具体而言,我们识别出了一种新颖的“关联(association)”模式,该模式全局地仅关注相同的词元(token)。基于这些观察,我们提出一个假设:在此处预训练之所以有助于 LEGO 任务,是由于某些结构化的注意力模式,并通过实验验证了该假设。我们还观察到在某些数据情形下,训练后的 transformer 找到了遵循推理链的“捷径”解,这损害了模型的鲁棒性,并且我们提出了防止该现象的方法。受我们在结构化注意力模式上的发现启发,我们提出了 LEGO 注意力模块,一种可即插即用的标准注意力头替代方案。这一架构改变显著减少了 Flops,并在大规模预训练上保持甚至提升了模型性能。
引用
@article{arxiv.2206.04301,
title = {Unveiling Transformers with LEGO: a synthetic reasoning task},
author = {Yi Zhang and Arturs Backurs and Sébastien Bubeck and Ronen Eldan and Suriya Gunasekar and Tal Wagner},
journal= {arXiv preprint arXiv:2206.04301},
year = {2023}
}