Vasu Shyam 在 2024 年 8 月 12 日的一篇新闻中提到，log-sum-exp 的梯度计算在 GPU 上比传统的 softmax（或 softargmax）更快。log-sum-exp 函数是一种平滑的最大值近似函数，常用于机器学习算法中。其梯度计算的加速意味着在处理大规模数据和复杂模型时，可以提高计算效率。
log-sum-exp 函数在语言模型中的应用尤为重要，特别是在基于 KV 缓存的解码过程中。具体来说，能量函数可以表示成 log-sum-exp 和 max 运算操作，由于这两者都满足结合律，可以按任意顺序进行，而不会影响最终结果。在此基础上，研究团队设计了新的并行化算法，先在各 GPU 上并行计算局部能量函数，再通过树状的 Allreduce 汇总各处结果，最后用自动微分取梯度，即可得到注意力的输出。整个过程仅需与计算能量函数相同的时间开销，而显存占用也几乎没有额外负担。
这种并行化算法的设计利用了 GPU 集群的两级拓扑特点 —— 即同节点内使用高速 NVLink，而节点间则依赖 IB 或以太网等。相比之下，传统的环形注意力难以将通信与计算很好地重叠，最终会被最慢的互联带宽所制约。
在实际应用中，这种新的并行化算法提高了大模型的推理速度。例如，某团队在实验中发现，使用这种方法后，推理速度提高了 4 倍，且上下文长度达到了 256K。
softmax 函数最初出现在统计学中，作为一种基于能级预测状态分布的方法。随着时间的推移，softmax 被引入到经济学和机器学习领域，成为生成概率分布的关键工具。在 Transformer 架构中，softmax 用于注意力机制，生成 token 的概率分布。然而，softmax 的计算成本较高，尤其是在处理长序列时，其指数计算和对序列长度求和的操作使得并行化变得困难。
为了应对这些挑战，研究人员提出了多种替代方案。例如，Google DeepMind 提出使用 ReLU 替代 softmax。在视觉 Transformer 中，使用 ReLU 除以序列长度的注意力可以接近或匹敌传统的 softmax 注意力。实验结果显示，ReLU 注意力在序列长度维度上实现了更好的并行化，其所需的 gather 操作比 softmax 注意力更少。
此外，Evan Miller 提出了一种名为 Softmax1 的改进版本。Softmax1 在负象限中有一个 “逃出口”，使得其导数总是正的，从而避免了异常值反馈循环的问题。Miller 的实验结果显示，Softmax_1 在处理异常值时表现更好，且其输出向量中的相对值保持不变。