图片速览 FlashAttention + Online softmax论文翻译 & Online normalizer calculation for softmax

FakeOccupational

已于 2025-03-28 15:53:32 修改

阅读量1.9k

点赞数 27

分类专栏：深度学习文章标签：笔记

于 2025-03-08 19:00:00 首次发布

本文链接：https://blog.csdn.net/ResumeProject/article/details/146052311

版权

深度学习专栏收录该内容

175 篇文章

订阅专栏

在这里插入图片描述

摘要

Softmax 函数在机器学习中无处不在，多个先前的研究提出了更快的替代方案。在本文中，我们提出了一种计算经典 Softmax 的方法，该方法可减少内存访问次数，并假设这种内存访问的减少应能提高 Softmax 在实际硬件上的性能。基准测试结果证实了这一假设：Softmax 的加速比最高可达 1.3 倍，而 Softmax 与 TopK 结合并融合后的加速比最高可达 5 倍。

1 Introduction

神经网络模型被广泛应用于语言建模任务，如机器翻译 [1] 和语音识别 [2]。这些模型在计算单词概率时，会考虑序列中已生成的部分。通常，概率由投影层（Projection layer）计算，该层将隐藏表示“投影”到输出词汇空间，随后通过 Softmax 函数将原始 logits 转换为概率向量。Softmax 不仅用于神经网络，例如，在多项逻辑回归（multinomial logistic regression）[3] 中也被应用。
许多先前的研究提出了计算单词概率的更快替代方案。Differentiated Softmax [4] 和 SVD-Softmax [5] 用计算效率更高的替代方法取代了投影层（通常只是矩阵乘法）。多种层次化 Softmax（Hierarchical Softmax）[6, 7, 8] 将单一的“投影+Softmax”组合拆分为多个更小的版本，并组织成树状结构。基于采样的近似方法，如重要性采样（Importance Sampling）[9]、噪声对比估计（Noise Contrastive Estimation, NCE）[10] 和 Blackout [11]，通过仅对原始向量的部分元素运行 Softmax 来加速训练。最后，自归一化 Softmax（Self-Normalized Softmax）[12] 通过优化目标函数，使 Softmax 归一化项接近 1，从而在推理时可以跳过计算该项。
这并不是一个详尽的列表，但希望它具有代表性。几乎所有方法仍然需要运行原始的 Softmax 函数，无论是对完整向量还是缩减后的向量进行计算。有两个例外无需计算 Softmax 归一化项：使用噪声对比估计（Noise Contrastive Estimation, NCE）进行训练，以及使用自归一化 Softmax（Self-Normalized Softmax）进行推理。除此之外，所有其他方法都将受益于原始 Softmax 的加速计算。
据我们所知，尚无针对提升原始 Softmax 函数性能的专门研究。我们试图弥补这一不足，并找到了一种减少内存访问次数来计算 Softmax 的方法。我们对其进行了基准测试，以验证内存访问的减少是否能够真正转化为实际硬件上的性能提升。

2 Original softmax

Function $y = S o f t ma x (x)$ 定义如下 :

$y_i = \frac{e^{x_i}}{ \sum_{ j=1}^Ve^{x_j}}$

其中 $\in \mathbb{R}^V$ 。
朴素实现方式（见算法 1）需要对输入向量进行两次扫描：一次用于计算归一化项(分母) $d_V$ ，另一次用于计算输出值 $y_i$ 。这实际上意味着每个向量元素需要进行三次内存访问：两次读取（load）和一次存储（store）。

在这里插入图片描述

不幸的是，在实际硬件上，由于表示数值的范围有限，算法 1 的第 3 行可能会因指数运算而导致溢出（overflow）或下溢（underflow）。有一种更安全的形式（公式 (1) 的变体），可以避免这一问题：

在这里插入图片描述

由于指数函数的性质，这种变换不会改变 Softmax 的归一化性质，最终的概率分布仍然相同：
- 防止溢出：如果 $x_i$ 的值非常大，直接计算 $e^{x_i}$ 可能会溢出，而 $e^{x_i - m}$ 会得到一个较小的值，保持数值稳定。
- 防止下溢：如果 $x_i$ 的值较小，直接计算 $e^{x_i}$ 可能会接近 0，影响计算精度，而 $e^{x_i - m}$ 能够保持有效的数值范围。
- 理论上可以减去任意常数 $c$ ，但减去最大值 $\max(x)$ 是最优选择，因为它能确保指数项的最大值始终为 1（即 $e^{m - m} = e^0 = 1$ ），避免数值不稳定。
所有主流深度学习框架在 Softmax 计算中都使用了这种安全版本，例如 TensorFlow [13] v1.7、PyTorch [14]（包含 Caffe2）v0.4.0、MXNet [15] v1.1.0、Microsoft Cognitive Toolkit [16] v2.5.1 以及 Chainer [17] v5.0.0a1。然而，安全 Softmax 需要对输入向量进行三次遍历：第一次计算最大值 $m_V$ ，第二次计算归一化项 (分母) $d_V$ ，第三次计算最终输出值 $y_i$ （见算法 2）。这导致每个向量元素总共需要进行 4 次内存访问。我们希望对此进行优化。

在这里插入图片描述

3 Onlinenormalizer calculation

算法 3 在一次遍历输入向量的过程中，同时计算最大值 $m$ 和归一化项 $d$ ，并且仅增加每个向量元素两个操作的额外计算开销。这种方法将 Softmax 函数计算中的内存访问次数从每个向量元素 4 次减少到 3 次。这个方法的灵感来自于数值稳定的方差计算在线算法，具体见 [18]。

在这里插入图片描述

本质上，算法在遍历输入数组的元素时，保持最大值 $m$ 和归一化项 $d$ 。在每次迭代时，它需要先根据新的最大值 $m_j$ 调整归一化项 $d$ ，然后再将新值加入归一化项。这种方法确保了最大值和归一化项在每次迭代时都保持正确更新，从而避免了额外的内存访问并提高了效率。
作者同时计算要归一化的最大值和归一化用的分母值，其通过数学归纳法的方式逐步进行：
- 首先是单个数据的情况：
- 然后推广到多个值：
关键公式为：
$d_j\leftarrow d_{j-1}\times e^{m_{j-1}-m_j}+e^{x_j-m_j}$
算法 3 被证明能够计算公式 (2) 中定义的 Softmax 函数。同时，它也是安全的：
$m_j$ 是运行中的最大值，且$m_j \in \mathbb{R} )。对于每个 ( j \in {1, V} )，( m_j ) 不会发生下溢或溢出。
$d_j$ 也有界： $\leq d_j \leq j$ ，对于每个 $\in \{1, V\}$ ，这一点可以通过数学归纳法轻松证明。对于 $d_j$ 使用 32 位浮点数存储，可以保证处理高达 $1.7 \times 10^{37}$ 个元素的向量 $x$ 而不发生溢出。这是一个相当大的数值，但如果向量更大，则需要使用 64 位浮点数来存储 $d_j$ 。

算法 2 提供了相同的保证： $\leq d_j \leq j$ ，对于每个 $\in \{1, V\}$ 。

在本文的其余部分，我们将称算法 3 为“在线 Softmax（Online Softmax）”。

3.1 Parallel online normalizer calculation

算法 3 的第 1-6 行定义了一种顺序计算归一化项的方法，通过一次遍历输入向量来完成。现代计算设备允许多个线程并行运行；为了充分利用设备的计算能力，我们需要一个并行版本的算法。为此，我们定义了一个广义版本的在线归一化器计算方法：
这种并行化方法将算法的计算过程分解成多个并行线程，每个线程负责计算输入向量的一部分，最后将结果合并。通过这种方式，可以显著提高计算效率，特别是在多核处理器或 GPU 上执行时。在并行版本中，我们将确保每个线程正确地更新归一化项 $d_j$ 和最大值 $m_j$ ，并且保持数值稳定性与顺序版本一致。
将公式 (3) 从左到右顺序应用，相当于执行算法 3 中的第 1-6 行。运算符 $\oplus$ 是结合律（计算的顺序不影响最终结果）的，这使得可以并行地评估公式 (3)。它也是交换律(可以灵活地调整数据分配策略)的，这为使并行实现更加高效提供了灵活性。为了简洁起见，我们省略了这两个性质的证明。

4 Softmax andtop-k fusion

在线 Softmax（算法 3）对每个向量元素进行三次内存访问：一次加载用于计算归一化项，一次加载和一次存储用于计算 Softmax 函数值 $y_i$ 。在自回归模型的推理过程中，结合 TopK 后进行搜索，而 TopK 不需要计算所有的 $y_i$ 值。这使得可以获得更大的性能提升。
TopK 函数生成一个包含 K 个整数索引的向量，这些索引指向输入向量中最大的 K 个值，同时也返回这些值：

$\text{TopK}(y) = (v, z) : v_i = y_{z_i}, \quad v_i \geq y_j, \, \forall i \in [1, K], \, \forall j \notin z$

其中， $\in \mathbb{R}^V$ ， $\in \mathbb{Z}^K$ ， $\in \mathbb{R}^K$ 。

TopK 至少需要加载输入向量的每个元素一次。单独运行安全 Softmax 和 TopK 会要求每个输入元素进行 5 次访问；而如果使用在线 Softmax 替代安全 Softmax（但仍然分别运行），每个元素则需要 4 次访问。

5 Benchmarking

在线归一化器计算减少了 Softmax 和 Softmax+TopK 函数的内存访问次数。Softmax 函数的每字节浮点操作数（flops per byte）非常低，这意味着即使是带有额外少量浮点运算的在线 Softmax，内存带宽也应该成为性能的瓶颈。因此，减少内存访问次数应该能转化为性能的提升，实验结果也证实了这一点。
我们使用 CUDA C 实现了一个针对 GPU 的基准测试。该基准测试利用了 CUB v1.8.0 进行快速并行归约。所有实验都在 NVIDIA Tesla V100 PCIe 16 GB 上运行，启用了 ECC 和持久模式，CUDA Toolkit 版本为 9.1。基准测试的源代码可在 github.com/NVIDIA/online-softmax 获取。

在这里插入图片描述