提出新型整数友好softmax替代,加速边缘设备推理。
Taming the Exponential: A Fast Softmax Surrogate for Integer-Native Edge Inference
- 用截断线性映射替代指数运算,实现无指数的软最大值近似。
- 在小模型低精度推理中速度超越现有方案,保持任务准确率。
- 专为AMD Versal AI引擎设计,适配整数向量计算单元,适合边缘部署。
Transformer模型中的多头注意力(MHA)模块,尤其在小型模型和低精度推理场景下,软最大值(Softmax)因指数运算与归一化带来显著计算开销。为此,本文提出头校准截断线性软最大值(HCCS),一种有界且单调的指数软最大值替代函数,通过最大值中心化的注意力得分进行截断线性映射。该方法生成稳定概率分布,保持原始得分排序并保证非负值。与以往替代方法不同,HCCS引入轻量化校准参数,基于代表性数据集离线优化,并针对每个注意力头独立校准,以保留各头的统计特性。本文还针对AMD Versal AI引擎设计了硬件驱动的HCCS实现。当前平台参考实现依赖bfloat16或查找表完成指数操作,限制吞吐量且未充分利用其高吞吐整数向量处理单元。而HCCS天然适配AI引擎的int8乘累加(MAC)单元,是首个为AMD AI引擎优化的int8软最大值替代方案,在量化感知重训练后,于小型或重度量化MHA负载上,性能显著优于其他参考实现,同时保持竞争性任务精度。
原文摘要 · Abstract (English)
Softmax can become a computational bottleneck in the Transformer model's Multi-Head Attention (MHA) block, particularly in small models under low-precision inference, where exponentiation and normalization incur significant overhead. As such, we suggest using Head-Calibrated Clipped-Linear Softmax (HCCS), a bounded, monotone surrogate to the exponential softmax function, which uses a clipped linear mapping of the max centered attention logits. This approximation produces a stable probability distribution, maintains the ordering of the original logits and has non-negative values. HCCS differs from previous softmax surrogates as it includes a set of lightweight calibration parameters that are optimized offline based on a representative dataset and calibrated for each individual attention head to preserve the statistical properties of the individual heads. We describe a hardware-motivated implementation of HCCS for high-throughput scenarios targeting the AMD Versal AI Engines. The current reference implementations from AMD for this platform rely upon either bfloat16 arithmetic or LUTs to perform the exponential operation, which might limit the throughput of the platform and fail to utilize the high-throughput integer vector processing units of the AI Engine. In contrast, HCCS provides a natural mapping to the AI Engines' int8 multiply accumulate (MAC) units. To the best of our knowledge, this is the first int8 optimized softmax surrogate for AMD AI engines that significantly exceeds the speed performance of other reference implementations while maintaining competitive task accuracy on small or heavily quantized MHA workloads after quantization-aware retraining.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。