With the rise of Transformer models in NLP and CV domain, Multi-Head Attention has been proven to be a game-changer. However, its expensive computation poses challenges to the model throughput and efficiency, especially for the long sequence tasks. Exploiting the sparsity in attention has been proven to be an effective way to reduce computation. Nevertheless, prior works do not consider the various distributions among different heads and lack a systematic method to determine the threshold. To address these challenges, we propose Low-Precision Approximate Attention with Head-wise Trainable Threshold for Efficient Transformer (LATTE). LATTE employs a headwise threshold-based filter with the low-precision dot product and computation reuse mechanism to reduce the computation of MHA. Moreover, the trainable threshold is introduced to provide a systematic method for adjusting the thresholds and enable end-to-end optimization. Experimental results indicate LATTE can smoothly adapt to both NLP and CV tasks, offering significant computation savings with only a minor compromise in performance. Also, the trainable threshold is shown to be essential for the leverage between the performance and the computation. As a result, LATTE filters up to 85.16% keys with only a 0.87% accuracy drop in the CV task and 89.91% keys with a 0.86 perplexity increase in the NLP task.

翻译：随着Transformer模型在自然语言处理和计算机视觉领域的兴起，多头注意力机制已被证明具有变革性意义。然而，其高昂的计算成本给模型吞吐量和效率带来挑战，尤其对于长序列任务。利用注意力稀疏性已被证明是降低计算量的有效方法，但现有工作未考虑不同注意力头之间的分布差异，且缺乏系统性的阈值确定方法。针对这些挑战，我们提出面向高效Transformer的头部可训练阈值低精度近似注意力机制（LATTE）。LATTE采用基于头部阈值的滤波器，结合低精度点积计算和计算复用机制以减少多头注意力计算量。此外，引入可训练阈值提供系统化阈值调整方法，实现端到端优化。实验结果表明，LATTE能无缝适配NLP和CV任务，在仅轻微牺牲性能的情况下显著节省计算量。同时，可训练阈值被证明是平衡性能与计算量关系的关键要素。最终，LATTE在CV任务中滤除高达85.16%的关键词且准确率仅下降0.87%，在NLP任务中滤除89.91%的关键词且困惑度仅增加0.86。