专栏名称: PaperWeekly

PaperWeekly是一个分享知识和交流学问的学术组织，关注的领域是自然语言处理的各个方向。我们热爱知识，分享知识，希望通过我们大家的努力为自然语言处理的发展做出一点点贡献。我们每周会分享一期特定话题的论文笔记和本周值得读的相关论文。

Softmax后传：寻找Top-K的光滑近似

PaperWeekly · 公众号 · 科研 · 2024-09-20 22:43

正文

©PaperWeekly 原创 · 作者 | 苏剑林

单位 | 科学空间

研究方向 | NLP、神经网络

Softmax，顾名思义是“soft 的 max”，是算子（准确来说是）的光滑近似，它通过指数归一化将任意向量转化为分量非负且和为1的新向量，并允许我们通过温度参数来调节它与（的 one hot 形式）的近似程度。除了指数归一化外，我们此前在《通向概率分布之路：盘点Softmax及其替代品》也介绍过其他一些能实现相同效果的方案。

我们知道，最大值通常又称 Top-1，它的光滑近似方案看起来已经相当成熟，那读者有没有思考过，一般的。Top-k 的光滑近似又是怎么样的呢？下面让我们一起来探讨一下这个问题。

问题描述

设向量，简单起见我们假设它们两两不相等，即。记为最大的 k 个分量的下标集合，即以及

。我们定义 Top-k 算子的映射：

说白了，如果属于最大的 k 个元素之一，那么对应的位置变成 1，否则变成 0，最终结果是一个 Multi-Hot 向量，比如。

从到实际上是一种硬指派（Hard Assignment）运算，它本质上是不连续的，没有保留关于的有效梯度，因此无法集成到模型中进行端到端训练。为了解决这个问题，我们就需要构造一个能够提供有效梯度信息的的光滑近似——在有些文献中也称为“可微 Top-k 算子（Differentiable Top-k Operator）”。

具体来说，我们定义集合

那么我们要做的事情就是构建的一个映射，并尽量满足如下性质：

可以验证作为的 Softmax 是满足如上性质的，所以提出上述性质实际就是希望构建出来的能够成为 Softmax 的自然推广。当然，构建 Top-k 的光滑近似本身就比 Top-1 的要难，所以如果遇到困难的话，不必要严格遵循以上性质，只要能表明所构建的映射确实具备的光滑近似的特性就行。

迭代构造

事实上，笔者很早之前就关注过该问题，首次讨论于 2019 年的文章《函数光滑化杂谈：不可导函数的可导逼近》中，当时称之为，并给出过一个迭代构造方案：

输入为，初始化为全 0 向量；

执行 )（保证所有元素非负）

对于，执行：

返回。

其实这个迭代的构造思路很简单，我们可以先将替换为来理解，此时算法流程就是先保证分量非负，然后识别出 Top-1，再将 Top-1 置零（最大值变成了最小值），接着识别出剩下的 Top-1，依此类推，最终的正好就是。而作为的光滑近似，在迭代过程中使用自然也就得到的光滑近似了。

无独有偶，笔者发现在 Stack Exchange 上的提问《Is there something like softmax but for top k values?》[1] 中有回复者提出过一个相同思路的方案，它先定义了加权 Softmax：

然后构建的迭代过程为：

输入为，初始化为全 0 向量；

对于，执行：

返回。

这跟笔者所提的迭代过程在思路上是完全一样的，只不过笔者把乘到了上而它则乘到了上，借助本身的非负性简化了流程。然而，这个迭代实际上是错误的，它不符合“趋近性”，比如当 k=2 时，代入取的极限并不是 Multi-Hot 向量，而是最大值变为 1.5、次大值变为 0.5、其余都变为 0 的向量，这是因为跟大致上是同阶的，乘到上并不能完全消除最大值。

作为梯度

迭代构造全凭经验，可能会隐藏一些难以发现的问题，比如看起来更简单的加权 Softmax 迭代实际上不合理。由于没有更贴近本质的原理指导，这些方案往往也难以理论分析，比如笔者构造的迭代，虽然测起来没有问题，但我们很难证明分量都在 [0,1] 内，也很难判断它是否满足单调性。

所以，我们希望得到一个更高观点的原理去指导和设计这个光滑近似。就在前几天，笔者突然意识到一个关键的事实

也就是说，最大的 k 个分量的和，它的梯度正好是。所以我们似乎可以改为找的光滑近似，然后求梯度就得到的光滑近似了，前者是一个标量，它的光滑近似找起来更容易一些，比如利用恒等式

即遍历所有 k 个分量之和取最大值，这样一来，问题变成了找 max 的光滑近似，这个我们早已解决（参考《寻求一个光滑的最大值函数》[4]），答案是：

对它求梯度，我们就得到的一个形式：

分母是所有 k 分量和的指数和，分子则是所有包含的 k 分量和的指数和。根据这个形式，我们可以轻松证明

所以这样定义的确实是属于的。事实上，我们还可以证明它同时满足单调性、不变性和单调性，并且它就是 Softmax，这些特性显示它确实是 Softmax 对于 Top-k 算子的自然推广，我们暂且称之为“GradTopK（Gradient-guided Soft Top-k operator）”。

不过还没到庆祝时刻，因为式（10）的数值计算问题还没解决。如果直接按照式（10）来计算的话，分母就涉及到项指数求和，计算量非常可观，所以必须找到一个高效的计算方法。我们已经分别记了分子、分母为，可以观察到分子满足递归式

结合对 i 求和等于的事实，我们可以构建一个递归计算过程：

其中，为了减少溢出风险，我们对两端都取了对数。现在只需要迭代k步就可以完成的计算，效率上是可以接受的。然而，即便做了对数化处理，上述递归也只能对小方差的或者比较小的 k 算一下，反之与最大的就会相当接近，当数值上无法区分时就会出现的 Bug，个人认为这是这种递归转化的根本困难。

一个非常简陋的参考实现：

import numpy as np

def GradTopK(x, k):
    for i in range(1, k + 1):
        logZs = x if i == 1 else x + logZ + np.log(1 - np.exp(logZs - logZ))
        logZ = np.logaddexp.reduce(logZs) - np.log(i)
    return np.exp(logZs - logZ)

k, x = 10, np.random.randn(100)
GradTopK(x, k)

待定常数

上一节通过梯度来构建 Top-k 光滑近似的思路，确实能给人一种高屋建瓴的美感，但可能也会有读者觉得它过于抽象，缺乏一种由表及里的直观感，同时对于大方差的或者比较大的 k 的数值不稳定性，也让我们难以完全满意现有结果。所以，接下来我们将探究一种自下而上的构建思路。

方法大意

该思路来源于 Stack Exchange 上的另一个帖子《Differentiable top-k function》[2] 的回复。设 f(x) 是任意的、光滑的、单调递增的函数，并且满足

。

看上去条件很多，但实际上这种函数构造起来没有什么难度，比如经典的 Sigmoid 函数，还有、等。接着我们考虑

跟我们想要的差多远呢？每个分量都在 [0,1] 内肯定是满足了，但是分量之和等于 k 无法保证，所以我们引入一个跟相关的待定常数来保证这一点：

也就是通过分量之和为 k 来反解出，我们可以称之为“ThreTopK（Threshold-adjusted Soft Top-k operator）”，如果读者已经阅读过《通向概率分布之路：盘点Softmax及其替代品》，就会发现这个做法跟 Sparsemax、Entmax- 是一样的。

ThreTopK 会是我们理想的吗？还真是！首先我们假设了 f 的单调性，所以单调性是满足的，其次

，也就是常数可以收纳到里边，所以不变性也是满足的。

最后，当时，我们可以找到一个适当的阈值，使得最大的 k 个分量趋于，剩下的分量趋于，从而就等于，也就是说满足趋近性。

解析求解

既然已经证明了 ThreTopK 的理论优越性，那么接下来要解决的就是的计算了，这个大部份情况下都只能诉诸数值计算方法，不过对于，我们可以求得一个解析解。

求解的思路跟之前的 Sparsemax 是一样的。不失一般性，假设的分量已经从大到小排好顺序，即，接着假设我们已经知道，那么此时

由此解得

由此我们可以看出，当 k=1 时，m 只能取 0，此时可以发现 ThreTopK 正好就是 Softmax；当 k > 1 时，我们无法事先确定 m 的值，所以只能遍历，根据上式计算，寻找满足。下面同样是一个非常简陋的参考实现：

import numpy as np

def ThreTopK(x, k):
    x_sort = np.sort(x)
    x_lamb = np.logaddexp.accumulate(x_sort)[-k:] - np.log(np.arange(k) + 1)
    x_sort_shift = np.pad(x_sort[-k:][1:], (0, 1), constant_values=np.inf)
    lamb = x_lamb[(x_lamb <= x_sort_shift) & (x_lamb >= x_sort[-k:])]
    return np.clip(np.exp(x - lamb), 0, 1)

k, x = 10, np.random.randn(100)
ThreTopK(x, k)

通用结果

从原理和代码都可以看出，时的 ThreTopK 几乎不会出现数值稳定性问题，并且 k=1 时能退化为 Softmax，这些都是它的优势。然而，实际上也算不上完全光滑（除了当 k=1 时不起作用），它在 x=0 处是不可导的。如果介意这一点，那么我们就需要选择处处可导的 f(x)，比如。

下面我们以为例，此时我们无法求出的解析解，不过由于的单调递增性，所以函数

关于是单调递减的，因此的数值求解并不困难，二分法、牛顿法均可。以二分法为例，不难看出

，这里是的逆函数，从这个初始区间出发，逐步二分到指定精度即可：

import numpy as np

def sigmoid(x):
    y = np.exp(-np.abs(x))
    return np.where(x >= 0, 1, y) / (1 + y)

def sigmoid_inv(x):
    return np.log(x / (1 - x))

def ThreTopK(x, k, epsilon=1e-4):
    low = x.min() - sigmoid_inv(k / len(x))
    high = x.max() - sigmoid_inv(k / len(x))
    while high - low > epsilon:
        lamb = (low + high) / 2
        Z = sigmoid(x - lamb).sum()
        low, high = (low, lamb) if Z else (lamb, high)
    return sigmoid(x - lamb)

k, x = 10, np.random.randn(100)
ThreTopK(x, k)