[论文解读] Hardware Accelerator for Multi-Head Attention and Position-Wise Feed-Forward in the Transformer

Siyuan Lu, Meiqi Wang+3 more|arXiv (Cornell University)|Sep 18, 2020

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

一句话总结

本论文提出了一种基于FPGA的硬件加速器，首次针对Transformer模型中计算最密集的两个组件——多头注意力（MHA）和位置前馈网络（FFN）——进行加速。通过采用矩阵分块、可重构的脉动阵列，以及对层归一化和Softmax等非线性函数的优化，该设计在XCVU13P FPGA上以200MHz频率和16.7W功耗下，相较于GPU实现，使MHA获得14.6倍的加速，FFN获得3.4倍的加速。

ABSTRACT

Designing hardware accelerators for deep neural networks (DNNs) has been much desired. Nonetheless, most of these existing accelerators are built for either convolutional neural networks (CNNs) or recurrent neural networks (RNNs). Recently, the Transformer model is replacing the RNN in the natural language processing (NLP) area. However, because of intensive matrix computations and complicated data flow being involved, the hardware design for the Transformer model has never been reported. In this paper, we propose the first hardware accelerator for two key components, i.e., the multi-head attention (MHA) ResBlock and the position-wise feed-forward network (FFN) ResBlock, which are the two most complex layers in the Transformer. Firstly, an efficient method is introduced to partition the huge matrices in the Transformer, allowing the two ResBlocks to share most of the hardware resources. Secondly, the computation flow is well designed to ensure the high hardware utilization of the systolic array, which is the biggest module in our design. Thirdly, complicated nonlinear functions are highly optimized to further reduce the hardware complexity and also the latency of the entire system. Our design is coded using hardware description language (HDL) and evaluated on a Xilinx FPGA. Compared with the implementation on GPU with the same setting, the proposed design demonstrates a speed-up of 14.6x in the MHA ResBlock, and 3.4x in the FFN ResBlock, respectively. Therefore, this work lays a good foundation for building efficient hardware accelerators for multiple Transformer networks.

研究动机与目标

• 为解决Transformer模型缺乏专用硬件加速器的问题，该模型虽主导现代自然语言处理，但具有高计算与内存需求。
• 通过降低延迟与功耗，实现在移动与嵌入式系统中高效推理Transformer模型。
• 设计一种可重构硬件架构，以高效加速两个最复杂的层：MHA ResBlock与FFN ResBlock。
• 通过将大矩阵分块并重用脉动阵列硬件，实现MHA与FFN之间的资源共享。
• 通过提前计算与硬件友好的近似方法，最小化层归一化与缩放掩码Softmax等关键非线性操作的延迟。

提出的方法

• 对MHA与FFN中的大矩阵应用矩阵分块，以实现脉动阵列在两个组件间的共享使用。
• 采用可重构脉动阵列作为核心计算单元，并设计了高效的数据流以最大化硬件利用率。
• 通过定点数运算与log-sum-exp技巧，对缩放掩码Softmax进行优化，以降低复杂度与延迟。
• 通过预计算均值，并利用代数恒等式在更少周期内计算方差，从而加速层归一化：$var(G,i) = E(G,i)^2 - \frac{1}{d_{model}}\sum_{k=1}^{d_{model}} G(i,k)^2$。
• 使用查找表实现层归一化中的$x \hat{~{}} (-0.5)$操作，以减少关键路径延迟。
• 整个设计以HDL实现，并在Xilinx xcvu13p-fhga2104-3-e FPGA上进行综合，采用8位整数量化。

实验结果

研究问题

• RQ1 单一硬件加速器能否高效处理Transformer架构中的多头注意力与位置前馈块？
• RQ2 如何通过矩阵分块与脉动阵列复用，最小化面积与功耗，同时保持高吞吐量？
• RQ3 为降低硬件中Softmax与层归一化等非线性函数的延迟，需要哪些优化？
• RQ4 在Transformer推理方面，FPGA加速相比GPU在延迟与能效方面能实现多大程度的超越？
• RQ5 对注意力与前馈层进行硬件优化实现后，经INT8量化，模型精度是否能保持不变？

主要发现

• 所提出的加速器在MHA ResBlock上相较NVIDIA V100 GPU实现获得14.6倍加速。
• 在相同设置下（batch size = 1，sequence length = 64），FFN ResBlock相较同一GPU基线实现3.4倍加速。
• FPGA设计在200MHz下运行，片上总功耗为16.7W，其中动态功耗13.3W，静态功耗3.4W。
• 通过精心设计的计算流程，实现了脉动阵列在MHA与FFN操作之间的高效重配置，最大化硬件利用率。
• 层归一化模块通过预计算均值并利用代数恒等式在更少周期内计算方差，显著缩短关键路径，降低延迟。
• 采用8位整数量化的Transformer base模型在IWSLT 2016德英翻译任务上保持BLEU得分为23.57，表明精度损失极小。