The top-k operation, i.e., finding the k largest or smallest elements from a collection of scores, is an important model component, which is widely used in information retrieval, machine learning, and data mining. However, if the top-k operation is implemented in an algorithmic way, e.g., using bubble algorithm, the resulting model cannot be trained in an end-to-end way using prevalent gradient descent algorithms. This is because these implementations typically involve swapping indices, whose gradient cannot be computed. Moreover, the corresponding mapping from the input scores to the indicator vector of whether this element belongs to the top-k set is essentially discontinuous. To address the issue, we propose a smoothed approximation, namely the SOFT (Scalable Optimal transport-based diFferenTiable) top-k operator. Specifically, our SOFT top-k operator approximates the output of the top-k operation as the solution of an Entropic Optimal Transport (EOT) problem. The gradient of the SOFT operator can then be efficiently approximated based on the optimality conditions of EOT problem. We apply the proposed operator to the k-nearest neighbors and beam search algorithms, and demonstrate improved performance.
Differentiable Ranks and Sorting using Optimal Transport
Google Research（谷歌研究院）
对数组进行排序是机器学习中的一个基本例程，用于计算基于秩的统计数据、累积分布函数（CDF）、分位数或选择最近的邻居和标签。然而，排序函数是分段常数（如果向量的条目被无限小地扰动，则向量的排序置换不会改变），因此没有梯度信息来反向传播。我们提出了一个框架来排序的元素，是算法可区分的。我们利用了这样一个事实，即排序可以被视为$\mathbb{R}$上的最优传输（OT）问题的一个特定实例，从输入值到预定义的排序值数组（例如，如果输入数组有$n$元素，则为$1,2,\dots,n$）。在此基础上，我们提出了广义CDF和分位数算子，通过改变目标预排序数组的大小和权重。因为这相当于使用OT的所谓Kantorovich公式，我们称这些量为K-排序，K-CDF和K-分位数。我们恢复微分算法，通过添加到OT问题的熵正则化，并近似使用几个Sinkhorn迭代。我们将这些运算符称为S-sort、S-CDF和S-quantiles，并将它们用于各种学习设置：我们将它们与最近提出的neuralsort进行基准测试[Grover et al. 2019]，提出了分位数回归的应用，并引入了top-k精度的可微公式，这些公式提供了最先进的性能。
Differentiable Expectation-Maximisation and Applications to Gaussian Mixture Model Optimal Transport
Université Paris Cité, CNRS, MAP5（巴黎大学、国家科学研究中心、MAP5）；CNRS（国家科学研究中心）；ENS Paris-Saclay（巴黎-萨克雷大学）；ECOLE POLYTECHNIQUE（巴黎高等理工学院）；Institut Polytechnique de Paris（巴黎理工学院）
Differentiable Adaptive Kalman Filtering via Optimal Transport
Shanghai Research Institute for Intelligent Autonomous Systems（上海智能自主系统研究所）；Tongji University（同济大学）；School of Computer Science and Technology（计算机科学与技术学院）；Shanghai Jiaotong University（上海交通大学）；Shanghai Zhangjiang Institute of Mathematics（上海张江数学研究所）；East China Normal University（华东师范大学）
AD-NEGF: An End-to-End Differentiable Quantum Transport Simulator for Sensitivity Analysis and Inverse Problems
Tianjin University（天津大学）；Noah’s Ark Lab, Huawei（华为诺亚方舟实验室）；University College London（伦敦大学学院）；Institute of Physics, Chinese Academy of Sciences（中国科学院物理研究所）
LACTOSE: Linear Array of Conditions, TOpologies with Separated Error-backpropagation -- The Differentiable "IF" Conditional for Differentiable Digital Signal Processing
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（可信人工智能中心，华沙技术大学）
No More Sliding Window: Efficient 3D Medical Image Segmentation with Differentiable Top-k Patch Sampling
Institute of Data Science（数据科学研究所）；National University of Singapore（新加坡国立大学）；Saw Swee Hock School of Public Health（Saw Swee Hock公共卫生学院）；Agency for Science, Technology and Research（科技研究局）；Centre for Clinical Magnetic Resonance Research（临床磁共振研究中心）；University of Oxford（牛津大学）
