LapSum - One Method to Differentiate Them All: Ranking, Sorting and Top-k Selection

We present a novel technique for constructing differentiable order-type operations, including soft ranking, soft top-k selection, and soft permutations. Our approach leverages an efficient closed-form formula for the inverse of a function LapSum, defined as a sum of Laplace distributions. This formulation ensures low computational and memory complexity in selecting the highest activations, enabling losses and gradients to be computed in O(n log n) time. Moreover, LapSum can easily be parallelized, both with respect to time and memory. Through extensive experiments, we demonstrate that our method outperforms state-of-the-art techniques for highdimensional vectors and large k values. Furthermore, we provide efficient implementations for both CPU and CUDA environments, underscoring the practicality and scalability of our method for large-scale ranking and differentiable ordering problems.