top-k 算子返回一个稀疏向量，其中非零值对应于输入的 k 个最大值。遗憾的是，由于它是一个不连续函数，因此很难纳入通过反向传播进行端到端训练的神经网络中。近期工作研究了基于正则化或扰动技术的可微松弛。然而，迄今为止，尚无方法同时具备完全可微性和稀疏性。在本文中，我们提出了新的可微且稀疏的 top-k 算子。我们将 top-k 算子视为排列多面体（permutation 的凸包）上的线性规划。随后，我们引入一个 p-范数正则化项来平滑该算子，并证明其计算可以归约为保序优化。我们的框架比现有框架显著更通用，例如能够表达按幅值选择值的 top-k 算子。在算法方面，除了 pool adjacent violator（PAV，池化相邻违例者）算法外，我们还提出了一种新的对 GPU/TPU 友好的 Dykstra 算法，用于求解保序优化问题。我们成功地将所提算子用于神经网络权重剪枝、视觉 Transformer 微调，以及作为稀疏混合专家（sparse mixture of experts）中的路由器。
英文摘要：
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.
Weak-PDE-Net: Discovering Open-Form PDEs via Differentiable Symbolic Networks and Weak Formulation
School of Mathematical Sciences, East China Normal University（东华大学数学科学学院）；Institute of Applied Physics and Computational Mathematics（应用物理与计算数学研究所）；Institute of Artificial Intelligence, Beihang University（北航人工智能研究所）；Shanghai Zhangjiang Institute of Mathematics（上海张江数学研究所）；Academy for Advanced Interdisciplinary Studies, Northeast Normal University（东北师范大学先进交叉学科研究院）
SoftMoE: Soft Differentiable Routing for Mixture-of-Experts in LLMs
AGH University of Krakow, Poland（克拉科夫AGH大学）；Faculty of Mathematics（数学系）；Computer Science, Jagiellonian University, Poland（计算机科学系，杰哥利安大学，波兰）；Centre for Credible Artificial Intelligence, Warsaw University of Technology（可信人工智能中心，华沙技术大学）
FX-DARTS: Designing Topology-unconstrained Architectures with Differentiable Architecture Search and Entropy-based Super-network Shrinking
Beijing Normal University（北京师范大学）；Southern University of Science and Technology（南方科技大学）；University of Illinois Chicago（伊利诺伊大学香槟分校）；Politecnico di Milano（米兰理工大学）；Universita’ Della Svizzera Italiana（瑞士意大利大学）
Leveraging Differentiable PDE Solvers for Semi-Neural Spatial Reconstruction From Sparse Measurements
Department of Mathematics, Technion, Haifa, Israel（数学系，技术离子学院，海法，以色列）；Department of Civil and Environmental Engineering, Technion, Haifa, Israel（土木与环境工程系，技术离子学院，海法，以色列）
Sparse Interventions in Language Models with Differentiable Masking
University of Edinburgh（爱丁堡大学）；Facebook AI Research（Facebook 人工智能研究院）；University of Amsterdam（阿姆斯特丹大学）；Innopolis University（因诺波利斯大学）；University of Osnabrück（奥斯纳布吕克大学）
DashAttention: Differentiable and Adaptive Sparse Hierarchical Attention
Tsinghua University（清华大学）；Instituto Superior Técnico, Universidade de Lisboa（里斯本大学理工学院）；Instituto de Telecomunicações（电信研究院）；Carnegie Mellon University（卡内基梅隆大学）；Sapienza University of Rome（罗马萨皮恩扎大学）；University of Edinburgh（爱丁堡大学）；TransPerfect（TransPerfect公司）；ELLIS Unit Lisbon（里斯本ELLIS单位）
CAST: Continuous and Differentiable Semi-Structured Sparsity-Aware Training for Large Language Models
Department of Computer Science and Technology, Institute for AI, BNRist Center, THBI Lab, Tsinghua-Bosch Joint ML Center, Tsinghua University（计算机科学与技术系、人工智能研究院、BNRist中心、THBI实验室、清华大学-博世联合机器学习中心、清华大学）
LACTOSE: Linear Array of Conditions, TOpologies with Separated Error-backpropagation -- The Differentiable "IF" Conditional for Differentiable Digital Signal Processing
MaxCutPool: differentiable feature-aware Maxcut for pooling in graph neural networks
Alma Mater Studiorum - University of Bologna（博洛尼亚大学）；UiT the Arctic University of Norway（挪威北极大学）；Fondazione Istituto Italiano di Tecnologia（意大利技术研究院）；NORCE Norwegian Research Centre AS（NORCE挪威研究中心）
Deep greedy unfolding: Sorting out argsorting in greedy sparse recovery algorithms
Department of Mathematics and Statistics, Concordia University（数学与统计学系，康科迪亚大学）；Department of Applied Mathematics and Theoretical Physics, University of Cambridge（应用数学与理论物理系，剑桥大学）
Instituto Superior Técnico, Universidade de Lisboa（里斯本大学高等技术学院）；Instituto de Telecomunicações（电信研究所）；Language Technology Lab, University of Amsterdam（阿姆斯特丹大学语言技术实验室）；Champalimaud Research（尚帕利莫德研究院）；Unbabel（昂巴贝尔公司）
