# Fast, Differentiable and Sparse Top-k: a Convex Analysis Perspective

Michael E. Sander  Joan Puigcerver  Josip Djolonga  Gabriel Peyré  Mathieu
Blondel

###### Abstract

The top-$k$ operator returns a sparse vector, where the non-zero values
correspond to the $k$ largest values of the input. Unfortunately, because it
is a discontinuous function, it is difficult to incorporate in neural networks
trained end-to-end with backpropagation. Recent works have considered
differentiable relaxations, based either on regularization or perturbation
techniques. However, to date, no approach is fully differentiable _and_
sparse. In this paper, we propose new differentiable _and_ sparse top-$k$
operators. We view the top-$k$ operator as a linear program over the
permutahedron, the convex hull of permutations. We then introduce a $p$-norm
regularization term to smooth out the operator, and show that its computation
can be reduced to isotonic optimization. Our framework is significantly more
general than the existing one and allows for example to express top-$k$
operators that select values _in magnitude_. On the algorithmic side, in
addition to pool adjacent violator (PAV) algorithms, we propose a new GPU/TPU-
friendly Dykstra algorithm to solve isotonic optimization problems. We
successfully use our operators to prune weights in neural networks, to fine-
tune vision transformers, and as a router in sparse mixture of experts.

top-k, permutahedron, isotonic optimization, smoothing

##  1 Introduction

Finding the top-$k$ values and their corresponding indices in a vector is a
widely used building block in modern neural networks. For instance, in sparse
mixture of experts (MoEs) (Shazeer et al., 2017; Fedus et al., 2022), a
top-$k$ router maps each token to a selection of $k$ experts (or each expert
to a selection of $k$ tokens). In beam search for sequence decoding (Wiseman &
Rush, 2016), a beam of $k$ possible output sequences is maintained and updated
at each decoding step. For pruning neural networks, the top-$k$ operator can
be used to sparsify a neural network, by removing weights with the smallest
magnitude (Han et al., 2015; Frankle & Carbin, 2018). Finally, top-$k$
accuracy (e.g., top-$3$ or top-$5$) is frequently used to evaluate the
performance of neural networks at inference time.

Figure 1: Illustration of our differentiable and sparse top-k mask. For $k=2$,
we consider $\bm{\theta}(s)=(3,1,-1+s,s)\in\mathbb{R}^{4}$ and plot
$\mathrm{topkmask}(\bm{\theta}(s))_{2}+\mathrm{topkmask}(\bm{\theta}(s))_{3}$
as a function of $s$. We compare the hard version (no regularization) with our
proposed operator using $p$-norm regularization: $p=2$ leads to differentiable
a.e. operator; $p=4/3$, leads to a differentiable operator. Both operators are
sparse: they are exactly 0 for some values of $s$.

However, the top-$k$ operator is a discontinuous piecewise affine function
with derivatives either undefined or constant (the related top-$k$ _mask_
operator, which returns a binary encoding of the indices corresponding to the
top-$k$ values, has null derivatives). This makes it hard to use in a neural
network trained with gradient backpropagation. Recent works have considered
differentiable relaxations, based either on regularization or perturbation
techniques (see §2 for a review). However, to date, no approach is
differentiable everywhere _and_ sparse. Sparsity is crucial in neural networks
that require _conditional_ computation. This is for instance the case in
sparse mixture of experts, where the top-$k$ operator is used to “route”
tokens to selected experts. Without sparsity, all experts would need to
process all tokens, leading to high computational cost. This is also the case
when selecting weights with highest magnitude in neural networks: the non-
selected weights should be exactly $0$ in order to optimize the computational
and memory costs.

In this work, we propose novel differentiable everywhere _and_ sparse top-$k$
operators (see Figure 1). We build upon the framework of Blondel et al.
(2020b), which casts sorting and ranking as linear programs over the
permutahedron, and uses a reduction to isotonic optimization. We significantly
generalize that framework in several ways. Specifically, we make the following
contributions:

  * •

After reviewing related work in §2 and background in §3, we introduce our
generalized framework in §4. We introduce a new nonlinearity $\varphi$,
allowing us to express new operators (such as the top-$k$ in magnitude). In
doing so, we also establish new connections between so-called $k$-support
norms and the permutahedron.

  * •

We introduce a regularization term to obtain a relaxed top-$k$ operator. In
particular, using $p$-norm regularization, we obtain the first differentiable
everywhere _and_ sparse top-$k$ operator (Figure 3).

  * •

In §5, we derive pool adjacent violator (PAV) algorithms for solving isotonic
optimization when using $\varphi$ and/or when using $p$-norm regularization.
We show that the Jacobian of our operator can be computed in closed form.

  * •

As a GPU/TPU friendly alternative to PAV, we propose a Dykstra algorithm to
solve isotonic optimization, which is easy to vectorize in the case $p=2$.

  * •

In §6, we chose to focus on three applications of our operators. First, we use
them to prune weights in a multilayer perceptron during training and show that
they lead to better accuracy than with a hard top-$k$. Second, we define
top-$k$ losses to fine-tune vision transformers (ViTs) and obtain better
top-$k$ accuracy than with the cross-entropy loss. Finally, we use our
operators as a router in vision mixture of experts, and show that they
outperform the hard top-$k$ router.

##  2 Related work

#### Differentiable loss functions.

Several works have proposed a differentiable loss function as a surrogate for
a discrete, discontinuous metric. For example, loss functions have been
proposed for top-$k$ accuracy (Lapin et al., 2015, 2016; Berrada et al., 2018;
Petersen et al., 2022) and various ranking metrics (Chapelle & Wu, 2010; Adams
& Zemel, 2011; Rolínek et al., 2020).

#### Differentiable operators.

In a different line of research, which is the main focus of this work, a
differentiable operator is proposed, which can be used either as an
intermediate layer in a neural network or as final output, fed into an
arbitrary loss function. For example, Niepert et al. (2021) proposed a
framework for computing gradients of discrete probability distributions or
optimization problems. Amos et al. (2019) and Qian et al. (2022) proposed a
smooth (but not sparse) top-$k$ operator based on binary entropy
regularization. Related projections on the capped simplex were proposed
(Martins & Kreutzer, 2017; Malaviya et al., 2018; Blondel, 2019) in different
contexts. These works use ad-hoc algorithms, while we use a reduction to
isotonic optimization. Cuturi et al. (2019) proposed a relaxation of the
sorting and ranking operators based on entropy-regularized optimal transport
and used it to obtain a differentiable (but again not sparse) top-$k$
operator. Its computation relies on Sinkhorn’s algorithm (Sinkhorn, 1967;
Cuturi, 2013), which in addition makes it potentially slow to compute and
differentiate. A similar approach was proposed by Xie et al. (2020). Petersen
et al. (2021) propose smooth differentiable sorting networks by combining
differentiable sorting functions with sorting networks. Other relaxation of
the sort and operators have been proposed by Grover et al. (2019) and Prillo &
Eisenschlos (2020).

The closest work to ours is that of Blondel et al. (2020b), in which sorting
and ranking are cast as linear programs over the permutahedron. To make these
operators differentiable, regularization is introduced in the formulation and
it is shown that the resulting operators can be computed via isotonic
optimization in $O(n\log n)$ time. Unfortunately, the proposed operators still
include kinks: they are not differentiable everywhere. The question of how to
construct a differentiable everywhere _and_ sparse relaxation with $O(n\log
n)$ time complexity is therefore still open. In this work, we manage to do so
by using $p$-norm regularization. Furthermore, by introducing a new
nonlinearity $\varphi$, we significantly generalize the framework of Blondel
et al. (2020b), allowing us for instance to express a new top-$k$ operator _in
magnitude_. We introduce a new GPU/TPU friendly Dykstra algorithm as an
alternative to PAV.

Instead of introducing regularization, another technique relies on
perturbation (Berthet et al., 2020). This technique has been used to obtain a
differentiable (but still _not sparse_) top-k for image patch selection
(Cordonnier et al., 2021).

#### Pruning weights with small magnitude.

Many recent works focus on neural network pruning, where parameters are
removed to significantly reduce the size of a model. See Blalock et al. (2020)
for a recent survey. A simple yet popular method for pruning neural networks
is by global magnitude pruning (Collins & Kohli, 2014; Han et al., 2015):
weights with lowest absolute value are set to $0$. While most of the pruning
techniques are performed after the model is fully trained (Blalock et al.,
2020), some works prune periodically during training (Gale et al., 2019).
However, to the best of our knowledge, pruning by magnitude is not done in a
differentiable fashion. In this work, we empirically show that pruning weights
with a differentiable (or differentiable almost everywhere) top-$k$ operator
in magnitude during training leads to faster convergence and better accuracy
than with a “hard” one.

#### Top-k operator for mixture of experts.

Sparse mixture of experts models (MoEs) (Shazeer et al., 2017) are a class of
deep learning models where only a small proportion of the model, known as
experts, is activated, depending on its input. Therefore, sparse MoEs are able
to increase the number of parameters without increasing the time complexity of
the model. Sparse MoEs have achieved great empirical successes in computer
vision (Riquelme et al., 2021; Zhou et al., 2022) as well as natural language
processing (Shazeer et al., 2017; Lewis et al., 2021; Fedus et al., 2021). At
the heart of the sparse MoE model is its routing mechanism, which determines
which inputs (or tokens) are assigned to which experts. In the sparse mixture
of experts literature, some works have recently proposed new top-$k$ operators
in the routing module. Hazimeh et al. (2021) proposed a binary encoding
formulation to select non-zero weights. However, their formulation does not
approximate the true top-$k$ operator and sparsity is only supported at
inference time, not during training. Liu et al. (2022) proposed an optimal
transport formulation supporting $k$-sparsity constraints and used it for
sparse mixture of experts. In this work, we propose to replace the hard
top-$k$ router, which is a discontinuous function, by our smooth relaxation.

##  3 Background

#### Notation.

We denote a permutation of $[n]$ by $\sigma=(\sigma_{1},\dots,\sigma_{n})$ and
its inverse by $\sigma^{-1}$. When seen as a vector, we denote it
$\bm{\sigma}$. We denote the set of all $n!$ permutations by $\Sigma$. Given a
vector $\bm{x}\in\mathbb{R}^{n}$, we denote the version of $\bm{x}$ permuted
according to $\sigma$ by
$\bm{x}_{\sigma}\coloneqq(x_{\sigma_{1}},\dots,x_{\sigma_{n}})$. We denote the
$i$-th largest value of $\bm{x}\in\mathbb{R}^{n}$ by $x_{[i]}$. Without loss
of generality, we always sort values in descending order. The conjugate of
$f(\bm{x})$ is denoted by
$f^{*}(\bm{y})\coloneqq\sup_{\bm{x}\in\mathbb{R}^{n}}\langle\bm{x},\bm{y}\rangle-f(\bm{x})$.

#### Review of operators.

We denote the argsort operator as the permutation $\bm{\sigma}$ sorting
$\bm{x}\in\mathbb{R}^{n}$, i.e.,

$\mathrm{argsort}(\bm{x})\coloneqq\bm{\sigma},\quad\text{where}\quad
x_{\sigma_{1}}\geq\dots\geq x_{\sigma_{n}}.$ (1)

We denote the sort operator as the values of $\bm{x}\in\mathbb{R}^{n}$ in
sorted order, i.e.,

$\mathrm{sort}(\bm{x})\coloneqq\bm{x}_{\sigma},\quad\text{where}\quad\bm{\sigma}=\mathrm{argsort}(\bm{x}).$
(2)

The value $[\mathrm{sort}(\bm{x})]_{i}$ is also known as the $i$-th order
statistic. We denote the rank operator as the function returning the positions
of the vector $\bm{x}\in\mathbb{R}^{n}$ in the sorted vector. It is formally
equal to the argsort’s inverse permutation:

$\mathrm{rank}(\bm{x})\coloneqq\bm{\sigma}^{-1},\quad\text{where}\quad\bm{\sigma}=\mathrm{argsort}(\bm{x}).$
(3)

Smaller rank $[\mathrm{rank}(\bm{x})]_{i}$ means that $x_{i}$ has higher
value. The top-k mask operator returns a bit-vector encoding whether each
value $x_{i}$ is within the top-$k$ values or not:

$[\mathrm{topkmask}(\bm{x})]_{i}\coloneqq\begin{cases}1,&\text{if
}[\mathrm{rank}(\bm{x})]_{i}\leq k\\\ 0,&\text{otherwise.}\end{cases}.$ (4)

The top-k operator returns the values themselves if they are within the
top-$k$ values or $0$ otherwise, i.e.,

$\mathrm{topk}(\bm{x})\coloneqq\bm{x}\circ\mathrm{topkmask}(\bm{x}),$ (5)

where $\circ$ denotes element-wise multiplication. The top-k in magnitude
operator is defined similarly as

$\mathrm{topkmag}(\bm{x})\coloneqq\bm{x}\circ\mathrm{topkmask}(|\bm{x}|).$ (6)

To illustrate, if $x_{3}\geq x_{1}\geq x_{2}$ and
$|x_{2}|\geq|x_{3}|\geq|x_{1}|$, then

  * •

$\mathrm{argsort}(\bm{x})=(3,1,2)$

  * •

$\mathrm{sort}(\bm{x})=(x_{3},x_{1},x_{2})$

  * •

$\mathrm{rank}(\bm{x})=(2,3,1)$

  * •

$\mathrm{topkmask}(\bm{x})=(1,0,1)$

  * •

$\mathrm{topk}(\bm{x})=(x_{1},0,x_{3})$,

  * •

$\mathrm{topkmag}(\bm{x})=(0,x_{2},x_{3})$,

where in the last three, we used $k=2$.

#### Permutahedron.

The permutahedron associated with a vector $\bm{w}\in\mathbb{R}^{n}$, a well-
known object in combinatorics (Bowman, 1972; Ziegler, 2012), is the convex
hull of the permutations of $\bm{w}$, i.e.,

$P(\bm{w})\coloneqq\operatorname*{conv}(\\{\bm{w}_{\sigma}\colon\sigma\in\Sigma\\})\subset\mathbb{R}^{n}.$
(7)

We define the linear maximization oracles (LMO) associated with $P(\bm{w})$ by

$\displaystyle f(\bm{x},\bm{w})$ $\displaystyle\coloneqq\max_{\bm{y}\in
P(\bm{w})}\langle\bm{x},\bm{y}\rangle$ (8)
$\displaystyle\bm{y}(\bm{x},\bm{w})$
$\displaystyle\coloneqq\operatorname*{argmax}_{\bm{y}\in
P(\bm{w})}~{}\langle\bm{x},\bm{y}\rangle=\nabla_{1}f(\bm{x},\bm{w}),$

where $\nabla_{1}f(\bm{x},\bm{w})$ is technically a subgradient of $f$ w.r.t.
$\bm{x}$. The LMO can be computed in $O(n\log n)$ time. Indeed, the
calculation of the LMO reduces to a sorting operation, as shown in the
following known proposition. A proof is included for completeness in Appendix
A.1. (Linear maximization oracles)

If $w_{1}\geq\dots\geq w_{n}$ (if not, sort $\bm{w}$), then

$f(\bm{x},\bm{w})=\sum_{i=1}^{n}w_{i}x_{[i]}\quad\text{and}\quad\bm{y}(\bm{x},\bm{w})=\bm{w}_{\mathrm{rank}(\bm{x})}.$
(9)

#### LP formulations.

Let us denote the reversing permutation by
$\bm{\rho}\coloneqq(n,n-1,\dots,1)$. Blondel et al. (2020b) showed that the
sort and rank operators can be formulated as linear programs (LP) over the
permutahedron:

$\displaystyle\mathrm{sort}(\bm{x})$
$\displaystyle=\bm{y}(\bm{\rho},\bm{x})=\operatorname*{argmax}_{\bm{y}\in
P(\bm{x})}\langle\bm{\rho},\bm{y}\rangle$ (10)
$\displaystyle\mathrm{rank}(\bm{x})$
$\displaystyle=\bm{y}(-\bm{x},\bm{\rho})=\operatorname*{argmax}_{\bm{y}\in
P(\bm{\rho})}\langle-\bm{x},\bm{y}\rangle.$

In the latter expression, the minus sign is due to the fact that we use the
convention that smaller rank indicates higher value (i.e., the maximum value
has rank $1$).

Although not mentioned by Blondel et al. (2020b), it is also easy to express
the top-$k$ mask operator as an LP

$\mathrm{topkmask}(\bm{x})=\bm{y}(\bm{x},\mathbf{1}_{k})=\operatorname*{argmax}_{\bm{y}\in
P(\mathbf{1}_{k})}\langle\bm{x},\bm{y}\rangle,$ (11)

where
$\mathbf{1}_{k}\coloneqq(\underbrace{1,...,1}_{k},\underbrace{0,\dots,0}_{n-k})$.
For this choice of $\bm{w}$, the permutahedron enjoys a particularly simple
expression

$P(\mathbf{1}_{k})=\\{\bm{y}\in\mathbb{R}^{n}\colon\langle\bm{y},\mathbf{1}\rangle=k,\bm{y}\in[0,1]^{n}\\}$
(12)

and $P(\mathbf{1}_{k}/k)$ is known as the capped simplex (Warmuth & Kuzmin,
2008; Blondel et al., 2020a). This is illustrated in Figure 2. To obtain
relaxed operators, Blondel et al. (2020b) proposed to introduce regularization
in (8) (see “recovering the previous framework” in the next section) and used
a reduction to isotonic optimization.

Figure 2: The permutahedron $P(\bm{w})$ is a polytope whose vertices are
permutations of $\bm{w}$. Depending, on the choice of $\bm{w}$, it can express
several known polytopes. When $\bm{w}=(1,0,0)$, $P(\bm{w})$ is the probability
simplex (light beige), which corresponds to the top-$1$ setting. When
$\bm{w}=(\frac{1}{2},\frac{1}{2},0)$, $P(\bm{w})$ is the capped probability
simplex (blue), which corresponds to the top-$k$ setting (here, with $k=2$).
When $\bm{w}=(\frac{2}{3},\frac{1}{3},0)$, $P(\bm{w})$ is an hexagon, which
corresponds to the partial ranking setting (gray).

##  4 Proposed generalized framework

In this section, we generalize the framework of Blondel et al. (2020b) by
adding an optional nonlinearity $\varphi(\bm{x})$. In addition to the
operators covered by the previous framework, this allows us to directly
express the top-$k$ in magnitude operator, which was not possible before. We
also support $p$-norm regularization, which allows to express differentiable
_and_ sparse operators when $1<p<2$.

#### Introducing a mapping $\varphi$.

Consider a mapping $\varphi(\bm{x})\coloneqq(\phi(x_{1}),\dots,\phi(x_{n}))$.
Given $\bm{x}\in\mathbb{R}^{n}$ and $\bm{w}\in\mathbb{R}^{n}$, we define

$\displaystyle f_{\varphi}(\bm{x},\bm{w})$ $\displaystyle\coloneqq
f(\varphi(\bm{x}),\bm{w})$ (13) $\displaystyle\bm{y}_{\varphi}(\bm{x},\bm{w})$
$\displaystyle\coloneqq\nabla_{1}f_{\varphi}(\bm{x},\bm{w}).$

When $\varphi(\bm{x})=\bm{x}$ (identity mapping), we clearly recover the
existing framework, i.e., $f_{\varphi}(\bm{x},\bm{w})=f(\bm{x},\bm{w})$ and
$\bm{y}_{\varphi}(\bm{x},\bm{w})=\bm{y}(\bm{x},\bm{w})$.

When $\varphi(\bm{x})\neq\bm{x}$, our framework starts to differ from the
previous one, since $\varphi$ affects differentiation. By the chain rule and
Danskin’s theorem (1966), we get

$\displaystyle\bm{y}_{\varphi}(\bm{x},\bm{w})$
$\displaystyle\coloneqq\nabla_{1}f_{\varphi}(\bm{x},\bm{w})$ (14)
$\displaystyle=\partial\varphi(\bm{x})^{\top}\bm{y}(\varphi(\bm{x}),\bm{w})$
$\displaystyle=(\phi^{\prime}(x_{1}),\dots,\phi^{\prime}(x_{n}))\circ\bm{y}(\varphi(\bm{x}),\bm{w}),$

where $\partial\varphi(\bm{x})\in\mathbb{R}^{n\times n}$ denotes the Jacobian
of $\varphi(\bm{x})$ and $\bm{y}(\bm{x},\bm{w})$ is given by Proposition 8.

#### Top-k in magnitude.

As we emphasized, one advantage of our proposed generalization is that we can
express the top-$k$ in magnitude operator. Indeed, with
$\phi(x)=\frac{1}{2}x^{2}$, we can see from (11) and (14) that we have for all
$\bm{x}\in\mathbb{R}^{n}$

$\mathrm{topkmag}(\bm{x})=\bm{y}_{\varphi}(\bm{x},\mathbf{1}_{k})=\nabla_{1}f_{\varphi}(\bm{x},\mathbf{1}_{k}).$
(15)

Obviously, for all $\bm{x}\in\mathbb{R}^{n}_{+}$, we also have
$\mathrm{topkmag}(\bm{x})=\mathrm{topk}(\bm{x})$. Top-k in magnitude is useful
for pruning weights with small magnitude in a neural network, as we
demonstrate in our experiments in §6.

#### Introducing regularization.

We now explain how to make our generalized operator differentiable. We
introduce convex regularization $R:\mathbb{R}^{n}\to\mathbb{R}$ in the dual
space:

$f^{*}_{\varphi,R}(\bm{y},\bm{w})\coloneqq
f_{\varphi}^{*}(\bm{y},\bm{w})+R(\bm{y}),$ (16)

where $f^{*}_{\varphi}$ is the conjugate of $f_{\varphi}$ in the first
argument. Going back to the primal space, we obtain a new relaxed operator. A
proof is given in Appendix A.2. (Relaxed operator)

Let $R\colon\mathbb{R}^{n}\to\mathbb{R}$ be a convex regularizer. Then

$\displaystyle f_{\varphi,R}(\bm{x},\bm{w})$
$\displaystyle\coloneqq\max_{\bm{y}\in\mathbb{R}^{n}}\langle\bm{y},\bm{x}\rangle-
f_{\varphi}^{*}(\bm{y},\bm{w})-R(\bm{y})$ (17)
$\displaystyle=\min_{\bm{u}\in\mathbb{R}^{n}}R^{*}(\bm{x}-\bm{u})+f_{\varphi}(\bm{u},\bm{w})$
$\displaystyle\bm{y}_{\varphi,R}(\bm{x},\bm{w})$
$\displaystyle\coloneqq\bm{y}^{\star}=\nabla
R^{*}(\bm{x}-\bm{u}^{\star})=\nabla_{1}f_{\varphi,R}(\bm{x},\bm{w}).$

The mapping $\varphi$ also affects conjugacy. We have the following
proposition. (Conjugate of $f_{\varphi}$ in the first argument)

If $\phi$ is convex and $\bm{w}\in\mathbb{R}_{+}^{n}$, then

$f_{\varphi}^{*}(\bm{y},\bm{w})=\min_{\bm{z}\in
P(\bm{w})}D_{\phi^{*}}(\bm{y},\bm{z}),$ (18)

where $D_{f}(\bm{y},\bm{z})\coloneqq\sum_{i=1}^{n}z_{i}f(y_{i}/z_{i})$. A
proof is given in Appendix A.3. The function $(y_{i},z_{i})\mapsto
y^{\prime}_{i}\phi^{*}(y_{i}/z_{i})$ is known as the perspective of $\phi^{*}$
and is jointly convex when $z_{i}>0$. The function $D_{f}(\bm{y},\bm{z})$ is
known as the $f$-divergence between $\bm{y}$ and $\bm{z}$. Therefore,
$f_{\varphi}^{*}(\bm{y},\bm{w})$ can be seen as the minimum “distance” between
$\bm{y}$ and $P(\bm{w})$ in the $\phi^{*}$-divergence sense.

#### Recovering the previous framework.

If $\phi(x)=x$, then

$\phi^{*}(y_{i}/z_{i})=\begin{cases}0,&\text{if }y_{i}=z_{i}\\\
\infty,&\text{otherwise}\end{cases}.$ (19)

This implies that
$f_{\varphi}^{*}(\bm{y},\bm{w})=f^{*}(\bm{y},\bm{w})=\delta_{P(\bm{w})}(\bm{y})$,
the indicator function of $P(\bm{w})$, which is $0$ if $\bm{y}\in P(\bm{w})$
and $\infty$ otherwise. In this case, we therefore obtain

$f_{\varphi,R}(\bm{x},\bm{w})=\max_{\bm{y}\in
P(\bm{w})}\langle\bm{y},\bm{x}\rangle-R(\bm{y}),$ (20)

which is exactly the relaxation of Blondel et al. (2020b).

#### Differentiable and sparse top-k operators.

To obtain a relaxed top-$k$ operator with our framework, we simply replace
$f_{\varphi}$ with $f_{\varphi,R}$ and $\bm{y}_{\varphi}$ with
$\bm{y}_{\varphi,R}$ in (15) to define

$\mathrm{topkmag}_{R}(\bm{x})\coloneqq\bm{y}_{\varphi,R}(\bm{x},\mathbf{1}_{k})=\nabla_{1}f_{\varphi,R}(\bm{x},\mathbf{1}_{k}).$
(21)

A relaxed top-$k$ mask can be defined in a similar way, but using $\phi(x)=x$
instead of $\phi(x)=\frac{1}{2}x^{2}$. For the regularization $R$, we propose
to use $p$-norms to the power $p$:

$R(\bm{y})=\frac{1}{p}\|\bm{y}\|_{p}^{p}\coloneqq\frac{1}{p}\sum_{i=1}^{n}|y_{i}|^{p}.$
(22)

The choice $p=2$ used in previous works leads to sparse outputs but is not
differentiable everywhere. Any $p$ between $1$ (excluded) and $2$ (excluded)
leads to differentiable _and_ sparse outputs. We propose to use $p=4/3$ to
obtain a differentiable everywhere operator, and $p=2$ to obtain a
differentiable a.e. operator, which is more convenient numerically. This is
illustrated in Figure 3.

#### Connection with $k$-support and OWL norms.

When $\phi(x)=\frac{1}{2}x^{2}$ and $\bm{w}=\mathbf{1}_{k}$, we obtain

$f_{\varphi}^{*}(\bm{y},\bm{w})=\frac{1}{2}\min_{\bm{z}\in[0,1]^{n}}\sum_{i=1}^{n}\frac{y_{i}^{2}}{z_{i}}\quad\text{s.t.}\quad\langle\bm{z},\mathbf{1}\rangle=k,$
(23)

which is known as the squared k-support norm (Argyriou et al., 2012; McDonald
et al., 2014; Eriksson et al., 2015). Our formulation is a generalization of
the squared $k$-support norm, as it supports other choices of $\bm{w}$ and
$\phi$. For instance, we use it to define a new notion of $k$-support
negentropy in Appendix B.1. When $\varphi(\bm{x})=|\bm{x}|$, we recover the
ordered weighted lasso (OWL) norm (Zeng & Figueiredo, 2014) as

$f_{\varphi}(\bm{x},\bm{w})=\sum_{i=1}^{n}w_{i}|x|_{[i]}.$ (24)

With that choice of $\varphi$, it is easy to see from (14) that (15) becomes a
_signed_ top-$k$ mask. Note that, interestingly, $k$-support and OWL norms are
not defined in the same space.

Figure 3: Example of our relaxed top-k operators. We take
$\phi(x)=\frac{1}{2}x^{2}$ and $k=2$. For an input
$\bm{x}=(x_{1},x_{2},\frac{1}{2},1)$, we plot
$y_{\varphi,R}(\bm{x},1_{k})_{1}+y_{\varphi,R}(\bm{x},1_{k})_{2}$ for $R=0$
(left), $R=\frac{\lambda}{2}\|\bm{x}\|^{2}_{2}$ (center) and
$R=\frac{\lambda}{p}\|\bm{x}\|^{p}_{p}$ with $p=\frac{4}{3}$ (right). We take
$\lambda=0.3$. While no regularization leads to a discontinuous mapping, the
$2$-norm regularization leads to continuity and a.e. differentiability, and
the ${\frac{4}{3}}$-norm regularization provides a continuously differentiable
mapping. We emphasize that, although in the left plot the graph looks
connected, it is actually a discontinuous function. Note that our relaxed
operators are sparse as they are exactly $0$ in the center.

#### Biconjugate interpretation.

Let us define the set of $k$-sparse vectors, which is nonconvex, as
$S_{k}\coloneqq\\{\bm{x}\in\mathbb{R}^{n}\colon\|\bm{x}\|_{0}\leq k\\},$ where
$\|\bm{x}\|_{0}$ is the number of non zero elements in $\bm{x}$. We saw that
if $\phi(x)=\frac{1}{2}x^{2}$ then $f_{\varphi}^{*}(\bm{y},\mathbf{1}_{k})$ is
the squared $k$-support norm. It is known to be the biconjugate (i.e., the
tightest convex relaxation) of the squared $L_{2}$ norm restricted to $S_{k}$
(Eriksson et al., 2015; Liu et al., 2022). We now prove a more general result:
$f_{\varphi}^{*}(\bm{y},\mathbf{1}_{k})$ is the biconjugate of
$\sum_{i=1}^{n}\phi^{*}(y_{i})$ restricted to $S_{k}$. (Biconjugate
interpretation)

Let
$\Phi_{k}(\bm{y})\coloneqq\sum_{i=1}^{n}\phi^{*}(y_{i})+\delta_{S_{k}}(\bm{y})$.
Suppose that $\phi$ is convex. Then the biconjugate of $\Phi_{k}$ is given by

$\Phi^{**}_{k}(\bm{y})=f_{\varphi}^{*}(\bm{y},\mathbf{1}_{k})$ (25)

See Appendix A.4 for a proof.

##  5 Algorithms

In this section, we propose efficient algorithms for computing our operators.
We first show that the calculation of our relaxed operator reduces to isotonic
optimization.

#### Reduction to isotonic optimization.

We now show how to compute $\bm{u}^{\star}$ in Proposition 16 by reduction to
isotonic optimization, from which $\bm{y}^{\star}$ can then be recovered by
$\bm{y}^{\star}=\nabla R^{*}(\bm{x}-\bm{u}^{\star})$. We first recall the case
$\varphi(\bm{x})=\bm{x}$, which was already proved in existing works (Lim &
Wright, 2016; Blondel et al., 2020b). (Reduction, $\varphi(\bm{x})=\bm{x}$
case)

Suppose that $R(\bm{y})=\sum_{i=1}^{n}r(y_{i})$. Let $\sigma$ be the
permutation sorting $\bm{x}$, $\bm{s}\coloneqq\bm{x}_{\sigma}$ and

$\bm{v}^{\star}=\operatorname*{argmin}_{v_{1}\geq\dots\geq
v_{n}}R^{*}(\bm{s}-\bm{v})+f(\bm{v},\bm{w}).$

Then $\bm{u}^{\star}$ from Proposition 16 is given by
$\bm{u}^{\star}=\bm{v}^{\star}_{\sigma^{-1}}$. The set
$\\{\bm{v}\in\mathbb{R}^{n}\colon v_{1}\geq\dots\geq v_{n}\\}$ is called the
monotone cone. Next, we show that a similar result is possible when
$\phi(\bm{x})$ and $r^{*}$ are both even functions (sign-invariant) and
increasing on $\mathbb{R}_{+}$. (Reduction, $\varphi(\bm{x})=\varphi(-\bm{x})$
case)

Suppose that $R(\bm{y})=\sum_{i=1}^{n}r(y_{i})$ and
$\varphi(\bm{x})=(\phi(x_{1}),\dots,\phi(x_{n}))$. Assume $\phi$ and $r^{*}$
are both even functions (sign-invariant) and increasing on $\mathbb{R}_{+}$.
Let $\sigma$ be the permutation sorting $|\bm{x}|$,
$\bm{s}\coloneqq|\bm{x}|_{\sigma}$ and

$\bm{v}^{\star}=\operatorname*{argmin}_{v_{1}\geq\dots\geq v_{n}\geq
0}R^{*}(\bm{s}-\bm{v})+f_{\varphi}(\bm{v},\bm{w}).$

Then, $\bm{u}^{\star}$ (Proposition 16) is equal to
$\mathrm{sign}(\bm{x})\circ\bm{v}^{\star}_{\sigma^{-1}}$. See Appendix A.6 for
a proof. Less general results are proved in (Zeng & Figueiredo, 2014; Eriksson
et al., 2015) for specific cases of $\phi$ and $R$. The set
$\\{v\in\mathbb{R}^{n}\colon v_{1}\geq\dots\geq v_{n}\geq 0\\}$ is called the
non-negative monotone cone. In practice, the additional non-negativity
constraint is easy to handle: we can solve the isotonic optimization problem
without it and truncate the solution if it is not non-negative (Németh &
Németh, 2012).

#### Pool adjacent violator (PAV) algorithms.

Under the conditions of Proposition 5, assuming $\bm{v}$ and $\bm{w}$ are both
sorted, we have from Proposition 8 that

$f(\bm{v},\bm{w})=\sum_{i=1}^{n}w_{i}v_{i},\quad
f_{\varphi}(\bm{v},\bm{w})=\sum_{i=1}^{n}w_{i}\phi(v_{i}).$ (26)

We then get that the problems in Proposition 5 and 5 are coordinate-wise
separable:

$\bm{v}^{\star}=\operatorname*{argmin}_{v_{1}\geq\dots\geq
v_{n}}\sum_{i=1}^{n}h_{i}(v_{i}),$ (27)

for $h_{i}(v_{i})=r^{*}(s_{i}-v_{i})+w_{i}\phi(v_{i})$. Such problems can be
solved in $O(n)$ time using the pool adjacent violator (PAV) algorithm (Best
et al., 2000). This algorithm works by partitioning the set $[n]$ into
disjoint sets $(B_{1},\cdots B_{m})$, starting from $m=n$ and $B_{i}=\\{i\\}$,
and by merging these sets until the isotonic condition is met. A pseudo-code
is available for completness in Appendix B.2. At its core, PAV simply needs a
routine to solve the “pooling” subproblem

$\gamma^{\star}_{B}=\operatorname*{argmin}_{\gamma\in\mathbb{R}}\sum_{i\in
B}h_{i}(\gamma)$ (28)

for any $B\subseteq[n]$. Once the optimal partition $(B_{1},\cdots B_{m})$ is
identified, we have that

$\bm{v}^{\star}=(\underbrace{\gamma_{B_{1}}^{\star},\dots,\gamma_{B_{1}}^{\star}}_{|B_{1}|},\cdots,\underbrace{\gamma_{B_{m}}^{\star},\dots,\gamma_{B_{m}}^{\star}}_{|B_{m}|})\in\mathbb{R}^{n}.$
(29)

Because Proposition 5 and 5 require to obtain the sorting permutation $\sigma$
beforehand, the total time complexity for our operators is $O(n\log n)$.

#### Example.

Suppose $\phi(x)=\frac{1}{2}x^{2}$ and $R=\frac{\lambda}{2}\|.\|^{2}_{2}$,
where $\lambda>0$. Since $R^{*}=\frac{1}{2\lambda}\|.\|^{2}$, this gives
$h_{i}(v_{i})=\frac{1}{2}(w_{i}v^{2}_{i}+\frac{1}{\lambda}(s_{i}-v_{i})^{2})$.
The solution of the sub-problem is then

$\gamma^{\star}_{B}=\frac{\sum_{i\in B}s_{i}}{\sum_{i\in B}(\lambda
w_{i}+1)}.$

Using this formula, when $\lambda$ is small enough, we can upper-bound the
error between the hard and relaxed operators:
${\|\mathrm{topkmag}_{R}(\bm{x})-\mathrm{topkmag}(\bm{x})\|_{\infty}}\leq\lambda\|\bm{x}\|_{\infty}$.
See Appendix A.6 for details and for the case $p=\frac{4}{3}$.

#### Dykstra’s alternating projection algorithm.

The PAV algorithm returns an exact solution of (27) in $O(n)$ time.
Unfortunately, it relies on element-wise dynamical array assignments, which
makes it potentially slow on GPUs and TPUs. We propose an alternative to
obtain faster computations. Our key insight is that by defining

$\displaystyle C_{1}$ $\displaystyle\coloneqq\\{\bm{v}\in\mathbb{R}^{n}\colon
v_{1}\geq v_{2},v_{3}\geq v_{4},\dots\\}$ (30) $\displaystyle C_{2}$
$\displaystyle\coloneqq\\{\bm{v}\in\mathbb{R}^{n}\colon v_{2}\geq
v_{3},v_{4}\geq v_{5},\dots\\},$

one has $\\{\bm{v}\in\mathbb{R}^{n}\colon v_{1}\geq\dots\geq
v_{n}\\}=C_{1}\cap C_{2}.$ We can therefore rewrite (27) as

$\bm{v}^{\star}=\operatorname*{argmin}_{\bm{v}\in C_{1}\cap
C_{2}}\sum_{i=1}^{n}h_{i}(v_{i}).$ (31)

In the case $R=\frac{1}{2}\|.\|^{2}$ and $\phi(x)=x$ or $\frac{1}{2}x^{2}$,
this reduces to a projection onto $C_{1}\cap C_{2}$, for which we can use
Dykstra’s celebrated projection algorithm (Boyle & Dykstra, 1986; Combettes &
Pesquet, 2011). Jegelka et al. (2013) have used this method for computing the
projection onto the intersection of submodular polytopes, whereas we project
onto a single polyhedral face. When $\phi(x)=x$, Dykstra’s projection
algorithm takes the simple form in Algorithm 5.

(Dykstra’s projection algorithm)

Starting from $\bm{v}^{0}=\bm{s}$, $\bm{p}^{0}=\bm{q}^{0}=\mathbf{0}$,
Dykstra’s algorithm iterates

$\displaystyle\bm{y}^{k}=\operatorname*{argmin}_{\bm{y}\in
C_{1}}\frac{1}{2}\|\bm{v}^{k}+\bm{p}^{k}-\bm{y}\|^{2}_{2}+\langle\bm{y},\bm{w}\rangle$
$\displaystyle\bm{p}^{k+1}=\bm{v}^{k}+\bm{p}^{k}-\bm{y}^{k}$
$\displaystyle\bm{v}^{k+1}=\operatorname*{argmin}_{\bm{v}\in
C_{2}}\frac{1}{2}\|\bm{y}^{k}+\bm{q}^{k}-\bm{v}\|^{2}_{2}+\langle\bm{v},\bm{w}\rangle$
$\displaystyle\bm{q}^{k+1}=\bm{y}^{k}+\bm{q}^{k}-\bm{v}^{k+1}.$

Each argmin calculation corresponds to a Euclidean projection onto $C_{1}$ or
$C_{2}$, which can be computed in closed form. Therefore, $\bm{v}^{k}$
provably converges to $\bm{v}^{\star}$. Each iteration of Dykstra’s algorithm
is learning-rate free, has linear time complexity and can be efficiently
written as a matrix-vector product using a mask, which makes it particularly
appealing for GPUs and TPUs. Interestingly, we find out that when
$\bm{w}=\mathbf{1}_{k}$, then Dykstra’s algorithm converges surprisingly fast
to the exact solution. We validate that Dykstra leads to a faster runtime than
PAV in Figure 4. We use $100$ iterations of Dykstra, and verify that we obtain
the same output as PAV.

For the general non-Euclidean $R^{*}$ case, i.e., $p\neq 2$, we can use block
coordinate ascent in the dual of (27). We can divide the dual variables into
two blocks, corresponding to $C_{1}$ and $C_{2}$; see Appendix A.8. In fact,
it is known that in the Euclidean case, Dykstra’s algorithm in the primal and
block coordinate ascent in the dual are equivalent (Tibshirani, 2017).
Therefore, although a vectorized implementation could be challenging for
$p\neq 2$, block coordinate ascent can be seen as an elegant way to generalize
Dykstra’s algorithm. Convergence is guaranteed as long as each $h_{i}$ is
strictly convex.

#### Differentiation.

The Jacobian of the solution of the isotonic optimization problems can be
expressed in closed form for any $p$-norm regularization. (Differentiation)

Let
$\bm{v}^{\star}=(\gamma_{B_{1}}^{\star},\dots,\gamma_{B_{1}}^{\star}\cdots,\gamma_{B_{m}}^{\star},\dots,\gamma_{B_{m}}^{\star})$
be the optimal solution of the isotonic optimization problem with
$R=\frac{1}{p}\|.\|^{p}_{p}$ and $p>0$. Then one has that $\bm{v}^{\star}$ is
differentiable with respect to $\bm{s}=(s_{1},\dots,s_{n})$. Furthermore, for
any $r\in\\{1,\dots,m\\}$ and $i\in B_{r}$,

$\frac{\partial\gamma^{\star}_{B_{r}}}{\partial
s_{i}}=\begin{cases}\frac{|\gamma^{\star}_{B_{r}}-s_{i}|^{q-2}}{\sum_{j\in{B_{r}}}|\gamma^{\star}_{B_{r}}-s_{j}|^{q-2}}&\text{if
}\phi(x)=x\\\ &\\\
\frac{(q-1)|\gamma^{\star}_{B_{r}}-s_{i}|^{q-2}}{\sum_{j\in{B_{r}}}(q-1)|\gamma^{\star}_{B_{r}}-s_{j}|^{q-2}+w_{j}}&\text{if
}\phi(x)=\frac{1}{2}x^{2}\end{cases}$ (32)

where $q$ is such that $\frac{1}{p}+\frac{1}{q}=1$. When $i\notin B_{r}$, one
simply has $\frac{\partial\gamma^{\star}_{B_{r}}}{\partial s_{i}}=0.$ One then
has $\partial v_{j}^{\star}/\partial
s_{i}=\partial\gamma^{\star}_{B_{r_{j}}}/{\partial s_{i}}$ where $r_{j}$ is
such that $v^{\star}_{j}=\gamma^{\star}_{B_{r_{j}}}$. See Appendix A.7 for a
proof. Thanks to Proposition 5, we do not need to solve a linear system to
compute the Jacobian of the solution $\bm{v}^{\star}$, in contrast to implicit
differentiation of general optimization problems (Blondel et al., 2021). In
practice, this also means that we do not need to perform backpropagation
through the unrolled iteratations of PAV or Dykstra’s projection algorithm to
obtain the gradient of a scalar loss function, in which our operator is
incorporated. In particular, we do not need to store the intermediate iterates
of these algorithms in memory. Along with PAV and Dykstra’s projection
algorithm, we implement the corresponding Jacobian vector product routines in
JAX, using Proposition 5.

Figure 4: Runtime comparison for computing our relaxed top-$k$ on a TPU using
PAV, Dykstra, as a function of the dimension $n$. For each $n$, we set
$k=\lceil n/10\rceil$. We also compare with the hard top-k computation.

##  6 Experiments

We now demonstrate the applicability of our top-k operators through
experiments. Our JAX (Bradbury et al., 2018) implementation is available at
the following URL. See Appendix C for additional experimental details.

#### Weight pruning in neural networks.

Figure 5: Test error with respect to training time when training an MLP on
MNIST. We compare the baseline (grey) with the case where $90\%$ of the
weights are set to 0 by magnitude pruning, using a differentiable a.e. (red),
fully differentiable (green) or a hard top-$k$ (blue).

We experimentally validate the advantage of using a smoothed top-k for weight
pruning in neural networks. We use a multilayer perceptron (MLP) with 2 hidden
layers and with ReLU activation. The width of the layers are respectively
$784$, $32$, $32$ followed by a linear classification head of width $10$. More
precisely, our model takes as input an image $\bm{a}\in\mathbb{R}^{784}$ and
outputs

$\bm{x}=W_{3}\sigma(W_{2}\sigma(W_{1}\bm{a}+\bm{b}_{1})+\bm{b}_{2})+\bm{b}_{3}\quad\text{(logits)},$

where $W_{1}\in\mathbb{R}^{32\times 784}$, $\bm{b}_{1}\in\mathbb{R}^{32}$,
$W_{2}\in\mathbb{R}^{32\times 32}$, $\bm{b}_{2}\in\mathbb{R}^{32}$,
$W_{1}\in\mathbb{R}^{10\times 32}$, $\bm{b}_{3}\in\mathbb{R}^{10}$ and
$\sigma$ is a ReLU. In order to perform weight pruning we parametrize each
$W_{i}$ as $W_{i}=\mathrm{topkmag}_{R}(W^{\prime}_{i})$ and learn
$W^{\prime}_{i}$ instead of learning $W_{i}$ directly. The output is then fed
into a cross-entropy loss. We compare the performance of the model when
applying a hard vs differentiable top-k operator to keep only $10\%$ of the
coefficients. For the differentiable top-k, we use a regularization
$R(\bm{y})=\frac{\lambda}{p}\|\bm{y}\|^{p}$ with $p\in\\{\frac{4}{3},2\\}$ and
$\lambda=10^{-4}.$ We find out that the model trained with the differentiable
top-k trains significantly faster than the one trained with the hard top-k. We
also verify that our relaxed top-k maintains the $10\%$ rate of non-zero
weights. Results on MNIST are displayed in Figure 5. We also compare with an
entropy-regularized approximation of the top-k operator using the framework
proposed in Cuturi et al. (2019) and adapted in Petersen et al. (2022). To
guarantee the sparsity of the weights, we use the ”straight-through” trick:
the hard top-k is run on the forward pass but we use the gradient of the
relaxed top-k in the backward pass. This method leads to a test error of
$5.9\%$, which is comparable to the results obtained with our differentiable
operators.

#### Smooth top-k loss.

To train a neural network on a classification task, one typically minimizes
the cross-entropy loss, whereas the performance of the network is evaluated
using a top-$k$ test accuracy. There is therefore a mismatch between the loss
used at train time and the metric used at evaluation time. Cuturi et al.
(2019) proposed to replace the cross-entropy loss with a differentiable
top-$k$ loss. In the same spirit, we propose to finetune a ViT-B/16
(Dosovitskiy et al., 2020) pretrained on the ImageNet21k dataset on CIFAR 100
(Krizhevsky et al., 2009) using a smooth and sparse top-$k$ loss instead of
the cross-entropy loss. We use a Fenchel-Young loss (Blondel et al., 2020a).
It takes as input the vector $\bm{a}$ and parameters $\theta$ of a neural
network $g_{\theta}$:

$\displaystyle\bm{x}$ $\displaystyle=g_{\theta}(\bm{a})\quad\text{(logits)}$
$\displaystyle\ell(\bm{x},\bm{t})$
$\displaystyle=f_{\varphi,R}(\bm{x},\mathbf{1}_{k})-\langle\bm{x},\bm{t}\rangle,$

where $f_{\varphi,R}(\bm{x},\mathbf{1}_{k})$ is given by Proposition 16 and
$\bm{t}$ is a one-hot encoding of the class of $\bm{a}$. We set $\phi(x)=x$ as
we want a top-$k$ mask. We consider $p$ norm regularizations for $R$, where
$p=2$ or $p=4/3$. We take $k=3$. We use the exact same training procedure as
described in Dosovitskiy et al. (2020) and use the corresponding pretrained
ViT model B/16, and train our model for $100$ steps. Results are reported in
Figure 6. We find that the ViT finetuned with the smooth top-$3$ loss
outperforms the one finetuned with the cross-entropy loss in terms of top-$k$
error, for various $k$.

Figure 6: Validation top-k accuracy when fine-tuning a ViT-B/16 on CIFAR 100
using either the cross-entropy loss or our smooth top-$3$ loss for training.
We have the following running times obtained with a TPUv3-8. Baseline: 9.5
sec/step, $p=2$: 9.7 sec/step and $p=4/3$: 10 sec/step.

#### Sparse MoEs.

Finally, we demonstrate the applicability of our proposed smooth top-$k$
operators on a large-scale classification task using vision sparse mixture of
experts (V-MoE) (Riquelme et al., 2021). Vision transformers (ViTs) are made
of a succession of self-attention layers and MLP layers. The idea of V-MoEs is
to replace MLPs in ViTs by a sparsely-gated mixture of MLPs called experts.
This way, only some of the experts are activated by a given patch token. At
the heart of the token-expert assignment lies a routing mechanism which
performs a top-$k$ operation on gate values. We focus on the MoE with expert
choice routing framework (Zhou et al., 2022), where each expert is assigned to
$k$ tokens. We train a S/32 variant of the V-MoE model, with $32\times 32$
patches on the JFT-300M dataset (Sun et al., 2017), a dataset with more than
305 million images. Our model has $32$ experts, each assigned to $k=28$ tokens
selected among $n=400$ at each MoE layer. We compare the validation accuracy
when using the baseline (hard top-$k$) with our relaxed operator. We use $p=2$
and Dykstra’s projection algorithm, as we found it was the fastest method on
TPU. We used the training procedure proposed by Zhou et al. (2022) to obtain a
fair comparison with the baseline. Due to the large size of the JFT-300M
dataset (305 million images), we performed one run, as in Liu et al. (2022).
We find that our approach improves validation performance. Results are
displayed in Figure 7.

Figure 7: Precision-at-1 on the JFT-300M dataset when using a hard top-$k$
(baseline, in blue) or a differentiable a.e. one (in red) in a sparse MoE with
a ViT-S/32. We zoom in on the last training steps, where our proposed method
outperforms the baseline. The runtime is 10 hours for the baseline and 15 for
the differentiable a.e. top-k (gradient calculation is the bottleneck here).

##  7 Discussion

#### Advantage of the non-linearity.

As an alternative to performing a relaxed top-k operator in magnitude of
$\bm{x}$, one can perform a differentiable top-k mask on $|\bm{x}|$, and then
multiply the output by $\bm{x}$. This alternative would also lead to a
differentiable top-k operator in magnitude. However, our operator has more
principled behavior at the limit cases. For instance, as $\lambda\to\infty$,
it is easy to see that the relaxed top-k mask converges to the vector
$(k/n)\times 1_{n}$. Therefore, a rescaling by $n/k$ is needed to obtain the
identity as $\lambda\to\infty$, in contrast to our top-k in magnitude. From a
theoretical point of view, the introduction of a non-linearity allows us to
draw connections with the $k$-support norm. It also has a bi-conjugate
interpretation, which we believe has an interest by itself.

#### Sensitivity to the choice of p.

The subproblem needed within PAV enjoys a closed form only for specific
choices of $p$. This is why we focused on $p=2$ and $p=\frac{4}{3}$ in our
experiments. However, we stress out that the proposed methods work for any
choice of $p$. As an example, we provide the same illustration as for Figure 1
in Figure 8.

Figure 8: Illustration of our differentiable and sparse top-k mask. Same setup
as for Figure 1, with more values for $p$.

##  8 Conclusion

In this work, we proposed a generalized framework to obtain fast,
differentiable (or differentiable a.e.) and sparse top-$k$ and top-$k$ masks
operators, including operators that select values in magnitude. Thanks to a
reduction to isotonic optimization, we showed that these operators can be
computed using either the Pool Adjacent Violators (PAV) algorithm or Dykstra’s
projection algorithm, the latter being faster on TPU hardware. We successfully
demonstrated the usefulness of our operators for weight pruning, top-$k$
losses and as routers in vision sparse mixture of experts.

#### Acknowledgments.

We thank Vincent Roulet and Joelle Barral for comments on a draft of this
paper. We thank Felipe Llinares-López for helpful feedbacks regarding the
experiments, as well as Fabian Pedregosa for fruitful mathematical
discussions. We also thank the anonymous reviewers for their feedback.

## References

  * Adams & Zemel (2011) Adams, R. P. and Zemel, R. S.  Ranking via sinkhorn propagation.  _arXiv e-prints_ , 2011.
  * Amos et al. (2019) Amos, B., Koltun, V., and Kolter, J. Z.  The limited multi-label projection layer.  _arXiv preprint arXiv:1906.08707_ , 2019.
  * Argyriou et al. (2012) Argyriou, A., Foygel, R., and Srebro, N.  Sparse prediction with the $k$-support norm.  _Advances in Neural Information Processing Systems_ , 25, 2012.
  * Beck (2017) Beck, A.  _First-order methods in optimization_.  SIAM, 2017.
  * Berrada et al. (2018) Berrada, L., Zisserman, A., and Kumar, M. P.  Smooth loss functions for deep top-k classification.  _International Conference on Learning Representations_ , 2018.
  * Berthet et al. (2020) Berthet, Q., Blondel, M., Teboul, O., Cuturi, M., Vert, J.-P., and Bach, F.  Learning with differentiable pertubed optimizers.  _Advances in neural information processing systems_ , 33:9508–9519, 2020.
  * Best et al. (2000) Best, M. J., Chakravarti, N., and Ubhaya, V. A.  Minimizing separable convex functions subject to simple chain constraints.  _SIAM Journal on Optimization_ , 10(3):658–672, 2000.
  * Blalock et al. (2020) Blalock, D., Gonzalez Ortiz, J. J., Frankle, J., and Guttag, J.  What is the state of neural network pruning?  _Proceedings of machine learning and systems_ , 2:129–146, 2020.
  * Blondel (2019) Blondel, M.  Structured prediction with projection oracles.  _Advances in neural information processing systems_ , 32, 2019.
  * Blondel et al. (2020a) Blondel, M., Martins, A. F., and Niculae, V.  Learning with fenchel-young losses.  _J. Mach. Learn. Res._ , 21(35):1–69, 2020a.
  * Blondel et al. (2020b) Blondel, M., Teboul, O., Berthet, Q., and Djolonga, J.  Fast differentiable sorting and ranking.  In _International Conference on Machine Learning_ , pp. 950–959. PMLR, 2020b.
  * Blondel et al. (2021) Blondel, M., Berthet, Q., Cuturi, M., Frostig, R., Hoyer, S., Llinares-López, F., Pedregosa, F., and Vert, J.-P.  Efficient and modular implicit differentiation.  _arXiv preprint arXiv:2105.15183_ , 2021.
  * Bowman (1972) Bowman, V.  Permutation polyhedra.  _SIAM Journal on Applied Mathematics_ , 22(4):580–589, 1972.
  * Boyle & Dykstra (1986) Boyle, J. P. and Dykstra, R. L.  A method for finding projections onto the intersection of convex sets in hilbert spaces.  In _Advances in order restricted statistical inference_ , pp. 28–47. Springer, 1986.
  * Bradbury et al. (2018) Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q.  JAX: composable transformations of Python+NumPy programs, 2018\.  URL http://github.com/google/jax.
  * Chapelle & Wu (2010) Chapelle, O. and Wu, M.  Gradient descent optimization of smoothed information retrieval metrics.  _Information retrieval_ , 13(3):216–235, 2010\.
  * Collins & Kohli (2014) Collins, M. D. and Kohli, P.  Memory bounded deep convolutional networks.  _arXiv preprint arXiv:1412.1442_ , 2014.
  * Combettes & Pesquet (2011) Combettes, P. L. and Pesquet, J.-C.  Proximal splitting methods in signal processing.  In _Fixed-point algorithms for inverse problems in science and engineering_ , pp. 185–212. Springer, 2011.
  * Cordonnier et al. (2021) Cordonnier, J.-B., Mahendran, A., Dosovitskiy, A., Weissenborn, D., Uszkoreit, J., and Unterthiner, T.  Differentiable patch selection for image recognition.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 2351–2360, 2021.
  * Cuturi (2013) Cuturi, M.  Sinkhorn distances: Lightspeed computation of optimal transport.  _Advances in neural information processing systems_ , 26, 2013.
  * Cuturi et al. (2019) Cuturi, M., Teboul, O., and Vert, J.-P.  Differentiable ranking and sorting using optimal transport.  _Advances in neural information processing systems_ , 32, 2019.
  * Danskin (1966) Danskin, J. M.  The theory of max-min, with applications.  _SIAM Journal on Applied Mathematics_ , 14(4):641–664, 1966.
  * Dosovitskiy et al. (2020) Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv preprint arXiv:2010.11929_ , 2020.
  * Eriksson et al. (2015) Eriksson, A., Thanh Pham, T., Chin, T.-J., and Reid, I.  The k-support norm and convex envelopes of cardinality and rank.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , pp. 3349–3357, 2015.
  * Fedus et al. (2021) Fedus, W., Zoph, B., and Shazeer, N.  Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021.
  * Fedus et al. (2022) Fedus, W., Dean, J., and Zoph, B.  A review of sparse expert models in deep learning.  _arXiv preprint arXiv:2209.01667_ , 2022.
  * Frankle & Carbin (2018) Frankle, J. and Carbin, M.  The lottery ticket hypothesis: Finding sparse, trainable neural networks.  _arXiv preprint arXiv:1803.03635_ , 2018.
  * Gale et al. (2019) Gale, T., Elsen, E., and Hooker, S.  The state of sparsity in deep neural networks.  _arXiv preprint arXiv:1902.09574_ , 2019.
  * Grover et al. (2019) Grover, A., Wang, E., Zweig, A., and Ermon, S.  Stochastic optimization of sorting networks via continuous relaxations.  In _International Conference on Learning Representations_ , 2019.
  * Han et al. (2015) Han, S., Pool, J., Tran, J., and Dally, W.  Learning both weights and connections for efficient neural network.  _Advances in neural information processing systems_ , 28, 2015.
  * Hazimeh et al. (2021) Hazimeh, H., Zhao, Z., Chowdhery, A., Sathiamoorthy, M., Chen, Y., Mazumder, R., Hong, L., and Chi, E.  Dselect-k: Differentiable selection in the mixture of experts with applications to multi-task learning.  _Advances in Neural Information Processing Systems_ , 34:29335–29347, 2021.
  * Järvelin & Kekäläinen (2017) Järvelin, K. and Kekäläinen, J.  Ir evaluation methods for retrieving highly relevant documents.  In _ACM SIGIR Forum_ , volume 51, pp. 243–250. ACM New York, NY, USA, 2017.
  * Jegelka et al. (2013) Jegelka, S., Bach, F., and Sra, S.  Reflection methods for user-friendly submodular optimization.  _Advances in Neural Information Processing Systems_ , 26, 2013.
  * Krizhevsky et al. (2009) Krizhevsky, A., Hinton, G., et al.  Learning multiple layers of features from tiny images.  2009\.
  * Kyrillidis et al. (2013) Kyrillidis, A., Becker, S., Cevher, V., and Koch, C.  Sparse projections onto the simplex.  In _International Conference on Machine Learning_ , pp. 235–243. PMLR, 2013.
  * Lapin et al. (2015) Lapin, M., Hein, M., and Schiele, B.  Top-k multiclass svm.  _Advances in Neural Information Processing Systems_ , 28, 2015.
  * Lapin et al. (2016) Lapin, M., Hein, M., and Schiele, B.  Loss functions for top-k error: Analysis and insights.  In _Proc. of CVPR_ , 2016.
  * Lewis et al. (2021) Lewis, M., Bhosale, S., Dettmers, T., Goyal, N., and Zettlemoyer, L.  Base layers: Simplifying training of large, sparse models.  In _International Conference on Machine Learning_ , pp. 6265–6274. PMLR, 2021.
  * Lim & Wright (2016) Lim, C. H. and Wright, S. J.  Efficient bregman projections onto the permutahedron and related polytopes.  In _Artificial Intelligence and Statistics_ , pp. 1205–1213. PMLR, 2016.
  * Liu et al. (2022) Liu, T., Puigcerver, J., and Blondel, M.  Sparsity-constrained optimal transport.  _arXiv preprint arXiv:2209.15466_ , 2022.
  * Malaviya et al. (2018) Malaviya, C., Ferreira, P., and Martins, A. F.  Sparse and constrained attention for neural machine translation.  _arXiv preprint arXiv:1805.08241_ , 2018.
  * Martins & Kreutzer (2017) Martins, A. F. and Kreutzer, J.  Learning what’s easy: Fully differentiable neural easy-first taggers.  In _Proceedings of the 2017 conference on empirical methods in natural language processing_ , pp. 349–362, 2017.
  * McDonald et al. (2014) McDonald, A. M., Pontil, M., and Stamos, D.  Spectral k-support norm regularization.  _Advances in neural information processing systems_ , 27, 2014.
  * Németh & Németh (2012) Németh, A. and Németh, S.  How to project onto the monotone nonnegative cone using pool adjacent violators type algorithms.  _arXiv preprint arXiv:1201.2343_ , 2012.
  * Niepert et al. (2021) Niepert, M., Minervini, P., and Franceschi, L.  Implicit mle: backpropagating through discrete exponential family distributions.  _Advances in Neural Information Processing Systems_ , 34:14567–14579, 2021.
  * Petersen et al. (2021) Petersen, F., Borgelt, C., Kuehne, H., and Deussen, O.  Differentiable sorting networks for scalable sorting and ranking supervision.  In _International Conference on Machine Learning_ , pp. 8546–8555. PMLR, 2021.
  * Petersen et al. (2022) Petersen, F., Kuehne, H., Borgelt, C., and Deussen, O.  Differentiable top-k classification learning.  In _International Conference on Machine Learning_ , pp. 17656–17668. PMLR, 2022.
  * Prillo & Eisenschlos (2020) Prillo, S. and Eisenschlos, J.  Softsort: A continuous relaxation for the argsort operator.  In _International Conference on Machine Learning_ , pp. 7793–7802. PMLR, 2020.
  * Qian et al. (2022) Qian, Y., Lee, J., Duddu, S. M. K., Dai, Z., Brahma, S., Naim, I., Lei, T., and Zhao, V. Y.  Multi-vector retrieval as sparse alignment.  _arXiv preprint arXiv:2211.01267_ , 2022.
  * Riquelme et al. (2021) Riquelme, C., Puigcerver, J., Mustafa, B., Neumann, M., Jenatton, R., Susano Pinto, A., Keysers, D., and Houlsby, N.  Scaling vision with sparse mixture of experts.  _Advances in Neural Information Processing Systems_ , 34:8583–8595, 2021.
  * Rolínek et al. (2020) Rolínek, M., Musil, V., Paulus, A., Vlastelica, M., Michaelis, C., and Martius, G.  Optimizing rank-based metrics with blackbox differentiation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 7620–7630, 2020.
  * Shazeer et al. (2017) Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., and Dean, J.  Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.  _arXiv preprint arXiv:1701.06538_ , 2017.
  * Sinkhorn (1967) Sinkhorn, R.  Diagonal equivalence to matrices with prescribed row and column sums.  _The American Mathematical Monthly_ , 74(4):402–405, 1967.
  * Sun et al. (2017) Sun, C., Shrivastava, A., Singh, S., and Gupta, A.  Revisiting unreasonable effectiveness of data in deep learning era.  In _Proceedings of the IEEE international conference on computer vision_ , pp. 843–852, 2017.
  * Tibshirani (2017) Tibshirani, R. J.  Dykstra’s algorithm, admm, and coordinate descent: Connections, insights, and extensions.  _Advances in Neural Information Processing Systems_ , 30, 2017.
  * Warmuth & Kuzmin (2008) Warmuth, M. K. and Kuzmin, D.  Randomized online pca algorithms with regret bounds that are logarithmic in the dimension.  _Journal of Machine Learning Research_ , 9(Oct):2287–2320, 2008.
  * Wiseman & Rush (2016) Wiseman, S. and Rush, A. M.  Sequence-to-sequence learning as beam-search optimization.  _arXiv preprint arXiv:1606.02960_ , 2016.
  * Xie et al. (2020) Xie, Y., Dai, H., Chen, M., Dai, B., Zhao, T., Zha, H., Wei, W., and Pfister, T.  Differentiable top-k with optimal transport.  _Advances in Neural Information Processing Systems_ , 33:20520–20531, 2020.
  * Zeng & Figueiredo (2014) Zeng, X. and Figueiredo, M. A.  The ordered weighted $l_{1}$ norm: Atomic formulation, projections, and algorithms.  _arXiv preprint arXiv:1409.4271_ , 2014.
  * Zhou et al. (2022) Zhou, Y., Lei, T., Liu, H., Du, N., Huang, Y., Zhao, V., Dai, A., Chen, Z., Le, Q., and Laudon, J.  Mixture-of-experts with expert choice routing.  _arXiv preprint arXiv:2202.09368_ , 2022.
  * Ziegler (2012) Ziegler, G. M.  _Lectures on polytopes_ , volume 152.  Springer Science & Business Media, 2012.

##  Appendix A Proofs

###  A.1 Linear Maximization Oracle - Proof of Proposition 8

Assuming $\bm{w}$ is sorted in descending order, we have for any
$\bm{x}\in\mathbb{R}^{n}$

$\displaystyle\sum_{i=1}^{n}w_{i}x_{[i]}$
$\displaystyle=\max_{\sigma\in\Sigma}\langle\bm{x}_{\sigma},\bm{w}\rangle$
(33) $\displaystyle=\max_{\pi\in\Sigma}\langle\bm{x},\bm{w}_{\pi}\rangle$
$\displaystyle=\max_{\bm{y}\in P(\bm{w})}\langle\bm{x},\bm{y}\rangle.$

In the first line, we used that the inner product is maximized by finding the
permutation $\sigma$ sorting $\bm{x}$ in descending order. In the second line,
we used that
$\langle\bm{x}_{\sigma},\bm{w}\rangle=\langle\bm{x},\bm{w}_{\pi}\rangle$, if
$\pi$ is the inverse permutation of $\sigma$. In the third line, we used the
fundamental theorem of linear programming, which guarantees that the solution
happens at one of the vertices of the polytope. To summarize, if $\sigma$ is
the permutation sorting $\bm{x}$ in descending order, then
$\bm{y}^{\star}=\bm{w}_{\sigma^{-1}}$.

###  A.2 Relaxed operator - Proof of Proposition 16

Recall that $f^{*}_{\varphi,R}(\bm{y},\bm{w})\coloneqq
f^{*}_{\varphi}(\bm{y},\bm{w})+R(\bm{y})$. We then have

$f_{\varphi,R}(\bm{x},\bm{w})=\max_{\bm{y}\in\mathbb{R}^{n}}\langle\bm{y},\bm{x}\rangle-
f_{\varphi}^{*}(\bm{y},\bm{w})-R(\bm{y}).$ (34)

It is well-known that if $h_{1}$ and $h_{2}$ are two convex functions, then
$(h_{1}+h_{2})^{*}$ is equal to the infimal convolution of $h_{1}^{*}$ with
$h_{2}^{*}$ (Beck, 2017, Theorem 4.17):

$(h_{1}+h_{2})^{*}(\bm{x})=(h_{1}^{*}\square
h_{2}^{*})(\bm{x})\coloneqq\min_{\bm{u}\in\mathbb{R}^{n}}h_{1}^{*}(\bm{u})+h_{2}^{*}(\bm{x}-\bm{u}).$
(35)

With $h_{1}=f_{\varphi}^{*}$ and $h_{2}=R$, we therefore get

$f_{\varphi,R}(\bm{x},\bm{w})=\min_{\bm{u}\in\mathbb{R}^{n}}R^{*}(\bm{x}-\bm{u})+f_{\varphi}(\bm{u},\bm{w}).$
(36)

Finally, the expression of $\bm{y}^{\star}$ follows from Danskin’s theorem
applied.

###  A.3 Conjugate - Proof of Proposition 18

We have

$\displaystyle f^{*}_{\varphi}(\bm{y},\bm{w})$
$\displaystyle=\max_{\bm{x}\in\mathbb{R}^{n}}\langle\bm{x},\bm{y}\rangle-
f_{\varphi}(\bm{x},\bm{w})$ (37)
$\displaystyle=\max_{\bm{x}\in\mathbb{R}^{n}}\langle\bm{x},\bm{y}\rangle-\max_{\bm{y}^{\prime}\in
P(\bm{w})}~{}\langle\varphi(\bm{x}),\bm{y}^{\prime}\rangle$
$\displaystyle=\max_{\bm{x}\in\mathbb{R}^{n}}\min_{\bm{y}^{\prime}\in
P(\bm{w})}\langle\bm{x},\bm{y}\rangle-\langle\varphi(\bm{x}),\bm{y}^{\prime}\rangle.$

If $\bm{w}\in\mathbb{R}_{+}^{n}$, then $\bm{y}^{\prime}\in\mathbb{R}^{n}_{+}$
for all $\bm{y}^{\prime}\in P(\bm{w})$. Then the function
$(\bm{x},\bm{y}^{\prime})\mapsto\langle\bm{x},\bm{y}\rangle-\langle\varphi(\bm{x}),\bm{y}^{\prime}\rangle$
is concave-convex and we can switch the min and the max to obtain

$\displaystyle f^{*}_{\varphi}(\bm{y},\bm{w})$
$\displaystyle=\min_{\bm{y}^{\prime}\in
P(\bm{w})}\max_{\bm{x}\in\mathbb{R}^{n}}\langle\bm{x},\bm{y}\rangle-\langle\varphi(\bm{x}),\bm{y}^{\prime}\rangle$
(38) $\displaystyle=\min_{\bm{y}^{\prime}\in
P(\bm{w})}\sum_{i=1}^{n}y^{\prime}_{i}\phi_{i}^{*}(y_{i}/y^{\prime}_{i}).$

###  A.4 Biconjugate interpretation - Proof of Proposition 4

One has

$\Phi^{*}_{k}(\bm{x})=\max_{\bm{y}\in
S_{k}}\langle\bm{x},\bm{y}\rangle-\sum^{n}_{i=1}\phi^{*}(y_{i}).$

As in (Kyrillidis et al., 2013), let $\Sigma_{k}$ be the set of subsets of
$[n]$ with cardinality smaller than $k$. Then

$\Phi^{*}_{k}(\bm{x})=\max_{I\subset\Sigma_{k}}\max_{\bm{y}\in\mathbb{R}^{n}}\sum_{i\in
I}x_{i}y_{i}-\phi^{*}(y_{i}).$

This gives

$\Phi^{*}_{k}(\bm{x})=\max_{I\subset\Sigma_{k}}\sum_{i\in
I}\phi(x_{i})=\sum_{i=1}^{k}\varphi(\bm{x})_{[i]}.$

Taking the conjugate gives the desired result.

###  A.5 Reduction to isotonic optimization

We focus on the case when $\varphi$ is sign-invariant, i.e.,
$\varphi(\bm{x})=\varphi(-\bm{x})$, since the case $\varphi(\bm{x})=\bm{x}$ is
already tackled in (Lim & Wright, 2016; Blondel et al., 2020b).

We first show that $\bm{u}^{\star}$ preserves the sign of $\bm{x}$. We do so
by showing that for any $\bm{u}\in\mathbb{R}^{n}$,
$\bm{u}^{\prime}\coloneqq\mathrm{sign}(\bm{x})\circ|\bm{u}|$ achieves smaller
objective value than $\bm{u}$. Recall that

$f_{\varphi,R}(\bm{x},\bm{w})=\min_{\bm{u}\in\mathbb{R}^{n}}R^{*}(\bm{x}-\bm{u})+f_{\varphi}(\bm{u},\bm{w})=\min_{\bm{u}\in\mathbb{R}^{n}}R^{*}(\bm{x}-\bm{u})+f(\varphi(\bm{u}),\bm{w}).$
(39)

Clearly, we have
$f(\varphi(\bm{u}^{\prime}),\bm{w})=f(\varphi(\bm{u}),\bm{w})$. Moreover, if
$R(\bm{y})=\sum_{i=1}^{n}r(y_{i})$, then
$R^{*}(\bm{x}-\bm{u})=\sum_{i=1}^{n}r^{*}(x_{i}-u_{i})$. If $r^{*}$ is sign-
invariant and increasing on $\mathbb{R}_{+}$, we then have

$\displaystyle r^{*}(x_{i}-u^{\prime}_{i})$
$\displaystyle=r^{*}(\mathrm{sign}(x_{i})(|x_{i}|-|u_{i}|))$ (40)
$\displaystyle=r^{*}(|x_{i}|-|u_{i}|)$ $\displaystyle\leq r^{*}(x_{i}-u_{i}),$

where we used the reverse triangle inequality
$||x_{i}-|u_{i}||\leq|x_{i}-u_{i}|$. We conclude that $\bm{u}^{\star}$ has the
same sign as $\bm{x}$. From now on, we can therefore assume that
$\bm{x}\in\mathbb{R}_{+}^{n}$, which implies that
$\bm{u}\in\mathbb{R}^{n}_{+}$.

Since $\bm{u}\in\mathbb{R}^{n}_{+}$ and $\phi$ is increasing on
$\mathbb{R}_{+}$, we have $u_{\sigma_{1}}\geq\dots\geq
u_{\sigma_{n}}\Rightarrow\phi(u_{\sigma_{1}})\geq\dots\geq\phi(u_{\sigma_{n}})$.
We know that for all $\bm{u}\in\mathbb{R}^{n}$ and $\bm{w}\in\mathbb{R}^{n}$,
we have
$f(\varphi(u),\bm{w})=\langle\varphi(\bm{u})_{\sigma},\bm{w}\rangle=\langle\varphi(\bm{u}_{\sigma}),\bm{w}\rangle$,
where $\sigma$ is the permutation sorting $\bm{u}$ in descending order. From
now on, let us fix $\sigma$ to the permutation sorting $\bm{u}^{\star}$. We
will show in the sequel that this is the same permutation as the one sorting
$\bm{x}$. We then have

$f_{\varphi,R}(\bm{x},\bm{w})=\min_{\bm{u}\in\mathbb{R}^{n}}R^{*}(\bm{x}-\bm{u})+\langle\varphi(\bm{u}_{\sigma}),\bm{w}\rangle.$
(41)

Using the change of variable
$\bm{v}=\bm{u}_{\sigma}\Leftrightarrow\bm{v}_{\sigma^{-1}}=\bm{u}$, we obtain

$f_{\varphi,R}(\bm{x},\bm{w})=\max_{v_{1}\geq\dots\geq
v_{n}}R^{*}(\bm{x}_{\sigma}-\bm{v})+\langle\varphi(\bm{v}),\bm{w}\rangle$ (42)

where we used that if $R(\bm{y})=\sum_{i=1}^{n}r(y_{i})$, then

$R^{*}(\bm{x}-\bm{u})=R^{*}(\bm{x}-\bm{v}_{\sigma^{-1}})=R^{*}(\bm{x}_{\sigma}-\bm{v}).$
(43)

Let $\bm{s}\coloneqq\bm{x}_{\sigma}$. It remains to show that
$s_{1}\geq\dots\geq s_{n}$, i.e., that $\bm{s}$ and $\bm{v}^{\star}$ are both
in descending order. Suppose $s_{j}>s_{i}$ for some $i<j$. Let
$\bm{s}^{\prime}$ be a copy of $\bm{s}$ with $s_{i}$ and $s_{j}$ swapped.
Since $R^{*}$ is convex, by (Blondel et al., 2020b, Lemma 4),

$R^{*}(\bm{s}-\bm{v}^{\star})-R^{*}(\bm{s}^{\prime}-\bm{v}^{\star})=r^{*}(s_{i}-v_{i}^{\star})+r^{*}(s_{j}-v_{j}^{\star})-r^{*}(s_{j}-v_{i}^{\star})-r^{*}(s_{i}-v_{j}^{\star})\geq
0,$ (44)

which contradicts the assumption that $\bm{v}^{\star}$ and the corresponding
$\sigma$ are optimal.

###  A.6 Subproblem derivation

#### Case $\phi(x)=\frac{1}{2}x^{2}$ and
$R(\bm{x})=\frac{\lambda}{2}\|\bm{x}\|^{2}$.

One has
$h_{i}(\gamma)=\frac{1}{2}(w_{i}\gamma^{2}+\frac{1}{\lambda}(s_{i}-\gamma)^{2})$
so that
$\frac{dh_{i}}{d\gamma}=(w_{i}+\frac{1}{\lambda})\gamma-\frac{1}{\lambda}s_{i}.$
Therefore,

$\frac{d\sum_{i\in B}h_{i}(\gamma)}{d\gamma}=\gamma\sum_{i\in
B}(w_{i}+\frac{1}{\lambda})-\frac{1}{\lambda}\sum_{i\in B}s_{i}.$

Since in addition $\sum_{i\in B}h_{i}$ is convex we obtain that its minimum is
given by canceling the derivative, hence

$\gamma_{B}^{\star}=\frac{\sum_{i\in B}s_{i}}{\sum_{i\in B}(\lambda
w_{i}+1)}.$

#### Remark.

This result shows that when $\lambda$ is small enough, one can control the
approximation error induced by our proposed operator in comparison to the hard
operator. For simplicity, let us focus on the case where there are no ties:
$\forall i\neq j$, $x_{i}\neq x_{j}$. This implies that
$s_{1}>s_{2}>\cdots>s_{n}$. In this case, for $\lambda$ small enough, we get

$\gamma_{\\{1\\}}^{\star}>\gamma_{\\{2\\}}^{\star}>\cdots>\gamma_{\\{n\\}}^{\star}$

so that the optimal partition in PAV’s algorithm is given by taking
$B_{i}=\\{i\\}$. Therefore, $v_{i}^{\star}=\frac{s_{i}}{\lambda w_{i}+1}.$ One
then has

$\bm{u}^{\star}=\mathrm{sign}(\bm{x})\circ\bm{v}^{\star}_{\sigma^{-1}}=\frac{\bm{x}}{\lambda\bm{w}_{\sigma^{-1}}+1}.$

Plugging it into $\bm{y}^{\star}$ given by Proposition 16 gives

$\bm{y}^{\star}=\frac{\bm{w}_{\sigma^{-1}}\circ\bm{x}}{\lambda\bm{w}_{\sigma^{-1}}+1}.$

Since the hard operator is given by $\bm{w}_{\sigma^{-1}}\circ\bm{x}$, the
approximation error $e_{\lambda}(\bm{x},\bm{w})$ in infinite norm is then
simply bounded by

$e_{\lambda}(\bm{x},\bm{w}):=\|\bm{w}_{\sigma^{-1}}\circ\bm{x}(\frac{1}{\lambda\bm{w}_{\sigma^{-1}}+1}-1)\|_{\infty}\leq\|\bm{w}\|_{\infty}\|\bm{x}\|_{\infty}\|\frac{\lambda\bm{w}}{\lambda\bm{w}+1}\|_{\infty}.$

In the $\mathrm{topkmag}$ case, $w=\mathbf{1}_{k}$, so that
$e_{\lambda}(\bm{x},\bm{w})\leq\lambda\|\bm{x}\|_{\infty}$.

#### Case $\phi(x)=x$ and $R(\bm{x})=\frac{\lambda}{p}\|\bm{x}\|^{p}$ with
$p=\frac{4}{3}$.

One has $h_{i}(\gamma)=w_{i}\gamma+\frac{1}{4\lambda^{3}}(s_{i}-\gamma)^{4}$
so that $\frac{dh_{i}}{d\gamma}=w_{i}+\frac{1}{\lambda^{3}}(\gamma-
s_{i})^{3}.$ Therefore,

$\frac{d\sum_{i\in B}h_{i}(\gamma)}{d\gamma}=\frac{1}{\lambda^{3}}\sum_{i\in
B}(\gamma-s_{i})^{3}+\sum_{i\in B}w_{i}.$

Since in addition $\sum_{i\in B}h_{i}$ is convex we obtain that its minimum is
given by canceling the derivative, and hence by solving the third-order
polynomial equation

$\frac{1}{\lambda^{3}}\sum_{i\in B}(\gamma-s_{i})^{3}+\sum_{i\in B}w_{i}=0.$

In practice, we solve this equation using the root solver from the numpy
library. Note that taking $p=\frac{4}{3}$ leads to an easier subproblem than
$p=\frac{3}{2}$, hence our choice for $p$.

#### Case $\phi(x)=\frac{1}{2}x^{2}$ and
$R(\bm{x})=\frac{\lambda}{p}\|\bm{x}\|^{p}$ with $p=\frac{4}{3}$.

The derivation is very similar to the previous case. Indeed, one has
$h_{i}(\gamma)=\frac{1}{2}w_{i}\gamma^{2}+\frac{1}{4\lambda^{3}}(s_{i}-\gamma)^{4}$
so that $\frac{dh_{i}}{d\gamma}=w_{i}\gamma+\frac{1}{\lambda^{3}}(\gamma-
s_{i})^{3}.$ Therefore,

$\frac{d\sum_{i\in B}h_{i}(\gamma)}{d\gamma}=\frac{1}{\lambda^{3}}\sum_{i\in
B}(\gamma-s_{i})^{3}+\gamma\sum_{i\in B}w_{i}.$

###  A.7 Differentiation - Proof of Proposition 5

#### Case $\phi(x)=x$.

One has $h_{i}(\gamma)=\frac{1}{q}|s_{i}-\gamma|^{q}+w_{i}v_{i}.$ For any
optimal $B$ in PAV one has the optimality condition

$\sum_{i\in
B}(\mathrm{sign}(\gamma^{\star}_{B}-s_{i})|\gamma^{\star}_{B}-s_{i}|^{q-1}+w_{i})=0.$

Using the implicit function theorem gives that $\gamma^{\star}_{B}$ is
differentiable, and differentiating with respect to any $s_{i}$ for $i\in B$
leads to

$\frac{\partial\gamma^{\star}_{B}}{\partial
s_{i}}=\frac{|\gamma^{\star}_{B}-s_{i}|^{q-2}}{\sum_{j\in
B}|\gamma^{\star}_{B}-s_{j}|^{q-2}}.$

#### Case $\phi(x)=\frac{1}{2}x^{2}$.

Similar calculations lead to

$\frac{\partial\gamma^{\star}_{B}}{\partial
s_{i}}=\frac{(q-1)|\gamma^{\star}_{B}-s_{i}|^{q-2}}{\sum_{j\in
B}(q-1)|\gamma^{\star}_{B}-s_{j}|^{q-2}+w_{j}}.$

###  A.8 Dual of isotonic optimization

$\displaystyle\min_{v_{1}\geq\dots\geq v_{n}}\sum_{i=1}^{n}h_{i}(v_{i})$
$\displaystyle=\min_{\bm{v}\in\mathbb{R}^{n}}\max_{\bm{\alpha}\in\mathbb{R}^{n-1}_{+}}\sum_{i=1}^{n}h_{i}(v_{i})-\alpha_{i}(v_{i}-v_{i+1})$
(45)
$\displaystyle=\min_{\bm{v}\in\mathbb{R}^{n}}\max_{\bm{\alpha}\in\mathbb{R}^{n-1}_{+}}\sum_{i=1}^{n}h_{i}(v_{i})-v_{i}(\alpha_{i}-\alpha_{i-1})$
$\displaystyle=\max_{\bm{\alpha}\in\mathbb{R}^{n-1}_{+}}-\left[\sum_{i=1}^{n}h^{*}_{i}(\alpha_{i}-\alpha_{i-1})\right]$

where $\alpha_{0}\coloneqq 0$ and $\alpha_{n}\coloneqq 0$ are constants (i.e.,
not optimized). An optimal solution $\bm{v}^{\star}$ is recovered from
$\bm{\alpha}^{\star}$ by
$v^{\star}_{i}=(h_{i}^{*})^{\prime}(\alpha_{i}^{\star}-\alpha_{i-1}^{\star})$.
Since $\bm{\alpha}$ is only constrained to be non-negative, we can solve the
dual by coordinate ascent. The subproblem associated with $\alpha_{i}$, for
$i\in\\{1,\dots,n-1\\}$, is

$\max_{\alpha_{i}\in\mathbb{R}_{+}}-h_{i}^{*}(\alpha_{i}-\alpha_{i-1})-h_{i+1}^{*}(\alpha_{i+1}-\alpha_{i}).$
(46)

The subproblem is a simple univariate problem with non-negative constraint.
Let us define $\alpha_{i}^{\prime}$ as the solution of
$(h_{i}^{*})^{\prime}(\alpha_{i}^{\prime}-\alpha_{i-1})-(h_{i+1}^{*})^{\prime}(\alpha_{i+1}-\alpha_{i}^{\prime})=0$.
The solution is then $\alpha_{i}^{\star}=[\alpha_{i}^{\prime}]_{+}$.

In practice, we can alternate between updating $\alpha_{1},\alpha_{3},\dots$
in parallel and $\alpha_{2},\alpha_{4},\dots$ in parallel. The dual variables
with odd coordinates correspond to the set
$C_{1}=\\{\bm{v}\in\mathbb{R}^{n}\colon v_{1}\geq v_{2},v_{3}\geq
v_{4},\dots\\}$ and the dual variables with even coordinates correspond to the
set $C_{2}=\\{\bm{v}\in\mathbb{R}^{n}\colon v_{2}\geq v_{3},v_{4}\geq
v_{5},\dots\\}$. Coordinate ascent converges to an optimal dual solution,
assuming each $h_{i}^{*}$ is differentiable, which is equivalent to each
$h_{i}$ being strictly convex.

Note that the subproblem can be rewritten in primal space as

$\displaystyle\max_{\alpha_{i}\in\mathbb{R}_{+}}-h_{i}^{*}(\alpha_{i}-\alpha_{i-1})-h_{i+1}^{*}(\alpha_{i+1}-\alpha_{i})$
(47)
$\displaystyle=\max_{\alpha_{i}\in\mathbb{R}_{+}}-\left[\max_{v_{i}}(\alpha_{i}-\alpha_{i-1})v_{i}-h_{i}(v_{i})\right]-\left[\max_{v_{i+1}}(\alpha_{i+1}-\alpha_{i})v_{i+1}-h_{i+1}(v_{i+1})\right]$
$\displaystyle=\min_{v_{i},v_{i+1}}\alpha_{i-1}v_{i}+h_{i}(v_{i})-\alpha_{i+1}v_{i+1}+h_{i+1}(v_{i+1})+\max_{\alpha_{i}\in\mathbb{R}_{+}}\alpha_{i}(v_{i+1}-v_{i})$
$\displaystyle=\min_{v_{i}\geq
v_{i+1}}h_{i}(v_{i})+h_{i+1}(v_{i+1})+\alpha_{i-1}v_{i}-\alpha_{i+1}v_{i+1}.$

In fact, in the Euclidean case, it is known that Dykstra’s algorithm in the
primal and block coordinate ascent in the dual are equivalent (Tibshirani,
2017). Therefore, block coordinate ascent can be seen as an elegant way to
generalize Dykstra’s algorithm to the non-Euclidean case.

##  Appendix B Additional material

###  B.1 k-support negentropies

When $\phi(x)=e^{x-1}$ and $\bm{w}=\mathbf{1}_{k}$, we obtain

$f_{\varphi}^{*}(\bm{y},\bm{w})=\min_{\bm{z}\in[0,1]^{n}}\sum_{i=1}^{n}y_{i}\log(\frac{y_{i}}{z_{i}})\quad\text{s.t.}\quad\langle\bm{z},\mathbf{1}\rangle=k.$
(48)

We call it a $k$-support negative entropy.

Figure 9: Contours of the k-support entropy on the simplex for $n=3$ and
$k=2$. Lighter colors indicate lower values.

###  B.2 PAV algorithm

We present the pseudo code for PAV, adapted from Lim & Wright (2016). Recall
that we define

$\gamma^{\star}_{B}=\operatorname*{argmin}_{\gamma\in\mathbb{R}}\sum_{i\in
B}h_{i}(\gamma).$ (49)

Algorithm 1 Pool Adjacent Violators (PAV)

Input: Convex functions $\\{h_{i}$ :
$\mathbb{R}\to\mathbb{R}\\}_{i\in[n]}$Initalize partitions
$P\leftarrow\\{\\{i\\}|i\in[n]\\}$Initialize
$v_{i}\leftarrow\gamma_{\\{i\\}}^{\star}$ for all $i\in[n]$

while there exists $i$ such that $v_{i}<{v_{i+1}}$ do

Find $B_{r_{i}}$ and $B_{r_{i+1}}$ in $P$ such that $i\in B_{r_{i}}$ and
$i+1\in B_{r_{i+1}}$ Remove $B_{r_{i}}$ and $B_{r_{i+1}}$ from $P$ Add
$B_{r_{i}}\cup B_{r_{i+1}}$ to $P$ Compute $\gamma_{B_{r_{i}}\cup
B_{r_{i+1}}}^{\star}$ by solving (28) Assign
$v_{r}\leftarrow\gamma_{B_{r_{i}}\cup B_{r_{i+1}}}^{\star}$ for all $r\in
B_{r_{i}}\cup B_{r_{i+1}}$

end while

return $v$

##  Appendix C Experimental details

###  C.1 Weight pruning in neural networks

For our experiment on the MNIST dataset, we train the MLP using SGD with a
batch size of $128$ and a constant learning rate of $10^{-2}$. We trained the
model for $30$ epochs. In terms of hardware, we use a single GPU.

###  C.2 Smooth top-k loss

For our experiment on the CIFAR-100 dataset, we train the ViT-B/16 using SGD
with a momentum of 0.9 and with a batch size of $512$.

For the cross-entropy loss, we follow the training procedure of Dosovitskiy et
al. (2020): warmup phase until the learning rate reaches $3\times 10^{-3}$.
The model is trained for $100$ steps using a cosine learning rate scheduler.
This choice of learning rate gave the best performance for this number of
training steps.

For our top-$3$ losses: warmup phase until the learning rate reaches $5\times
10^{-3}$. The model is trained for $100$ steps using a cosine learning rate
scheduler.

In terms of hardware, we use $8$ TPUs.

###  C.3 Sparse MoEs

We train the V-MoE S/32 model (Riquelme et al., 2021) on the JFT-300M dataset
(Sun et al., 2017). JFT is a multilabel dataset, and thus accuracy is not an
appropriate metric since each image may have multiple labels. Therefore, we
measure the quality of the models using the commonly-used _precision-at-1_
metric (Järvelin & Kekäläinen, 2017). The training procedure is analogous to
the one described in Riquelme et al. (2021), except that we replace the
routing algorithm. In particular, we use the Expert Choice Routing algorithm
described in Zhou et al. (2022) as our baseline, and replace the non-
differentiable top-$k$ operation used there with our differentiable approach
(we perform 10 iterations of Dykstra’s algorithm).

We use exactly the same hyperparameters as described in Riquelme et al.
(2021), except for the fact that Expert Choice Routing does not require any
auxiliary loss. Specifically, we train for 7 epochs using a batch size of 4
096. We use the Adam optimizer ($\beta_{1}$ = 0.9, $\beta_{2}$ = 0.999), with
a peak learning rate of $10^{-3}$, warmed up for 10 000 steps and followed by
linear decay. We use mild data augmentations (random cropping and horizontal
flipping) and weight decay of $10^{-1}$ in all parameters as means of
regularization. We trained both models on TPUv2-128 devices.