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算子返回一个稀疏向量，其中非零值对应输入中的k个最大值。不幸的是，由于它是一个不连续函数，难以融入通过反向传播进行端到端训练的神经网络中。近期的工作考虑了基于正则化或扰动技术的可微松弛方法。然而，迄今为止，尚无方法同时具备完全可微性和稀疏性。在本文中，我们提出了新的可微且稀疏的Top-k算子。我们将Top-k算子视为排列多面体（即排列的凸包）上的线性规划问题。接着，我们引入p-范数正则化项来平滑该算子，并证明其计算可简化为等调优化问题。我们的框架相比现有方法具有显著更广的泛化性，例如能够表达按幅度选择值的Top-k算子。在算法层面，除了池化相邻违反者（PAV）算法外，我们还提出一种新的适用于GPU/TPU的Dykstra算法来解决等调优化问题。我们成功地将所提出的算子用于神经网络剪枝、视觉变压器微调，以及作为稀疏混合专家模型中的路由器。