[论文解读] SoftSort: A Continuous Relaxation for the argsort Operator

Sebastian Prillo, Julian Martin Eisenschlos|arXiv (Cornell University)|Jun 29, 2020

Domain Adaptation and Few-Shot Learning参考文献 13被引用 6

一句话总结

本文提出 SoftSort，一种用于 argsort 操作的简单且高效的连续松弛方法，可在深度学习中实现可微排序。通过计算每个排序元素到其负 L1 距离的 softmax，SoftSort 实现了最先进性能，比先前的 NeuralSort 方法更快，并在保持强大优化特性的同时简化了数学分析。

ABSTRACT

While sorting is an important procedure in computer science, the argsort operator - which takes as input a vector and returns its sorting permutation - has a discrete image and thus zero gradients almost everywhere. This prohibits end-to-end, gradient-based learning of models that rely on the argsort operator. A natural way to overcome this problem is to replace the argsort operator with a continuous relaxation. Recent work has shown a number of ways to do this, but the relaxations proposed so far are computationally complex. In this work we propose a simple continuous relaxation for the argsort operator which has the following qualities: it can be implemented in three lines of code, achieves state-of-the-art performance, is easy to reason about mathematically - substantially simplifying proofs - and is faster than competing approaches. We open source the code to reproduce all of the experiments and results.

研究动机与目标

• 解决 argsort 操作中零梯度的问题，该问题阻碍了可微模型的端到端训练。
• 开发一种计算高效且易于实现的 argsort 连续松弛方法。
• 与 NeuralSort 等先前方法相比，显著简化可微排序的数学推理。
• 在减少计算开销的同时，实现与现有方法相当或更优的性能。

提出的方法

• SoftSort 通过计算每个元素与输入向量中第 r 大值之间负 L1 距离的行向 softmax，得到软排序排列。
• 该方法使用温度超参数 τ 控制松弛的锐度，当 τ → 0 时，其收敛至真实 argsort。
• 核心运算定义为 SoftSort^d_τ(s) = softmax( -d(sort(s)1^T, 1s^T) / τ )，其中 d 为半度量（例如 |x−y|）。
• 该方法利用 PyTorch 和 TensorFlow 等标准自动微分框架，可无缝集成到深度学习流水线中。
• SoftSort 是行随机的，可通过行向 argmax 投影为有效排列矩阵，保持可解释性。
• 该方法具有模块化和可扩展性，支持使用不同的距离函数 d，同时保持理论保证。

实验结果

研究问题

• RQ1 能否设计一种既简单又高效、同时保持高性能的 argsort 连续松弛方法？
• RQ2 SoftSort 在速度和优化效率方面是否优于现有的可微排序方法（如 NeuralSort）？
• RQ3 通过新型松弛方法，可微排序的数学分析能否得到显著简化？
• RQ4 SoftSort 在涉及重排和排序任务的多样化基准测试中表现如何？
• RQ5 不同距离函数 d 对 SoftSort 的收敛性和性能有何影响？

主要发现

• SoftSort 在多个涉及重排、匹配的任务基准测试中达到最先进性能，与 NeuralSort 表现相当或略优。
• 在 CPU 和 GPU 上，对于输入大小 n ≤ 4000 的情况，SoftSort 的速度约为 NeuralSort 实现的 6 倍。
• 即使在优化后，SoftSort 在 CPU 上比优化版 NeuralSort 快 80%，在 GPU 上快 40%。
• SoftSort（d=|·|）与 NeuralSort 的学习曲线几乎完全一致，表明其优化动态相近。
• 在合成排序任务中，SoftSort（d=|·|）的收敛速度慢于 d=|·|²，表明其对距离函数选择较为敏感。
• 在最坏输入顺序下未观察到显著性能下降，且通过输入打乱可进一步缓解潜在性能问题。