[论文解读] Energon: Towards Efficient Acceleration of Transformers Using Dynamic Sparse Attention

Zhe Zhou, Junlin Liu+2 more|arXiv (Cornell University)|Oct 18, 2021

Advanced Neural Network Applications被引用 7

一句话总结

Energon 提出了一种协同设计框架，通过混合精度多轮过滤（MP-MRF）算法实现动态稀疏注意力，从而加速 Transformer 模型。该算法在运行时使用低精度操作识别关键的查询-键对，并仅对选定的对进行高精度计算。该方法在 NLP 和 CV 工作负载上相较 CPU 和 GPU 基线实现了最高 168× 的加速和 10⁴× 的能效提升。

ABSTRACT

In recent years, transformer models have revolutionized Natural Language Processing (NLP) and shown promising performance on Computer Vision (CV) tasks. Despite their effectiveness, transformers' attention operations are hard to accelerate due to the complicated data movement and quadratic computational complexity, prohibiting the real-time inference on resource-constrained edge-computing platforms. To tackle this challenge, we propose Energon, an algorithm-architecture co-design approach that accelerates various transformers using dynamic sparse attention. With the observation that attention results only depend on a few important query-key pairs, we propose a Mix-Precision Multi-Round Filtering (MP-MRF) algorithm to dynamically identify such pairs at runtime. We adopt low bitwidth in each filtering round and only use high-precision tensors in the attention stage to reduce overall complexity. By this means, we significantly mitigate the computational cost with negligible accuracy loss. To enable such an algorithm with lower latency and better energy efficiency, we also propose an Energon co-processor architecture. Elaborated pipelines and specialized optimizations jointly boost the performance and reduce power consumption. Extensive experiments on both NLP and CV benchmarks demonstrate that Energon achieves $168 imes$ and $8.7 imes$ geo-mean speedup and up to $10^4 imes$ and $10^3 imes$ energy reduction compared with Intel Xeon 5220 CPU and NVIDIA V100 GPU. Compared to state-of-the-art attention accelerators SpAtten and $A^3$, Energon also achieves $1.7 imes, 1.25 imes$ speedup and $1.6 imes, 1.5 imes $ higher energy efficiency.

研究动机与目标

• 解决 Transformer 模型中自注意力机制带来的高计算与内存开销，尤其在资源受限的边缘平台上的问题。
• 克服先前稀疏注意力方法（如 SpAtten 和 A³）的局限性，包括固定剪枝策略、高带外内存访问开销，或需要模型微调重训练的问题。
• 实现高效、动态且硬件感知的注意力计算，在显著降低复杂度的同时保持高精度。
• 设计一种协同处理器架构，通过优化的流水线和内存访问模式，高效加速 MP-MRF 算法。
• 在多样化的 NLP 和 CV 基准测试中验证该方法的有效性，包括云环境和边缘场景下的真实部署。

提出的方法

• 提出混合精度多轮过滤（MP-MRF）算法，通过低比特宽运算在运行时动态识别关键查询-键对。
• 在多轮中使用低精度（如 4 位）过滤，逐步缩小关键注意力对的搜索空间，再对最终筛选出的对执行高精度注意力计算。
• 仅对过滤后的关键对使用高精度（如 16 位）张量执行最终注意力计算，从而最小化整体计算开销。
• 设计 Energon 协同处理器，配备专用流水线和内存优化，以加速 MP-MRF 及稀疏注意力执行。
• 构建性能模型以指导硬件配置，并分析不同工作负载和核心数量下的性能瓶颈。
• 将 Energon 加速器集成至边缘和服务器平台（Energon-edge 和 Energon-server），评估端到端性能与可扩展性。

实验结果

研究问题

• RQ1 动态、运行时识别的稀疏注意力是否能在不造成显著精度损失的前提下降低 Transformer 模型的计算复杂度？
• RQ2 混合精度、多轮过滤策略是否能以极低的精度开销有效识别关键查询-键对？
• RQ3 Energon 协同处理器架构如何优化稀疏注意力工作负载的内存访问与流水线效率？
• RQ4 在多样化的 NLP 和 CV 任务中，Energon 相较于 SOTA 加速器（如 SpAtten 和 A³）在速度和能效方面能实现多大程度的超越？
• RQ5 当处理核心数量增加时，Energon 的可扩展性如何？性能瓶颈出现在何处？

主要发现

• 在云工作负载中，Energon 相较于 Intel Xeon 5220 CPU 实现了 168× 的几何平均加速和 10⁴× 的能效降低；在边缘平台中，相较 ARM-A72 CPU 实现了 440× 加速和 2951× 能耗降低。
• 在 NLP 基准测试中，Energon 相较于 NVIDIA V100 GPU 实现了最高 8.7× 加速和 10³× 能耗降低，在性能与能效方面均优于 SpAtten 和 A³。
• MP-MRF 算法将计算量减少 4× 至 8×，且精度损失可忽略不计，实现了无需模型微调的动态剪枝。
• 配备 Energon 的系统因注意力计算被卸载并实现流水线处理，端到端延迟降低 1.21×，吞吐量提升 1.55×。
• 可扩展性分析表明，核心数从 4 增加到 16 时，端到端延迟显著降低，但超过 16 个核心后性能增益趋于平缓，主要受限于内存带宽。
• 与 SpAtten 相比，Energon 在 Task-A、C 和 D 上分别实现了 16×、8× 和 8× 的剪枝率，而 SpAtten 仅实现 2× 剪枝率且无需微调，展现出更优的稀疏性与精度权衡。