[论文解读] Efficient Softmax Approximation for Deep Neural Networks with Attention Mechanism

Ihor Vasyltsov, Wooseok Chang|arXiv (Cornell University)|Nov 21, 2021

Advanced Neural Network Applications参考文献 33被引用 6

一句话总结

本文提出两种基于查找表（LUT）的方法，用于在基于注意力机制的深度神经网络中实现高效的Softmax近似，从而在无需昂贵除法器的情况下实现硬件加速。该方法在多种自然语言处理和视觉任务中，采用8位量化时实现低于1%的精度损失，展现出强大的泛化能力与极低的内存开销（700字节）。

ABSTRACT

There has been a rapid advance of custom hardware (HW) for accelerating the inference speed of deep neural networks (DNNs). Previously, the softmax layer was not a main concern of DNN accelerating HW, because its portion is relatively small in multi-layer perceptron or convolutional neural networks. However, as the attention mechanisms are widely used in various modern DNNs, a cost-efficient implementation of softmax layer is becoming very important. In this paper, we propose two methods to approximate softmax computation, which are based on the usage of LookUp Tables (LUTs). The required size of LUT is quite small (about 700 Bytes) because ranges of numerators and denominators of softmax are stable if normalization is applied to the input. We have validated the proposed technique over different AI tasks (object detection, machine translation, sentiment analysis, and semantic equivalence) and DNN models (DETR, Transformer, BERT) by a variety of benchmarks (COCO17, WMT14, WMT17, GLUE). We showed that 8-bit approximation allows to obtain acceptable accuracy loss below $1.0\%$.

研究动机与目标

• 解决基于注意力机制的模型（如Transformer和BERT）中Softmax层日益增长的计算成本，尤其是在资源受限的边缘环境中的问题。
• 克服在边缘加速器中将Softmax计算卸载到主机CPU所导致的效率低下问题，该问题会阻碍低延迟自然语言处理工作负载的性能。
• 开发一种硬件友好的Softmax近似方法，避免复杂的除法运算，从而在设备端推理中降低功耗与延迟。
• 在多种人工智能任务和模型上验证该方法，以确保其在不同序列长度和模型架构下的泛化能力与鲁棒性。

提出的方法

• 通过预计算的查找表（LUT）对Softmax函数进行分段常数近似，以替代直接计算。
• 利用输入归一化稳定分子与分母的取值范围，从而在多种模型中实现紧凑的LUT大小（约700字节）。
• 实现两种变体：一种使用单个LUT完成整个Softmax计算，另一种为分子和分母分别使用独立的LUT以提升精度。
• 将基于LUT的Softmax集成到硬件加速器（如NPU）中，消除对高面积除法器单元的需求，从而降低功耗与延迟。
• 在训练后量化（PTQ-D）阶段应用该方法，无需微调，确保与现有量化工作流的兼容性。
• 优化LUT大小与分辨率（如256–512字节），在精度与内存占用之间取得平衡，尤其适用于激活值动态范围较大的模型（如DETR+DC5）。

实验结果

研究问题

• RQ1 紧凑的LUT-based Softmax近似是否能在边缘AI加速器中实现高精度，同时最小化硬件复杂度？
• RQ2 所提出方法在多种基于注意力机制的模型（如BERT、DETR、Transformer）和任务（自然语言处理、目标检测）中表现如何？
• RQ3 LUT大小与量化精度（如8位）对具有高幅值Softmax输入的模型中精度损失的影响如何？
• RQ4 与标准的PTQ-D量化相比，基于LUT的近似在精度退化与泛化能力方面表现如何？
• RQ5 该方法在激活分布不均衡的模型（如具有膨胀卷积层的DETR）中泛化程度如何？

主要发现

• 在所有评估任务（目标检测、机器翻译、情感分析、语义等效性）中，采用8位量化时，所提出的LUT-based Softmax近似方法的精度损失均低于1.0%。
• 在机器翻译（WMT14/WMT17）任务中，BLEU分数与全精度基线的偏差低于0.5%，表明性能退化可忽略不计。
• 在情感分析与语义等效性（GLUE基准）任务中，LUT方法导致的精度下降始终小于仅使用PTQ-D量化带来的下降。
• 对于具有膨胀卷积层的DETR模型（DETR+DC5），当LUT大小从256字节增加到512字节时，8位精度下的精度损失从9%降至3%，显示出良好的可扩展性。
• 该方法在多种模型上泛化良好：DETR、BERT与Transformer均表现出一致的性能，精度损失极小，证实了其鲁棒性。
• LUT大小保持较小（约700字节），使其在内存受限的边缘设备上部署时，不会带来显著的面积或功耗开销。