[论文解读] Stochastic Optimization of Sorting Networks via Continuous Relaxations

Aditya Grover, Eric Wang+2 more|arXiv (Cornell University)|Mar 21, 2019

Stochastic Gradient Optimization Techniques参考文献 35被引用 36

一句话总结

本文提出 NeuralSort，一种可微连续的排序松弛，用于实现端到端梯度优化，并给出对置换的 Plackett-Luce 分布的再参数化梯度估计器。

ABSTRACT

Sorting input objects is an important step in many machine learning pipelines. However, the sorting operator is non-differentiable with respect to its inputs, which prohibits end-to-end gradient-based optimization. In this work, we propose NeuralSort, a general-purpose continuous relaxation of the output of the sorting operator from permutation matrices to the set of unimodal row-stochastic matrices, where every row sums to one and has a distinct arg max. This relaxation permits straight-through optimization of any computational graph involve a sorting operation. Further, we use this relaxation to enable gradient-based stochastic optimization over the combinatorially large space of permutations by deriving a reparameterized gradient estimator for the Plackett-Luce family of distributions over permutations. We demonstrate the usefulness of our framework on three tasks that require learning semantic orderings of high-dimensional objects, including a fully differentiable, parameterized extension of the k-nearest neighbors algorithm.

研究动机与目标

• 在排序不可微的神经网络流水线中，动机是使排序成为一个可微的操作。
• 提出 NeuralSort，使排序松弛为单峰行随机矩阵，并带有一个温度控制的近似。
• 通过重参数化的 PL 采样方法，启用对置换的基于梯度的随机优化。
• 演示其在高维语义排序任务中的适用性（例如：按推断标签对图像排序、分位数回归、可微分的 kNN）。

提出的方法

• 将 NeuralSort 定义为从排列矩阵到单峰行随机矩阵的松弛。
• 使用基于 softmax 的 argmax 松弛来获得一个带有温度参数 tau 的可微排序代理。
• 在温和假设下证明该松弛是单峰的，并且当 tau -> 0+ 时收敛到排列矩阵。
• 通过 Gumbel 扰动导出一个对 Plackett-Luce 分布可重参数化的采样器，并将 NeuralSort 作为排序的可微代理在梯度估计中使用。
• 为涉及置换分布的目标函数开发随机计算图，并提供 REINFORCE 与重参数化梯度估计器。
• 展示使用 NeuralSort 的端到端可微分 kNN 及其他语义排序任务。

实验结果

研究问题

• RQ1 可微的排序算子松弛是否能够在依赖排序的流水线中实现端到端的基于梯度的优化？
• RQ2 在通过 Plackett-Luce 分布对置换进行优化时，如何高效地估计梯度？
• RQ3 NeuralSort及其随机变体是否在需要对高维输入进行语义排序的任务中提升学习效果？
• RQ4 是否可以通过用 NeuralSort 替代排序来实现端到端训练一个可微分的 kNN？
• RQ5 随着温度参数变化，NeuralSort 的收敛性质是什么？

主要发现

• NeuralSort 产生一个可微排序代理，它是单峰行随机的，并且当温度 tau 趋向于零时收敛到真正的排列。
• 该松弛算子使直通梯度估计成为可能，并支持高效地投影回排列矩阵。
• 使用 Gumbel 扰动推导出 Plackett-Luce 分布的重参数化梯度估计器，梯度通过 NeuralSort 代理传播。
• 经验结果表明，基于 NeuralSort 的方法在涉及对高维对象排序、分位数回归和可微分 kNN 的任务上超过基线（如 Sinkhorn、Gumbel-Sinkhorn、vanilla RS）。
• 随机 NeuralSort 和确定性 NeuralSort 在多个序列规模上的排序准确性和回归指标方面提供了强劲的性能。
• 该框架通过 NeuralSort 对候选邻居进行排序，从而实现对 kNN 的端到端可微学习。