[论文解读] Neural Architecture Search on Efficient Transformers and Beyond

Zexiang Liu, Dong Li+5 more|arXiv (Cornell University)|Jul 28, 2022

Advanced Neural Network Applications被引用 11

一句话总结

该论文提出了一种神经架构搜索（NAS）框架，用于自动发现高效Transformer模型的最优架构，特别针对基于核函数的cosFormer。通过引入结合Softmax与线性注意力的混合注意力搜索空间，该方法在显著提升计算效率的同时，实现了与标准Transformer相当的性能，其准确率优于纯高效模型，且FLOPS更低。

ABSTRACT

Recently, numerous efficient Transformers have been proposed to reduce the quadratic computational complexity of standard Transformers caused by the Softmax attention. However, most of them simply swap Softmax with an efficient attention mechanism without considering the customized architectures specially for the efficient attention. In this paper, we argue that the handcrafted vanilla Transformer architectures for Softmax attention may not be suitable for efficient Transformers. To address this issue, we propose a new framework to find optimal architectures for efficient Transformers with the neural architecture search (NAS) technique. The proposed method is validated on popular machine translation and image classification tasks. We observe that the optimal architecture of the efficient Transformer has the reduced computation compared with that of the standard Transformer, but the general accuracy is less comparable. It indicates that the Softmax attention and efficient attention have their own distinctions but neither of them can simultaneously balance the accuracy and efficiency well. This motivates us to mix the two types of attention to reduce the performance imbalance. Besides the search spaces that commonly used in existing NAS Transformer approaches, we propose a new search space that allows the NAS algorithm to automatically search the attention variants along with architectures. Extensive experiments on WMT' 14 En-De and CIFAR-10 demonstrate that our searched architecture maintains comparable accuracy to the standard Transformer with notably improved computational efficiency.

研究动机与目标

• 为解决手工地设计的高效Transformer架构在性能上不如标准Transformer架构的问题，这些架构通常是从标准Transformer设计中改编而来。
• 探究NAS是否能够发现专为高效注意力机制（如线性注意力）量身定制的模型架构。
• 探索在将Softmax注意力替换为线性注意力时，高效Transformer中准确率与效率之间的权衡关系。
• 提出并验证一种新颖的混合注意力机制，该机制在单个模型中结合Softmax与线性注意力，以平衡性能与效率。
• 在机器翻译与图像分类等真实任务上，展示所提NAS框架的有效性。

提出的方法

• 采用RankNAS框架，通过成对排序与硬件感知约束加速神经架构搜索，降低搜索成本。
• 设计一种新型搜索空间，使NAS算法能够联合搜索架构超参数（如头数、前馈网络维度）与每层的注意力类型（Softmax与线性注意力）。
• 以cosFormer作为目标高效Transformer模型，将其Softmax注意力替换为线性注意力，以实现O(N)复杂度的高效计算。
• 在两个基准任务上进行NAS：WMT’14 En-De机器翻译与CIFAR-10图像分类，以评估泛化能力。
• 提出mFormer架构，根据搜索结果动态选择每层使用Softmax或线性注意力，实现混合注意力机制。
• 应用搜索空间剪枝与早停策略，在不牺牲模型质量的前提下提升NAS过程的效率。

实验结果

研究问题

• RQ1 神经架构搜索能否发现比为标准Transformer设计的手工地架构更适合高效Transformer的架构？
• RQ2 在相同搜索空间下，纯高效Transformer（如cosFormer）与标准Transformer在准确率与FLOPS上的表现如何比较？
• RQ3 混合注意力机制（结合Softmax与线性注意力）是否能比单独使用任一类型带来更优的准确率-效率平衡？
• RQ4 在不同任务中，NAS发现的高效Transformer最优模型中出现了哪些架构模式？
• RQ5 所提出的NAS框架能否在显著降低FLOPS的同时，实现与标准Transformer相当甚至更高的准确率，且保持竞争力？

主要发现

• 高效cosFormer模型的最优架构相比标准Transformer，参数量减少18%，FLOPS降低18%，但准确率更低，表明存在性能-效率失衡。
• mFormer模型采用Softmax与线性注意力的混合机制，在BLEU分数上与标准Transformer相比仅低1.51分，同时FLOPS降低23%，参数量减少20%。
• 在CIFAR-10上，mFormer比cosFormer*准确率高出5.19个百分点，FLOPS仅增加1.38G，展现出更优的平衡性。
• 若将mFormer中的所有线性注意力替换为Softmax注意力，准确率略有提升，但FLOPS显著增加，证实了效率与成本之间的权衡。
• 若将Transformer*中的所有Softmax注意力替换为线性注意力，准确率出现大幅下降（尤其在图像分类任务中），表明线性注意力在对准确率敏感的任务中存在局限性。
• mFormer架构具有“浅而薄”的设计特征：25%的层使用线性注意力（如第3、7、11层），其余层使用Softmax注意力，且头数与前馈网络维度呈负相关。