[论文解读] Transformer Acceleration with Dynamic Sparse Attention

Liu Liu, Zheng Qu+3 more|arXiv (Cornell University)|Oct 21, 2021

Advanced Neural Network Applications参考文献 26被引用 12

一句话总结

本文提出动态稀疏注意力（DSA），一种通过动态识别并利用输入相关的稀疏注意力模式来降低Transformer计算成本的方法，在不损失模型准确率的前提下实现高效推理。通过在近似注意力得分上使用轻量级预测路径，DSA实现了高达95%的稀疏度和4.35倍的理论计算量减少，仅带来1.33%的额外开销，在V100 GPU上实现1.94倍的注意力输出加速和14.6倍的softmax计算加速，稀疏度达90%时表现优异。

ABSTRACT

Transformers are the mainstream of NLP applications and are becoming increasingly popular in other domains such as Computer Vision. Despite the improvements in model quality, the enormous computation costs make Transformers difficult at deployment, especially when the sequence length is large in emerging applications. Processing attention mechanism as the essential component of Transformer is the bottleneck of execution due to the quadratic complexity. Prior art explores sparse patterns in attention to support long sequence modeling, but those pieces of work are on static or fixed patterns. We demonstrate that the sparse patterns are dynamic, depending on input sequences. Thus, we propose the Dynamic Sparse Attention (DSA) that can efficiently exploit the dynamic sparsity in the attention of Transformers. Compared with other methods, our approach can achieve better trade-offs between accuracy and model complexity. Moving forward, we identify challenges and provide solutions to implement DSA on existing hardware (GPUs) and specialized hardware in order to achieve practical speedup and efficiency improvements for Transformer execution.

研究动机与目标

• 为解决Transformer中自注意力机制在长序列长度下的高计算成本问题。
• 克服静态稀疏注意力模式无法捕捉输入相关动态稀疏性的局限。
• 开发一种实用且硬件高效的算法，在保持完整注意力表示能力的同时减少计算量。
• 通过动态稀疏性和数据流优化，实现大规模Transformer在现有及专用硬件上的高效部署。

提出的方法

• 在标准自注意力机制中引入轻量级预测路径，以识别注意力得分中的动态稀疏模式。
• 使用低成本的近似注意力得分作为稀疏度预测器的输入，以降低计算开销。
• 将预测的稀疏性表示为可学习的二值掩码，以跳过softmax和输出计算中的冗余操作。
• 将细粒度动态稀疏性扩展至块状和向量级稀疏性等结构化模式，以支持硬件优化。
• 设计定制化内核，并利用数据重用机会，提升GPU和专用加速器上的内存访问效率。
• 集成低精度推理用于稀疏性预测，并通过计算重排实现token级并行化。

实验结果

研究问题

• RQ1 能否有效识别并利用自注意力中的动态稀疏模式，在不损失模型准确率的前提下减少计算量？
• RQ2 与静态或固定稀疏模式相比，动态稀疏性在准确率与效率之间的权衡表现如何？
• RQ3 DSA在现有GPU架构和专用硬件加速器上的实际性能提升如何？
• RQ4 算法稀疏性如何与硬件协同设计，以最大化加速比和能效？

主要发现

• DSA在LRA基准测试中实现了高达95%的注意力权重稀疏度，且未造成可测量的准确率损失。
• 稀疏性预测仅引入1.17%至1.33%的额外计算开销，实现高达4.35倍的理论计算量减少。
• 在NVIDIA V100 GPU上，DSA在90%稀疏度下实现注意力得分计算1.15倍加速、softmax计算14.6倍加速，以及注意力输出计算1.94倍加速。
• 结合向量级稀疏性和优化内核设计，DSA在专用硬件上实现总内存访问量最高减少2.54倍。
• 该方法通过保留关键注意力权重，维持了完整注意力的表达能力，区别于低秩近似或激进剪枝方法。
• 硬件专用化通过低精度预测、细粒度稀疏计算和改进的数据局部性，进一步提升了性能。