Softmax's Sequel: In Search of a Smooth Approximation to Top-K
Softmax, as its name suggests, is a "soft" version of max — a smooth approximation of the $\max$ operator (or more precisely, $\text{argmax}$). It converts an arbitrary vector $\boldsymbol{x}\in\mathbb{R}^n$ into a new vector with nonnegative components summing to 1, via exponential normalization, and it lets us tune, through a temperature parameter, how closely this approximates the (one-hot form of) $\text{argmax}$. Besides exponential normalization, we previously introduced other schemes that achieve the same effect in The Road to Probability Distributions: A Survey of Softmax and Its Alternatives].
As we know, the maximum is often called Top-1, and its smooth approximation scheme already looks quite mature. But has the reader ever wondered what a smooth approximation to the general Top-$k$ might look like? Let's explore this question together below.
Problem Statement
Let the vector be $\boldsymbol{x}=(x_1,x_2,\cdots,x_n)\in\mathbb{R}^n$; for simplicity we assume its components are pairwise distinct, i.e. $i\neq j \Leftrightarrow x_i\neq x_j$. Let $\Omega_k(\boldsymbol{x})$ denote the index set of the $k$ largest components of $\boldsymbol{x}$, i.e. $|\Omega_k(\boldsymbol{x})|=k$ and $\forall i\in \Omega_k(\boldsymbol{x}), j \not\in \Omega_k(\boldsymbol{x})\Rightarrow x_i > x_j$. We define the Top-$k$ operator $\mathcal{T}_k$ as the mapping $\mathbb{R}^n\mapsto\{0,1\}^n$:
\begin{equation} [\mathcal{T}_k(\boldsymbol{x})]_i = \left\{\begin{aligned}1,\,\, i\in \Omega_k(\boldsymbol{x}) \\ 0,\,\, i \not\in \Omega_k(\boldsymbol{x})\end{aligned}\right. \end{equation}
In plain terms, if $x_i$ is among the top $k$ elements, the corresponding position becomes 1, otherwise it becomes 0; the final result is a multi-hot vector, e.g. $\mathcal{T}_2([3,2,1,4]) = [1,0,0,1]$. more
The mapping from $\boldsymbol{x}$ to $\mathcal{T}_k(\boldsymbol{x})$ is essentially a hard assignment operation. It is fundamentally discontinuous and carries no useful gradient with respect to $\boldsymbol{x}$, so it cannot be integrated into a model for end-to-end training. To address this, we need to construct a smooth approximation of $\mathcal{T}_k(\boldsymbol{x})$ that provides useful gradient information — in some of the literature this is called the "Differentiable Top-$k$ Operator" (Differentiable Top-$k$).
Specifically, we define the set
\begin{equation}\Delta_k^{n-1} = \left\{\boldsymbol{p}=(p_1,p_2,\cdots,p_n)\left|\, p_1,p_2,\cdots,p_n\in[0,1],\sum_{i=1}^n p_i = k\right.\right\}\end{equation}
Then what we want to do is construct a mapping $\mathcal{ST}_k(\boldsymbol{x})$ from $\mathbb{R}^n\mapsto \Delta_k^{n-1}$, satisfying, as far as possible, the following properties:
\begin{align}&{\color{red}{monotonicity}}:\quad [\mathcal{ST}_k(\boldsymbol{x})]_i \geq [\mathcal{ST}_k(\boldsymbol{x})]_j \,\,\Leftrightarrow\,\, x_i \geq x_j \\[8pt] &{\color{red}{invariance}}:\quad \mathcal{ST}_k(\boldsymbol{x}) = \mathcal{ST}_k(\boldsymbol{x} + c),\,\,\forall c\in\mathbb{R} \\[8pt] &{\color{red}{convergence}}:\quad \lim_{\tau\to 0^+}\mathcal{ST}_k(\boldsymbol{x}/\tau) = \mathcal{T}_k(\boldsymbol{x}) \\ \end{align}
One can verify that Softmax, as $\mathcal{ST}_1(\boldsymbol{x})$, satisfies these properties, so proposing these properties is essentially asking that the constructed $\mathcal{ST}_k(\boldsymbol{x})$ be a natural generalization of Softmax. Of course, constructing a smooth approximation to Top-$k$ is inherently harder than doing so for Top-1, so if we run into difficulties, we need not adhere strictly to the above properties — it suffices to show that the constructed mapping genuinely has the character of a smooth approximation to $\mathcal{T}_k(\boldsymbol{x})$.
Iterative Construction
In fact, I looked into this problem long ago, first discussing it in a 2019 post Miscellany on Function Smoothing: Differentiable Approximations of Non-Differentiable Functions], where I called it $\text{soft-}k\text{-max}$ and gave an iterative construction:
Input $\boldsymbol{x}$, initialize $\boldsymbol{p}^{(0)}$ as an all-zero vector;
Execute $\boldsymbol{x} = \boldsymbol{x} - \min(\boldsymbol{x})$ (to guarantee all elements are nonnegative)
For $i=1,2,\dots,k$, do:
$\boldsymbol{y} = (1 - \boldsymbol{p}^{(i-1)})\otimes\boldsymbol{x}$;
$\boldsymbol{p}^{(i)} = \boldsymbol{p}^{(i-1)} + \text{softmax}(\boldsymbol{y})$
Return $\boldsymbol{p}^{(k)}$.
The idea behind this iterative construction is actually quite simple: we can understand it by first replacing $\text{softmax}(\boldsymbol{y})$ with $\mathcal{T}_1(\boldsymbol{y})$. In that case, the algorithm first ensures all components are nonnegative, then identifies the Top-1, sets the Top-1 to zero (turning the maximum into the minimum), then identifies the Top-1 among what remains, and so on; the final $\boldsymbol{p}_k$ is exactly $\mathcal{T}_k(\boldsymbol{x})$. Since $\text{softmax}(\boldsymbol{y})$ is a smooth approximation of $\mathcal{T}_1(\boldsymbol{y})$, using $\text{softmax}(\boldsymbol{y})$ in the iteration naturally yields a smooth approximation of $\mathcal{T}_k(\boldsymbol{x})$.
Coincidentally, I found that on Stack Exchange, in the question Is there something like softmax but for top k values?], a respondent proposed a scheme with exactly the same idea. It first defines a weighted softmax:
\begin{equation}[\text{softmax}(\boldsymbol{x};\boldsymbol{w})]_i = \frac{w_i e^{x_i}}{\sum\limits_{i=1}^n w_i e^{x_i}}\end{equation}
and then builds the following iterative process:
Input $\boldsymbol{x}$, initialize $\boldsymbol{p}^{(0)}$ as an all-zero vector;
For $i=1,2,\dots,k$, do:
$\boldsymbol{p}^{(i)} = \boldsymbol{p}^{(i-1)} + \text{softmax}(\boldsymbol{x}; 1 - \boldsymbol{p}^{(i-1)})$
Return $\boldsymbol{p}^{(k)}$.
This is exactly the same idea as the iteration I proposed, except that I multiplied $1 - \boldsymbol{p}_{i-1}$ into $\boldsymbol{x}$, while this scheme multiplies it into $e^{\boldsymbol{x}}$, exploiting the nonnegativity of $e^{\boldsymbol{x}}$ itself to simplify the process. However, this iteration is actually incorrect — it fails to satisfy the "convergence" property. For example, when $k=2$, substituting into $\boldsymbol{x}/\tau$ and taking the limit as $\tau\to 0^+$ does not give a multi-hot vector; instead, the largest value becomes 1.5, the second largest becomes 0.5, and the rest become 0. This is because $1-p_{\max}$ is roughly the same order as $e^{-x_{\max}}$, so multiplying $1-p_{\max}$ into $e^{x_{\max}}$ does not fully eliminate the maximum.
As a Gradient
Iterative constructions rely entirely on empirical intuition, which can hide problems that are hard to spot — as we saw, the seemingly simpler weighted-softmax iteration turns out to be unsound. Since there is no principle rooted more deeply in the essence of the problem to guide us, such schemes are also difficult to analyze theoretically. For instance, with the iteration I constructed, although it tests fine in practice, it's hard to prove that all components of $\boldsymbol{p}_k$ lie within $[0,1]$, and equally hard to determine whether monotonicity holds.
So, we would like a more principled, higher-level viewpoint to guide the design of this smooth approximation. Just a few days ago, I suddenly realized a key fact:
\begin{equation}\mathcal{T}_k(\boldsymbol{x}) = \nabla_{\boldsymbol{x}} \sum_{i\in\Omega_k(\boldsymbol{x})} x_i\end{equation}
That is, the sum of the $k$ largest components has a gradient that is exactly $\mathcal{T}_k(\boldsymbol{x})$. So it seems we could instead look for a smooth approximation of $\sum\limits_{i\in\Omega_k(\boldsymbol{x})} x_i$, and then take its gradient to obtain a smooth approximation of $\mathcal{T}_k(\boldsymbol{x})$. The former is a scalar, and it's much easier to find a smooth approximation for it — for example, using the identity
\begin{equation}\sum_{i\in\Omega_k(\boldsymbol{x})} x_i = \max_{i_1 < \cdots < i_k} (x_{i_1} + \cdots + x_{i_k})\end{equation}
i.e., taking the maximum over the sums of all $k$-element subsets. This turns the problem into finding a smooth approximation of $\max$, which we already solved (see In Search of a Smooth Maximum Function]) — the answer being $\text{logsumexp}$:
\begin{equation}\max_{i_1 < \cdots < i_k} (x_{i_1} + \cdots + x_{i_k})\approx \log\sum_{i_1 < \cdots < i_k} e^{x_{i_1} + \cdots + x_{i_k}}\triangleq \log Z_k\end{equation}
Taking its gradient, we obtain a form for $\mathcal{ST}_k(\boldsymbol{x})$:
\begin{equation}[\mathcal{ST}_k(\boldsymbol{x})]_i = \frac{\sum\limits_{i_2 < \cdots < i_k} e^{x_i+x_{i_2} + \cdots + x_{i_k}}}{\sum\limits_{i_1 < \cdots < i_k} e^{x_{i_1} +x_{i_2}+ \cdots + x_{i_k}}}\triangleq \frac{Z_{k,i}}{Z_k}\label{eq:k-max-grad}\end{equation}
The denominator is the exponential sum over all $k$-element subset sums, and the numerator is the exponential sum restricted to those $x_i$-element subsets that include $k$. From this form, we can easily prove
\begin{equation}0 < [\mathcal{ST}_k(\boldsymbol{x})]_i < 1,\quad \sum_{i=1}^n [\mathcal{ST}_k(\boldsymbol{x})]_i = k\end{equation}
so the $\mathcal{ST}_k(\boldsymbol{x})$ defined this way indeed belongs to $\Delta_k^{n-1}$. In fact, we can also show it satisfies monotonicity, invariance, and the convergence property, and that $\mathcal{ST}_1(\boldsymbol{x})$ it reduces exactly to Softmax. These properties show it is indeed a natural generalization of Softmax to the Top-$k$ operator; let's call it "GradTopK" (Gradient-guided Soft Top-k operator).
But it's not yet time to celebrate, because the numerical computation of $\eqref{eq:k-max-grad}$ has not been resolved. If we compute directly from $\eqref{eq:k-max-grad}$, the denominator involves an exponential sum over $C_n^k$ terms, which is a very substantial computation, so we need an efficient computation method. We've denoted the numerator and denominator $Z_{k,i},Z_k$; we can observe that the numerator $Z_{k,i}$ satisfies the recursion
\begin{equation}Z_{k,i} = e^{x_i}(Z_{k-1} - Z_{k-1,i})\end{equation}
Combined with the fact that summing $Z_{k,i}$ over $i$ equals $kZ_k$, we can build a recursive computation procedure:
\begin{equation}\begin{aligned} \log Z_{k,i} =&\, x_i + \log(e^{\log Z_{k-1}} - e^{\log Z_{k-1,i}}) \\ \log Z_k =&\, \left(\log\sum_{i=1}^n e^{\log Z_{k,i}}\right) - \log k \\ \end{aligned}\end{equation}
where $\log Z_{1,i} = x_i$; to reduce the risk of overflow, we take logarithms on both sides. Now we only need $k$ iterations to complete the computation of $\mathcal{ST}_k(\boldsymbol{x})$, which is acceptably efficient. However, even after log-transforming, this recursion only works for $\boldsymbol{x}$ with small variance, or for relatively small $k$; otherwise, $\log Z_{k-1}$ and the largest $\log Z_{k-1,i}$ become quite close, and when they become numerically indistinguishable, the $\log 0$ bug appears — in my view this is the fundamental difficulty of this style of recursive transformation.
A rather rough reference implementation:
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)
Undetermined Constants
The approach in the previous section, constructing a smooth Top-$k$ approximation via the gradient, does have a certain top-down elegance, but some readers may also find it overly abstract, lacking the intuitiveness of moving from surface to essence. Moreover, the numerical instability for $\boldsymbol{x}$ with large variance or relatively large $k$ still leaves us less than fully satisfied with the current result. So next, we will explore a bottom-up approach to construction.
Basic Idea
This approach comes from a reply to another Stack Exchange thread, Differentiable top-k function]. Let $f(x)$ be any smooth, monotonically increasing function on $\mathbb{R}\mapsto [0,1]$ satisfying $\lim\limits_{x\to\infty}f(x) = 1,\lim\limits_{x\to-\infty}f(x) = 0$. This may look like a lot of conditions, but such functions are actually quite easy to construct — for example, the classic sigmoid function $\sigma(x)=1/(1+e^{-x})$, or $\text{clip}(x,0,1)$, $\min(1, e^x)$, and so on. Next consider
\begin{equation}f(\boldsymbol{x}) = [f(x_1),f(x_2),\cdots,f(x_n)]\end{equation}
How far is $f(\boldsymbol{x})$ from the $\mathcal{ST}_k(\boldsymbol{x})$ we actually want? Each component lying in $[0,1]$ is certainly satisfied, but there's no guarantee that the components sum to $k$, so we introduce an undetermined constant $\lambda(x)$ depending on $\boldsymbol{x}$ to enforce this:
\begin{equation}\mathcal{ST}_k(\boldsymbol{x}) \triangleq f(\boldsymbol{x} - \lambda(\boldsymbol{x})),\quad \sum_{i=1}^n f(x_i - \lambda(\boldsymbol{x})) = k\end{equation}
That is, we solve for $\lambda(\boldsymbol{x})$ by requiring the component sum to equal $k$. We can call this "ThreTopK" (Threshold-adjusted Soft Top-k operator); readers who have already read The Road to Probability Distributions: A Survey of Softmax and Its Alternatives] will notice this is the same approach as Sparsemax and Entmax-$\alpha$.
Will ThreTopK turn out to be the ideal $\mathcal{ST}_k(\boldsymbol{x})$ we're after? It really does! First, since we assumed the monotonicity of $f$, monotonicity holds. Second, $f(\boldsymbol{x} - \lambda(\boldsymbol{x}))=f(\boldsymbol{x}+c - (c+\lambda(\boldsymbol{x})))$, meaning the constant can be absorbed into $\lambda(\boldsymbol{x})$, so invariance also holds. Finally, when $\tau\to 0^+$, we can find a suitable threshold $\lambda(\boldsymbol{x}/\tau)$ such that the $k$ largest components of $\boldsymbol{x}/\tau-\lambda(\boldsymbol{x}/\tau)$ tend to $\infty$ while the rest tend to $-\infty$, so $f(\boldsymbol{x}/\tau-\lambda(\boldsymbol{x}/\tau))$ tends to $\mathcal{T}_k(\boldsymbol{x})$ — that is, the convergence property is satisfied.
Analytical Solution
Having established the theoretical merits of ThreTopK, the next task is to compute $\lambda(\boldsymbol{x})$, which for the most part can only be handled numerically — but for $f(x)=\min(1, e^x)$, we can find an analytical solution.
The solving approach is the same as before for Sparsemax. Without loss of generality, assume the components of $\boldsymbol{x}$ are already sorted in decreasing order, i.e. $x_1 > x_2 > \cdots > x_n$, and suppose we already know $x_m \geq \lambda(\boldsymbol{x}) \geq x_{m+1}$. Then
\begin{equation}k = \sum_{i=1}^n \min(1, e^{x_i - \lambda(\boldsymbol{x})}) = m + \sum_{i=m+1}^n e^{x_i - \lambda(\boldsymbol{x})}\end{equation}
from which we solve
\begin{equation}\lambda(\boldsymbol{x})=\log\left(\sum_{i=m+1}^n e^{x_i}\right) - \log(k-m)\end{equation}
From this we can see that when $k=1$, $m$ can only take the value $0$, and in this case ThreTopK reduces exactly to Softmax. When $k > 1$, we cannot determine the value of $m$ in advance, so we must enumerate $m=0,1,\cdots,k-1$, compute $\lambda(\boldsymbol{x})$ via the above formula, and find the $\lambda(\boldsymbol{x})$ satisfying $x_m \geq \lambda(\boldsymbol{x}) \geq x_{m+1}$. Below is again a rather rough reference implementation:
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)
General Result
As is clear both from the theory and the code, ThreTopK with $f(x)=\min(1, e^x)$ rarely encounters numerical stability issues, and it degenerates to Softmax when $k=1$ — both advantages. However, $\min(1, e^x)$ is actually not fully smooth either (except when $k=1$, in which case $\min$ plays no role): it is non-differentiable at $x=0$. If this bothers us, then we need to choose an everywhere-differentiable $f(x)$, such as $\sigma(x)$.
Let's take $f(x)=\sigma(x)$ as an example below. In this case we can't obtain an analytical solution for $\lambda(\boldsymbol{x})$, but since $\sigma(x)$ is monotonically increasing, the function
\begin{equation}F(\lambda)\triangleq \sum_{i=1}^n \sigma(x_i - \lambda)\end{equation}
is monotonically decreasing in $\lambda$, so numerically solving for $F(\lambda(\boldsymbol{x}))=k$ is not difficult — either bisection or Newton's method will do. Taking bisection as an example, it's easy to see that $\lambda(\boldsymbol{x})\in[x_{\min} - \sigma^{-1}(k/n), x_{\max} - \sigma^{-1}(k/n)]$, where $\sigma^{-1}$ is the inverse function of $\sigma$. Starting from this initial interval, we can repeatedly bisect down to the desired precision:
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 < k else (lamb, high)
return sigmoid(x - lamb)
k, x = 10, np.random.randn(100)
ThreTopK(x, k)
So the numerical computation of $\lambda(\boldsymbol{x})$ isn't particularly difficult; the real difficulty is that when we compute $\lambda(\boldsymbol{x})$ numerically, we typically lose the gradient of $\lambda(\boldsymbol{x})$ with respect to $\boldsymbol{x}$, which affects end-to-end training. To address this, we can compute $\nabla_{\boldsymbol{x}}\lambda(\boldsymbol{x})$ by hand and then customize the backward pass. Specifically, we take the partial derivative with respect to some $x_j$ on both sides of
\begin{equation}\sum_{i=1}^n \sigma(x_i - \lambda(\boldsymbol{x})) = k\end{equation}
obtaining
\begin{equation}\sigma'(x_j - \lambda(\boldsymbol{x}))-\sum_{i=1}^n \sigma'(x_i - \lambda(\boldsymbol{x}))\frac{\partial\lambda(\boldsymbol{x})}{\partial x_j} = 0\end{equation}
and thus
\begin{equation}\frac{\partial\lambda(\boldsymbol{x})}{\partial x_j} = \frac{\sigma'(x_j - \lambda(\boldsymbol{x}))}{\sum\limits_{i=1}^n \sigma'(x_i - \lambda(\boldsymbol{x}))}\end{equation}
where $\sigma'$ is the derivative of $\sigma$. Now that we have an expression for $\nabla_{\boldsymbol{x}}\lambda(\boldsymbol{x})$, every term of which is computable ($\lambda(\boldsymbol{x})$ has already been found numerically as well), we can directly specify it as the result of the backward pass. A relatively simple and general implementation trick is the $\text{stop_gradient}$ (henceforth $\text{sg}$) trick: when implementing the model, replace $\lambda(\boldsymbol{x})$ with
\begin{equation}\boldsymbol{x}\cdot\text{sg}[\nabla_{\boldsymbol{x}}\lambda(\boldsymbol{x})] + \text{sg}[\lambda(\boldsymbol{x}) - \boldsymbol{x}\cdot\nabla_{\boldsymbol{x}}\lambda(\boldsymbol{x})]\end{equation}
where $\cdot$ denotes the inner product of vectors. This way, in the forward pass it's as if $\text{sg}$ didn't exist, so the result is just $\lambda(\boldsymbol{x})$; but in the backward pass, the gradient contributed by the part being subtracted, $\text{sg}$, is zero, so the gradient we get is exactly the specified $\nabla_{\boldsymbol{x}}\lambda(\boldsymbol{x})$. In this way we've customized the gradient of $\lambda(\boldsymbol{x})$, independent of how $\lambda(\boldsymbol{x})$ was actually computed.
Best of Both Worlds
We've now seen that $f(x)=\min(1,e^x)$ has an analytical solution but is not globally smooth, while $f(x)=\sigma(x)$ is smooth enough but more complicated to solve. Is there a choice that combines the advantages of both? There indeed is — I found that the following $f(x)$ is globally smooth, and $\lambda(\boldsymbol{x})$ can be solved analytically for it:
\begin{equation}f(x) = \left\{\begin{aligned}1 - e^{-x}/2,\quad x\geq 0 \\ e^x / 2,\quad x < 0\end{aligned}\right.\end{equation}
which can also be written as $f(x) = (1 - e^{-|x|})\text{sign}(x)/2+1/2$. One can verify that $f(x)$ is in fact also a sigmoid-shaped function; although it's piecewise, both it and its derivative are continuous at $x=0$, so it is smooth enough.
The solving approach is the same as before: without loss of generality assume $x_1 > x_2 > \cdots > x_n$, and suppose we already know $x_m \geq \lambda(\boldsymbol{x}) \geq x_{m+1}$. Then
\begin{equation}\begin{aligned} k =&\, \sum_{i=1}^m (1 - e^{-(x_i - \lambda(\boldsymbol{x}))}/2) + \sum_{i=m+1}^n e^{x_i - \lambda(\boldsymbol{x})}/2 \\ =&\, m - \frac{1}{2}e^{\lambda(\boldsymbol{x})}\sum_{i=1}^m e^{-x_i} + \frac{1}{2}e^{-\lambda(\boldsymbol{x})}\sum_{i=m+1}^n e^{x_i} \end{aligned}\end{equation}
from which we solve
\begin{equation}\lambda(\boldsymbol{x})=\log\sum_{i=m+1}^n e^{x_i} - \log\left(\sqrt{(k-m)^2 + \left(\sum_{i=1}^m e^{-x_i}\right)\left(\sum_{i=m+1}^n e^{x_i}\right)}+(k-m)\right)\end{equation}
Then we enumerate $m=0,1,\cdots,n-1$ and find the $\lambda(\boldsymbol{x})$ satisfying $x_m \geq \lambda(\boldsymbol{x}) \geq x_{m+1}$. Readers may also want to try proving that when $k=1$, ThreTopK under this $f(x)$ also reduces exactly to Softmax.
Reference implementation:
import numpy as np
def ThreTopK(x, k):
x_sort = np.sort(x)
lse1 = np.logaddexp.accumulate(x_sort)
lse2 = np.pad(np.logaddexp.accumulate(-x_sort[::-1])[::-1], (0, 1), constant_values=-np.inf)[1:]
m = np.arange(len(x) - 1, -1, -1)
x_lamb = lse1 - np.log(np.sqrt((k - m)**2 + np.exp(lse1 + lse2)) + (k - m))
x_sort_shift = np.pad(x_sort[1:], (0, 1), constant_values=np.inf)
lamb = x_lamb[(x_lamb <= x_sort_shift) & (x_lamb >= x_sort)]
return (1 - np.exp(-np.abs(x - lamb))) * np.sign(x - lamb) * 0.5 + 0.5
k, x = 10, np.random.randn(100)
ThreTopK(x, k)
Summary
This post explored the problem of finding a smooth approximation to the Top-k operator, a general extension of the smooth approximation to Top-1 that Softmax provides. We proposed three construction approaches — iterative construction, gradient guidance, and undetermined constants — and analyzed their respective strengths and weaknesses.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.