提出可微排序与选 top-k 的新方法,速度快且适合大规模数据
LapSum -- One Method to Differentiate Them All: Ranking, Sorting and Top-k Selection
- 基于拉普拉斯分布之和的反函数,实现可微排序与选 top-k
- 在高维向量和大 k 值下优于现有方法,计算复杂度 O(n log n)
- 支持 CPU 与 CUDA,适用于大规模可微排序任务
我们提出一种新型可微顺序操作构建技术,涵盖软排名、软 top-k 选择和软排列。该方法利用拉普拉斯分布之和(LapSum)函数的闭式反函数,确保在选择最高激活值时具有低计算与内存开销,可在 O(n log n) 时间内完成损失与梯度计算。通过大量实验,证明该方法在高维向量与大 k 值场景下优于当前最优技术。此外,我们提供了针对 CPU 与 CUDA 的高效实现,验证了该方法在大规模排名与可微排序问题中的实用性与可扩展性。
原文摘要 · Abstract (English)
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 the function LapSum, defined as the 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. Through extensive experiments, we demonstrate that our method outperforms state-of-the-art techniques for high-dimensional 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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。