[论文解读] Pay Attention when Required

Swetha Mandava, Szymon Migacz+1 more|arXiv (Cornell University)|Sep 9, 2020

Topic Modeling参考文献 31被引用 5

一句话总结

本文提出PAR Transformer，一种更高效的Transformer架构，在WikiText-103上保持最先进困惑度的同时，将自注意力模块减少63%。通过可微神经架构搜索，发现仅在前三分之二层需要自注意力机制，从而在多个数据集和模型（包括BERT）上实现35%更低的推理延迟，且准确率无下降。

ABSTRACT

Transformer-based models consist of interleaved feed-forward blocks - that capture content meaning, and relatively more expensive self-attention blocks - that capture context meaning. In this paper, we explored trade-offs and ordering of the blocks to improve upon the current Transformer architecture and proposed PAR Transformer. It needs 35% lower compute time than Transformer-XL achieved by replacing ~63% of the self-attention blocks with feed-forward blocks, and retains the perplexity on WikiText-103 language modelling benchmark. We further validated our results on text8 and enwiki8 datasets, as well as on the BERT model.

研究动机与目标

• 通过重新思考自注意力与前馈网络模块的排列与比例，优化Transformer架构。
• 在不牺牲性能的前提下，降低大规模语言模型的计算成本与推理延迟。
• 通过自动化搜索识别最优架构模式，重点关注模块的排序与组合方式。
• 在多个数据集（WikiText-103、enwiki8、text8）和模型类型（BERT、Transformer-XL）上验证发现的泛化能力。
• 提供一种可解释的自动化方法，利用可微神经架构搜索设计高效Transformer家族。

提出的方法

• 采用可微神经架构搜索（NAS），搜索空间包含三种模块类型：自注意力、前馈网络和恒等映射。
• 构建一个超网络，其中每一层均为三种模块类型的可微加权组合，训练期间使用Gumbel-Softmax进行采样。
• 训练单一超网络，联合优化准确率与效率，实现在3^32种可能架构中的线性复杂度搜索。
• 利用学习到的注意力概率推导出设计规则：自注意力仅需存在于前三分之二层，且自注意力与前馈网络模块的最优比例为p:1（p=5）。
• 将推导出的规则应用于构建PAR Transformer与PAR BERT模型，显著减少自注意力模块数量，同时保持性能不变。
• 在多个基准测试（WikiText-103、enwiki8、text8）和任务（语言建模、问答、情感分析）上验证结果，使用标准超参数与推理设置。

实验结果

研究问题

• RQ1 Transformer编码器中自注意力与前馈网络模块的最优排列与比例为何？
• RQ2 在性能增益方面，额外的自注意力模块在何时开始出现收益递减？
• RQ3 减少自注意力模块数量是否可在显著降低推理成本的同时保持模型准确率？
• RQ4 最优架构模式是否在不同数据集与模型架构（如BERT、Transformer-XL）间具有泛化能力？
• RQ5 可微NAS能否用于推导出可解释的、高效的Transformer架构设计规则？

主要发现

• 与Transformer-XL相比，PAR Transformer将自注意力模块减少63%，在WikiText-103上实现35%更低的推理延迟，且困惑度完全一致。
• 最优架构仅在前三分之二层使用自注意力，且前馈网络与自注意力模块的最优比例为5:1。
• 在WikiText-103数据集上，PAR Transformer Base与Transformer-XL Base的测试困惑度相同，但推理延迟降低至0.65倍。
• PAR设计规则可泛化至其他数据集：PAR Transformer在enwiki8与text8上实现相当性能，同时计算量更低。
• PAR BERT Base在预训练损失以及MRPC/SST-2微调任务上与BERT Base性能相当，SQuAD v1.1上仅出现1%准确率下降，尽管其自注意力模块数量仅为一半。
• PAR BERT在推理延迟上与DistilBERT相当，但参数量更多、层数更多，且训练范式更简单，表明其在无需蒸馏的情况下实现了架构层面的高效性。