驯服指数函数：面向整数原生边缘推理的快速 Softmax 替代函数
机器学习
2026-04-03 v1 硬件体系结构
摘要
Softmax 在 Transformer 模型的多头注意力（MHA）模块中可能成为计算瓶颈，特别是在低精度推理下的小型模型中，其中指数运算和归一化产生显著开销。因此，我们建议使用头校准截断线性 Softmax（HCCS），这是一种有界、单调的指数 softmax 函数替代函数，使用对最大中心化注意力 logit 的截断线性映射。该近似产生稳定的概率分布，保持原始 logit 的排序，并具有非负值。HCCS 与之前的 softmax 替代函数不同之处在于，它包含一组轻量级校准参数，这些参数基于代表性数据集离线优化，并针对每个单独的注意力头进行校准，以保留各头的统计特性。我们描述了 HCCS 面向高吞吐量场景的硬件导向实现，目标是 AMD Versal AI 引擎。AMD 该平台的当前参考实现依赖于 bfloat16 算术或查找表（LUT）来执行指数运算，这可能限制平台的吞吐量，且无法利用 AI 引擎的高吞吐量整数向量处理单元。相比之下，HCCS 提供了到 AI 引擎的 int8 乘加（MAC）单元的自然映射。据我们所知，这是第一个针对 AMD AI 引擎的 int8 优化 softmax 替代函数，在量化感知重训练后，在小型或重度量化 MHA 工作负载上显著超越其他参考实现的速度性能，同时保持具有竞争力的任务精度。
引用
@article{arxiv.2604.02292,
title = {Taming the Exponential: A Fast Softmax Surrogate for Integer-Native Edge Inference},
author = {Dimitrios Danopoulos and Enrico Lupi and Michael Kagan and Maurizio Pierini},
journal= {arXiv preprint arXiv:2604.02292},
year = {2026}
}