[论文解读] Fast Transformers with Clustered Attention

Apoorv Vyas, Angelos Katharopoulos+1 more|arXiv (Cornell University)|Jul 9, 2020

Speech Recognition and Synthesis参考文献 28被引用 7

一句话总结

本文提出聚类注意力（clustered attention），一种自注意力机制的线性复杂度近似方法，通过将查询分组为聚类，并仅对聚类中心计算注意力，显著降低计算成本。该方法在自动语音识别任务上达到最先进性能，并能仅用25个聚类近似完整BERT模型，在GLUE和SQuAD基准测试中无性能下降。

ABSTRACT

Transformers have been proven a successful model for a variety of tasks in sequence modeling. However, computing the attention matrix, which is their key component, has quadratic complexity with respect to the sequence length, thus making them prohibitively expensive for large sequences. To address this, we propose clustered attention, which instead of computing the attention for every query, groups queries into clusters and computes attention just for the centroids. To further improve this approximation, we use the computed clusters to identify the keys with the highest attention per query and compute the exact key/query dot products. This results in a model with linear complexity with respect to the sequence length for a fixed number of clusters. We evaluate our approach on two automatic speech recognition datasets and show that our model consistently outperforms vanilla transformers for a given computational budget. Finally, we demonstrate that our model can approximate arbitrarily complex attention distributions with a minimal number of clusters by approximating a pretrained BERT model on GLUE and SQuAD benchmarks with only 25 clusters and no loss in performance.

研究动机与目标

• 为解决Transformer中自注意力机制的二次方计算复杂度问题，该问题限制了其在长序列上的应用。
• 开发一种在大幅降低推理与训练时间的同时保持高性能的模型方法。
• 实现对预训练模型（如BERT）的高效近似，且精度下降最小。
• 证明聚类注意力能够处理真实NLP任务中复杂的稀疏注意力模式。
• 展示该方法具有高效可扩展性，训练时间与二氧化碳排放量最高可减少50%。

提出的方法

• 使用K均值聚类与局部敏感哈希（LSH）将查询分组，以减少注意力计算次数。
• 仅对每个聚类中心计算一次注意力，显著减少点积运算次数。
• 为提升精度，方法会识别并精确计算每个查询-聚类中注意力分数最高的键。
• 该方法在聚类数量固定时，相对于序列长度保持线性复杂度。
• 该方法在训练与推理阶段均可应用，并可用于从预训练模型蒸馏完整注意力机制。
• 推导了理论边界，以分析与完整自注意力机制的近似误差。

实验结果

研究问题

• RQ1 聚类注意力机制是否能在保持序列建模任务性能的前提下实现线性复杂度？
• RQ2 聚类注意力在多大程度上能近似预训练BERT模型中发现的复杂注意力分布？
• RQ3 在与标准Transformer相同的计算预算下，该方法是否能保持性能？
• RQ4 该方法是否能在无显著精度损失的前提下减少训练时间与能耗？
• RQ5 在下游任务中，实现接近完整注意力机制的近似所需最少聚类数是多少？

主要发现

• 在Switchboard语音识别数据集上，i-clustered注意力在每次前向传播约50秒的计算预算下，词错误率（WER）比完整注意力低超过2个百分点。
• 在12层网络中，i-clustered将每轮训练时间减少48%（1.91小时 vs. 3.84小时），总收敛时间减少44%（132.13小时 vs. 228.05小时）。
• 在GLUE与SQuAD基准测试中，仅使用25个聚类的i-clustered在所有任务上均达到与完整BERT相当的性能，仅在SQuAD任务上略逊一筹（F1: 0.876 vs. 0.904）。
• 使用25个聚类的聚类注意力在GLUE任务上实现近乎零性能损失，仅在SQuAD与RTE任务上出现轻微下降，而这两项任务需要复杂的注意力模式。
• 该方法使长序列的GPU训练时间减少50%，直接转化为更低的二氧化碳排放与能耗。
• 理论分析证实，近似误差是受控的，且可通过为每个聚类选择高注意力键来最小化。