[论文解读] Online normalizer calculation for softmax

Maxim Milakov, Natalia Gimelshein|arXiv (Cornell University)|May 8, 2018

Parallel Computing and Optimization Techniques参考文献 8被引用 20

一句话总结

本文提出了一种用于 Softmax 函数的在线归一化计算方法，通过在单次遍历中同时计算最大值和归一化项，将每个元素的内存访问次数从四次减少到三次，从而实现更快的推理速度。在 Tesla V100 上的基准测试显示，仅 Softmax 情况下最高可实现 1.3 倍的加速，而融合 Softmax+TopK 时最高可达 5 倍加速，性能提升主要源于内存带宽压力的降低。

ABSTRACT

The Softmax function is ubiquitous in machine learning, multiple previous works suggested faster alternatives for it. In this paper we propose a way to compute classical Softmax with fewer memory accesses and hypothesize that this reduction in memory accesses should improve Softmax performance on actual hardware. The benchmarks confirm this hypothesis: Softmax accelerates by up to 1.3x and Softmax+TopK combined and fused by up to 5x.

研究动机与目标

• 减少 Softmax 计算中的内存访问开销，因为这是深度学习推理中的性能瓶颈。
• 解决尽管存在众多替代方案，但对经典 Softmax 函数缺乏针对性优化的问题。
• 通过将归一化逻辑与选择逻辑合并，实现 Softmax 和 TopK 操作的高效融合。
• 通过单次遍历算法，在保持数值稳定性的同时减少内存带宽使用。
• 在现代硬件上展示可测量的性能提升，特别是 GPU 架构。

提出的方法

• 提出一种单次遍历算法，在一次迭代中同时计算最大值和 Softmax 归一化项，将每个元素的内存访问次数从四次减少到三次。
• 通过在指数运算前从所有 logits 中减去最大值，采用数值稳定的公式，以防止溢出/下溢。
• 使用增量更新规则计算归一化项：$ d_j = d_{j-1} \cdot e^{m_{j-1} - m_j} + e^{x_j - m_j} $，其中 $ m_j $ 为运行最大值。
• 同时维护运行最大值 $ m_j $ 和归一化和 $ d_j $，在处理每个元素时逐步更新。
• 通过在相同遍历过程中跟踪 top-k 值，将 Softmax 与 TopK 操作融合，消除冗余的内存访问。
• 通过最小化内存带宽使用并支持核函数融合，优化 GPU 性能。

实验结果

研究问题

• RQ1 减少 Softmax 计算中的内存访问次数是否能在现代硬件上带来可测量的性能提升？
• RQ2 是否可能在不牺牲数值稳定性的前提下，通过单次遍历同时计算 Softmax 归一化项和最大值？
• RQ3 与标准的两遍或三遍实现相比，所提出的在线归一化方法在性能和精度方面表现如何？
• RQ4 Softmax 和 TopK 操作的融合能在多大程度上提升端到端推理速度？
• RQ5 在 GPU 架构上，性能提升是否随向量大小和批处理大小而扩展？

主要发现

• 所提出的在线归一化方法将每个元素的内存访问次数从四次减少到三次，在 Tesla V100 上对大向量尺寸的 Softmax 单独操作最高实现 1.3 倍加速。
• 当与 TopK 融合时，组合的 Softmax+TopK 操作最高实现 5 倍加速，其中 2.5 倍来自融合，2 倍来自在线归一化。
• 性能提升在大批次场景中最为显著，此时内存带宽成为性能瓶颈。
• 即使在小批次设置下，由于延迟和内存访问开销降低，该在线 Softmax 仍可实现 1.5 倍至 2.5 倍的加速。
• 该方法保持了数值稳定性，并与现有深度学习框架兼容，可作为即插即用的优化方案。
• 性能增益与其他 Softmax 优化技术（如分层 Softmax、SVD-Softmax 和重要性采样）正交，可进一步结合以实现更高速度提升。