通向概率分布之路：盘点Softmax及其替代品

2024-06-27 06:06

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 科学空间

研究方向 | NLP、神经网络

不论是在基础的分类任务中，还是如今无处不在的注意力机制中，概率分布的构建都是一个关键步骤。具体来说，就是将一个维的任意向量，转换为一个元的离散型概率分布。众所周知，这个问题的标准答案是 Softmax，它是指数归一化的形式，相对来说比较简单直观，同时也伴有很多优良性质，从而成为大部分场景下的“标配”。

尽管如此，Softmax 在某些场景下也有一些不如人意之处，比如不够稀疏、无法绝对等于零等，因此很多替代品也应运而生。在这篇文章中，我们将简单总结一下 Softmax 的相关性质，并盘点和对比一下它的部分替代方案。

Softmax回顾

首先引入一些通用记号：是需要转为概率分布的 n 维向量，它的分量可正可负，也没有限定的上下界。定义为全体元离散概率分布的集合，即

之所以标注而不是，是因为约束定义了维空间中的一个维子平面，再加上的约束，的集合就只是该平面的一个子集，即实际维度只有。

基于这些记号，本文的主题就可以简单表示为探讨的映射，其中我们习惯称之为 Logits 或者 Scores。

基本定义

Softmax 的定义很简单：

Softmax 的来源和诠释都太多了，比如能量模型、统计力学或者单纯作为的光滑近似等，所以我们很难考证它的最早出处，也不去做这个尝试了。很多时候我们也会加上一个温度参数，即考虑，但 tau 本身也可以整合到的定义之中，因此这里不特别分离出参数。

Softmax 的分母我们通常记为，它的对数就是大多数深度学习框架都自带的运算，他它是的一个光滑近似：

当取1时，就可以写出

，方差越大近似程度越高，更进一步的讨论可以参考《寻求一个光滑的最大值函数》[1]。

两点性质

除了将任意向量转换为概率分布外，Softmax 还满足两点性质

单调性意味着 Softmax 是保序的，的最大值/最小值跟的最大值/最小值相对应；不变性说的是的每个分量都加上同一个常数，Softmax 的结果不变，这跟的性质是一样的，即同样有

。

因此，根据这两点性质我们可以得出，Softmax实际是一个光滑近似（更准确来说是的光滑近似），更具体地我们有

这大概就是 Softmax 这个名字的来源。注意不要混淆了，Softmax 是而不是的光滑近似，的光滑近似是才对。

梯度计算

对于深度学习来说，了解一个函数的性质重要方式之一是了解它的梯度，对于 Softmax，我们在《从梯度最大化看Attention的Scale操作》曾经算过：

这样排列成的矩阵也称为 Softmax 的雅可比矩阵（Jacobian Matrix）[2]，它的 L1 范数有一个简单的形式

当是 one hot 分布时，上式等于 0，这意味着 Softmax 的结果越接近 one hot，它的梯度消失现象越严重，所以至少初始化阶段，我们不能将 Softmax 初始化得接近 one hot。同时上式最右端也联系到了 Rényi 熵的概念，它跟常见的香侬熵类似。

参考实现

Softmax 的直接实现很简单，直接取然后归一化就行，Numpy 的参考代码为：

1def softmax(x):
2    y = np.exp(x)
3    return y / y.sum()

然而，如果中存在较大的分量，那么算时很容易溢出，因此我们通常都要利用 Softmax 的不变性，先将每个分量减去所有分量的最大值，然后再算 Softmax，这样每个取的分量都不大于 0，确保不会溢出：

1def softmax(x):
2    y = np.exp(x - x.max())
3    return y / y.sum()

损失函数

构建概率分布的主要用途之一是用于构建单标签多分类任务的输出，即假设有一个分类任务，是模型的输出，那么我们希望通过来预测每个类的概率。为了训练这个模型，我们需要一个损失函数，假设目标类别是，常见的选择是交叉熵损失：

我们可以求得它的梯度：

注意是给定的，所以实际表达的是目标分布，而全体就是本身，所以上式可以更直观地写成：

也就是说，它的梯度正好是目标分布与预测分布之差，只要两者不相等，那么梯度就一直存在，优化就可以持续下去，这是交叉熵的优点。当然，某些情况下这也是缺点，因为 Softmax 只有在才会得到 one hot，换言之正常情况下都不会出现 one hot，即优化一直不会完全停止，那么就有可能导致过度优化，这也是后面的一些替代品的动机。

除了交叉熵之外，还有其他一些损失可用，比如，这可以理解为准确率的光滑近似的相反数，但它可能会有梯度消失问题，所以它的优化效率往往不如交叉熵，一般只适用于微调而不是从零训练，更多讨论可以参考《如何训练你的准确率？》。

Softmax变体

介绍完 Softmax，我们紧接着总结一下本博客以往讨论过 Softmax 的相关变体工作，比如 Margin Softmax、Taylor Softmax、Sparse Softmax 等，它们都是在 Softmax 基础上的衍生品，侧重于不同方面的改进，比如损失函数、、稀疏性、长尾性等。

Margin Softmax

首先我们介绍起源于人脸识别的一系列 Softmax 变体，它们可以统称为 Margin Softmax，后来也被应用到 NLP 的 Sentence Embedding 训练之中，本站曾在《基于GRU和am-softmax的句子相似度模型》[3] 讨论过其中的一个变体 AM-Softmax，后来则在《从三角不等式到Margin Softmax》有过更一般的讨论。

尽管 Margin Softmax 被冠以 Softmax 之名，但它实际上更多是一种损失函数改进。以 AM-Softmax 为例，它有两个特点：第一，以形式构造 Logits，即

的形式，此时的温度参数是必须的，因为单纯的值域为，不能拉开类概率之间的差异；第二，它并不是简单地以为损失，而是做了加强：

直观来看，就是交叉熵希望是所有分量中最大的一个，而 AM-Softmax 则不仅希望最大，还希望它至少比第二大的分量多出，这里的就称为 Margin。

为什么要增加对目标类的要求呢？这是应用场景导致的。刚才说了，Margin Softmax 起源于人脸识别，放到 NLP 中则可以用于语义检索，也就是说它的应用场景是检索，但训练方式是分类。如果单纯用分类任务的交叉熵来训练模型，模型编码出来的特征不一定能很好地满足检索要求，所以要加上 Margin 使得特征更加紧凑一些。更具体的讨论请参考《从三角不等式到Margin Softmax》一文，或者查阅相关论文。

Taylor Softmax

接下来要介绍的，是在《exp(x)在x=0处的偶次泰勒展开式总是正的》讨论过的 Taylor Softmax，它利用了的泰勒展开式的一个有趣性质：

对于任意实数及偶数，总有，即在处的偶次泰勒展开式总是正的。

利用这个恒正性，我们可以构建一个 Softmax 变体（是任意偶数）：

由于是基于的泰勒展开式构建的，所以在一定范围内 Taylor Softmax 与 Softmax 有一定的近似关于，某些场景下可以用 Taylor Softmax 替换 Softmax。那么 Taylor Softmax 有什么特点呢？答案是更加长尾，因为 Taylor Softmax 是多项式函数归一化，相比指数函数衰减得更慢，所以对于尾部的类别，Taylor Softmax 往往能够给其分配更高的概率，可能有助于缓解 Softmax 的过度自信现象。

Taylor Softmax 的最新应用，是用来替换 Attention 中的 Softmax，使得原本的平方复杂度降低为线性复杂度，相关理论推导可以参考《Transformer升级之路：作为无限维的线性Attention》[4]。

该思路的最新实践是一个名为 Based 的模型，它利用来线性化 Attention，声称比 Attention 高效且比 Mamba 效果更好，详细介绍可以参考博客《Zoology（Blogpost 2）: Simple, Input-Dependent, and Sub-Quadratic Sequence Mixers》[5] 和《BASED: Simple linear attention language models balance the recall-throughput tradeoff》[6]。

Sparse Softmax

Sparse Softmax 是笔者参加 2020 年法研杯时提出的一个简单的 Softmax 稀疏变体，首先发表在《SPACES：“抽取-生成”式长文本摘要（法研杯总结）》，后来也补充了相关实验，写了篇简单的论文《Sparse-softmax: A Simpler and Faster Alternative Softmax Transformation》[7]。

我们知道，在文本生成中，我们常用确定性的 Beam Search 解码，或者随机性的 TopK/TopP Sampling 采样，这些算法的特点都是只保留了预测概率最大的若干个 Token 进行遍历或者采样，也就等价于将剩余的 Token 概率视为零，而训练时如果直接使用 Softmax 来构建概率分布的话，那么就不存在绝对等于零的可能，这就让训练和预测出现了不一致性。Sparse Softmax 就是希望能处理这种不一致性，思路很简单，就是在训练的时候也把 Top- 以外的 Token 概率置零：

其中是将从大到小排列后前个元素的原始下标集合。简单来说，就是在训练阶段就进行与预测阶段一致的阶段操作。这里的选取方式也可以按照 Nucleus Sampling 的 Top- 方式来操作，看具体需求而定。但要注意的是，Sparse Softmax 强行截断了剩余部分的概率，意味着这部分 Logits 无法进行反向传播了，因此 Sparse Softmax 的训练效率是不如 Softmax 的，所以它一般只适用于微调场景，而不适用于从零训练。

Perturb Max

这一节我们介绍一种新的构建概率分布的方式，这里称之为 Perturb Max，它是 Gumbel Max 的一般化，首次介绍于《从重参数的角度看离散概率分布的构建》，此外在论文《EXACT: How to Train Your Accuracy》[8] 也有过相关讨论，至于更早的出处笔者则没有进一步考究了。

问题反思

首先我们知道，构建一个的映射并不是难事，只要是（实数到非负实数）的映射，如，那么只需要让

就是满足条件的映射了。如果要加上“两点性质” [9] 中的“单调性”呢？那么也不难，只需要保证的单调递增函数，这样的函数也有很多，比如。但如果再加上“不变性”呢？我们还能随便写出一个满足不变性的映射吗？（反正我不能）

可能有读者疑问：为什么非要保持单调性和不变性呢？的确，单纯从拟合概率分布的角度来看，这两点似乎都没什么必要，反正都是“力大砖飞”，只要模型足够大，那么没啥不能拟合的。但从 “Softmax 替代品”这个角度来看，我们希望新定义的概率分布同样能作为的光滑近似，那么就要尽可能多保持跟相同的性质，这是我们希望保持单调性和不变性的主要原因。

Gumbel Max

Perturb Max 借助于Gumbel Max 的一般化来构造这样的一类分布。不熟悉 Gumbel Max 的读者，可以先到《漫谈重参数：从正态分布到Gumbel Softmax》[10] 了解一下 Gumbel Max。简单来说，Gumbel Max 就是发现：

怎么理解这个结果呢？首先，这里的是指的每个分量都是从 Gumbel 分布独立 [11] 重复采样出来的；接着，我们知道给定向量，本来是确定的结果，但加了随机噪声之后，的结果也带有随机性了，于是每个都有自己的概率；最后，Gumbel Max 告诉我们，如果加的是 Gumbel 噪声，那么i的出现概率正好是。

Gumbel Max 最直接的作用，就是提供了一种从分布中采样的方式，当然如果单纯采样还有更简单的方法，没必要“杀鸡用牛刀”。Gumbel Max 最大的价值是“重参数（Reparameterization）”，它将问题的随机性从带参数的离散分布转移到了不带参数的上，再结合 Softmax 是的光滑近似，我们得到是 Gumbel Max 的光滑近似，这便是 Gumbel Softmax，是训练“离散采样模块中带有可学参数”的模型的常用技巧。

一般噪声

Perturb Max 直接源自 Gumbel Max：既然 Softmax 可以从 Gumbel 分布中导出，那么如果将 Gumbel 分布换为一般的分布，比如正态分布，不就可以导出新的概率分布形式了？也就是说直接定义

重复 Gumbel Max 的推导，我们可以得到

其中是的累积概率函数。对于一般的分布，哪怕是简单的标准正态分布，上式都很难得出解析解，所以只能数值估计。为了得到确定性的计算结果，我们可以用逆累积概率函数的方式进行均匀采样，即先从均匀选取，然后通过求解来得到。

从 Perturb Max 的定义或者最后的形式我们都可以断言 Perturb Max 满足单调性和不变性，这里就不详细展开了。那它在什么场景下有独特作用呢？说实话，还真不知道，《EXACT: How to Train Your Accuracy》[8] 一文用它来构建新的概率分布并优化准确率的光滑近似，但笔者自己的实验显示没有特别的效果。个人感觉，可能在某些需要重参数的场景能够表现出特殊的作用吧。

Sparsemax

接下来要登场的是名为 Sparsemax 的概率映射，出自 2016 年的论文《From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification》[12]，它跟笔者提出的 Sparse Softmax 一样，都是面向稀疏性的改动，但作者的动机是用在 Attention 中提供更好的可解释性。跟 Sparse Softmax 直接强行截断 Top- 个分量不同，Sparsemax 提供了一个更为自适应的稀疏型概率分布构造方式。

基本定义

原论文将 Sparsemax 定义为如下优化问题的解：

通过拉格朗日乘数法就可以求出精确解的表达式。然而，这种方式并不直观，而且也不容易揭示它跟 Softmax 的联系。下面提供笔者构思的一种私以为更加简明的引出方式。

首先，我们可以发现，Softmax 可以等价地表示为

其中是使得的各分量之和为 1 的常数，对于 Softmax 我们可以求出。

然后，在 Taylor Softmax 那一节我们说了，的偶次泰勒展开总是正的，因此可以用偶次泰勒展开来构建 Softmax 变体。但如果是奇数次呢？比如，它并不总是非负的，但我们可以加个强行让它变成非负的，即，用这个近似替换掉式（20）的，就得到了 Sparsemax：

其中依然是使得的各分量之和为1的常数，并且常数 1 也可以整合到之中，所以上式也等价于

求解算法

到目前为止，Sparsemax 还只是一个形式化的定义，因为的具体计算方法尚不清楚，这就是本节需要探讨的问题。不过即便如此，单靠这个定义我们也不难看出 Sparsemax 满足单调性和不变性两点性质，如果还觉得不大肯定的读者，可以自行尝试证明一下它。

现在我们转向的计算。不失一般性，我们假设各分量已经从大到小排序好，即，接着我们不妨先假设已知，那么很显然

根据的定义，我们有

这就可以求出。当然，我们无法事先知道，但我们可以遍历，利用上式求一遍，取满足

，这也可以等价地表示为求满足的最大的，然后返回对应的

参考实现：

1def sparsemax(x):
2    x_sort = np.sort(x)[::-1]
3    x_lamb = (np.cumsum(x_sort) - 1) / np.arange(1, len(x) + 1)
4    lamb = x_lamb[(x_sort >= x_lamb).argmin() - 1]
5    return np.maximum(x - lamb, 0)

梯度计算

方便起见，我们引入记号

那么可以写出

从这个等价形式可以看出，跟 Sparse Softmax 一样，Sparsemax 同样也只对部分类别有梯度，可以直接算出雅可比矩阵：

由此可以看出，对于在里边的类别，Sparsemax 倒是不会梯度消失，因为此时它的梯度恒为常数，但它总的梯度大小，取决于的元素个数，它越少则越稀疏，意味着梯度也越稀疏。

损失函数

最后我们来讨论 Sparsemax 作为分类输出时的损失函数。比较直观的想法就是跟 Softmax 一样用交叉熵，但 Sparsemax 的输出可能是严格等于0的，所以为了防止错误，还要给每个分量都加上，最终的交叉熵形式为，但这一来有点丑，二来它还不是凸函数，所以并不是一个理想选择。

事实上，交叉熵在 Softmax 中之所以好用，是因为它的梯度恰好有（12）的形式，所以对于 Sparsemax，我们不妨同样假设损失函数的梯度为，然后反推出损失函数该有的样子，即：

从右往左验证比较简单，从左往右推可能会有些困难，但不多，反复拼凑一下应该就能出来了。第一个常数是为了保证损失函数的非负性，我们可以取一个极端来验证一下：假设优化到完美，那么应该也是 one hot，此时，并且，于是

所以要多加上常数。

Entmax-α

Entmax- 是 Sparsemax 的一般化，它的动机是 Sparsemax 往往会过度稀疏，这可能会导致学习效率偏低，导致最终效果下降的问题，所以 Entmax- 引入了参数，提供了 Softmax（）到Sparsemax（）的平滑过度。

Entmax- 出自论文《Sparse Sequence-to-Sequence Models》[13]，作者跟 Sparsemax 一样是 Andre F. T. Martins，这位大佬围绕着稀疏 Softmax、稀疏 Attention 做了不少工作，有兴趣的读者可以在他的主页 [14] 查阅相关工作。

基本定义

跟 Sparsemax 一样，原论文将 Entmax- 定义为类似（19）的优化问题的解，但这个定义涉及到 Tsallis entropy [15] 的概念（也是 Entmax 的 Ent 的来源），求解还需要用到拉格朗日乘数法，相对来说比较复杂，这里不采用这种引入方式。

我们的介绍同样是基于上一节的近似，对于 Softmax 和 Sparsemax，我们有

Sparsemax 太稀疏，背后的原因也可以理解为近似精度不够高，我们可以从中演化出更高精度的近似

只要，那么最右端就是一个比更好的近似（想想为什么）。利用这个新近似，我们就可以构建

这里是为了对齐原论文的表达方式，事实上用表示更简洁一些。同样地，常数 1 也可以收入到定义之中，所以最终定义可以简化为

求解算法

对于一般的，求解是比较麻烦的事情，通常只能用二分法求解。

首先我们记，并且不失一般性假设，然后我们可以发现 Entmax- 是满足单调性和不变性的，借助不变性我们可以不失一般性地设（如果不是，每个都减去即可）。

现在可以检验，当时，的所有分量之和大于等于 1，当时，的所有分量之和等于 0，所以最终能使分量之和等于 1 的必然在内，然后我们就可以使用二分法来逐步逼近最优的。

对于某些特殊的，我们可以得到一个求精确解的算法。Sparsemax 对应，我们前面已经给出了求解过程，另外一个能给解析解的例子是，这也是原论文主要关心的例子，如果不加标注，那么 Entmax 默认就是 Entmax-1.5。跟 Sparsemax 一样的思路，我们先假设已知，于是有

这只不过是关于的一元二次方程，可以解得

当我们无法事先知道时，可以遍历，利用上式求一遍，取满足那一个，但注意这时候不等价于求满足的最大的。

完整的参考实现：

1def entmat(x):
2    x_sort = np.sort(x / 2)[::-1]
3    k = np.arange(1, len(x) + 1)
4    x_mu = np.cumsum(x_sort) / k
5    x_sigma2 = np.cumsum(x_sort**2) / k  - x_mu**2
6    x_lamb = x_mu - np.sqrt(np.maximum(1. / k - x_sigma2, 0))
7    x_sort_shift = np.pad(x_sort[1:], (0, 1), constant_values=-np.inf)
8    lamb = x_lamb[(x_sort > x_lamb) & (x_lamb > x_sort_shift)]
9    return np.maximum(x / 2 - lamb, 0)**2

其他内容

Entmax- 的梯度跟 Sparsemax 大同小异，这里就不展开讨论了，读者自行推导一下或者参考原论文就行。至于损失函数，同样从梯度出发反推出损失函数也是可以的，但其形式有点复杂，有兴趣了解的读者可以参考原论文《Sparse Sequence-to-Sequence Models》[13] 以及《Learning with Fenchel-Young Losses》[16]。

不过就笔者看来，直接用算子来定义损失函数更为简单通用，可以避免求原函数的复杂过程：