[论文解读] SpAtten: Efficient Sparse Attention Architecture with Cascade Token and Head Pruning

Hanrui Wang, Zhekai Zhang+1 more|arXiv (Cornell University)|Dec 17, 2020

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

一句话总结

SpAtten 提出了一种软硬件协同设计的稀疏注意力架构，通过级联的标记与头剪枝，以及渐进式量化，显著减少了注意力机制中的内存访问和计算量。通过动态地实时剪除不重要的标记与头，并根据置信度自适应调整位宽，SpAtten 在不损失准确率的前提下，实现了相较于 TITAN Xp GPU 的 4059 倍能效提升和 162 倍加速。

ABSTRACT

The attention mechanism is becoming increasingly popular in Natural Language Processing (NLP) applications, showing superior performance than convolutional and recurrent architectures. However, attention becomes the compution bottleneck because of its quadratic computational complexity to input length, complicated data movement and low arithmetic intensity. Moreover, existing NN accelerators mainly focus on optimizing convolutional or recurrent models, and cannot efficiently support attention. In this paper, we present SpAtten, an efficient algorithm-architecture co-design that leverages token sparsity, head sparsity, and quantization opportunities to reduce the attention computation and memory access. Inspired by the high redundancy of human languages, we propose the novel cascade token pruning to prune away unimportant tokens in the sentence. We also propose cascade head pruning to remove unessential heads. Cascade pruning is fundamentally different from weight pruning since there is no trainable weight in the attention mechanism, and the pruned tokens and heads are selected on the fly. To efficiently support them on hardware, we design a novel top-k engine to rank token and head importance scores with high throughput. Furthermore, we propose progressive quantization that first fetches MSBs only and performs the computation; if the confidence is low, it fetches LSBs and recomputes the attention outputs, trading computation for memory reduction. Extensive experiments on 30 benchmarks show that, on average, SpAtten reduces DRAM access by 10.0x with no accuracy loss, and achieves 1.6x, 3.0x, 162x, 347x speedup, and 1,4x, 3.2x, 1193x, 4059x energy savings over A3 accelerator, MNNFast accelerator, TITAN Xp GPU, Xeon CPU, respectively.

研究动机与目标

• 解决通用平台（如 CPU 和 GPU）上注意力机制效率低下的问题，这些平台面临高内存带宽和低算术强度的挑战。
• 克服现有加速器仅针对 CNN 和 RNN 优化但不适用于注意力机制的局限性。
• 通过利用人类语言中的结构冗余性以及注意力头的冗余性，减少注意力模型中的 DRAM 访问与计算量。
• 设计一种领域专用加速器，支持推理时动态、输入相关的标记与头剪枝，实现实时效率提升。
• 通过最小化内存与能耗开销，实现 BERT 和 GPT-2 等大型 NLP 模型在资源受限设备上的高效部署。

提出的方法

• 提出级联标记剪枝，基于跨层累积注意力概率移除不重要的标记，剪枝比率根据句子长度自适应调整。
• 引入级联头剪枝，通过跨层累积头幅度得分对注意力头进行排序并移除冗余头。
• 设计一种高吞吐、并行的 top-k 引擎，以 O(n) 时间复杂度实时排序标记与头的重要性得分，支持实时剪枝。
• 实现渐进式量化：首先仅用最高有效位（MSBs）计算；若置信度较低，则获取最低有效位（LSBs）并重新计算，以计算开销换取内存减少。
• 构建专用硬件架构，包含流水线数据通路、优化的内存层次结构，以及对动态稀疏性与可变位宽的支持，将理论节能效果实际化。
• 软硬件协同设计，使剪枝与量化为输入相关，并在推理时应用，确保系统适应性强且准确率损失最小。

实验结果

研究问题

• RQ1 动态、输入相关的标记与头剪枝是否能在不损失模型准确率的前提下，减少注意力计算与内存访问？
• RQ2 硬件加速器如何高效支持使用低延迟 top-k 引擎实时排序与选择重要标记与头？
• RQ3 渐进式量化在保持推理置信度的前提下，能在多大程度上减少内存带宽与能耗？
• RQ4 与 A³ 和 MNNFast 等先进加速器相比，级联剪枝与渐进式量化在性能与能效方面表现如何？
• RQ5 所提出的协同设计方法是否可泛化应用于判别式（如 BERT）与生成式（如 GPT-2）Transformer 模型？

主要发现

• SpAtten 在 30 个基准测试中平均将 DRAM 访问减少 10.0 倍，且无准确率损失。
• 与 A³ 和 MNNFast 加速器相比，SpAtten 分别实现 1.6 倍与 3.0 倍的加速，同时实现 1.4 倍与 3.2 倍的能效提升。
• 在 TITAN Xp GPU 上，SpAtten 实现 162 倍加速与 1193 倍能效节省。
• 在 Xeon CPU 上，SpAtten 实现 347 倍加速与 4059 倍能效节省。
• top-k 引擎实现了高效、高吞吐的标记与头重要性排序，时间复杂度为 O(n)，对实时剪枝至关重要。
• 渐进式量化支持动态位宽自适应，通过迭代优化，在减少内存带宽的同时保持推理置信度。