PaperlayerView exact revision on arXiv

Source-linked AI summary

SoftSort: A Continuous Relaxation for the argsort Operator

Sebastian Prillo, Julian Martin Eisenschlos

arXiv:2006.16038v1cs.LGstat.ML

TL;DR

Argsort is discrete and has zero gradients almost everywhere, while existing continuous relaxations are described as computationally complex. SoftSort relaxes Pargsort using row-wise softmaxes of negative distances. It achieves state-of-the-art performance on multiple reordering benchmarks and is faster than NeuralSort while requiring the same number of gradient steps to converge.

  • Problem

    Argsort has a discrete image and zero gradients almost everywhere, while prior continuous relaxations are described as computationally complex.

    Source: §1 · p. 1

  • Method

    SoftSort relaxes Pargsort by applying a row-wise softmax to negative distances from each score to the corresponding ranked score.

    Source: §1 · p. 2, §3 · p. 4

  • Results

    SoftSort achieves state-of-the-art performance on multiple reordering benchmarks and is faster than NeuralSort while requiring the same gradient steps to converge.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors conclude that SoftSort advances differentiable sorting by simplifying previous approaches while retaining similar properties to NeuralSort.

    Source: §6 · p. 9

    The n × n matrix construction and application account for most computation; sparse SoftSort variants are proposed as future work to improve asymptotic speed.

    Source: §5.5 · p. 9

Abstract

from arXiv · show

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.

1. Introduction

Argsort is discrete and has zero gradients almost everywhere, limiting its use in gradient-based learning. SoftSort relaxes the permutation-matrix form of argsort with a simple, differentiable operator.

1. Introduction

  • Argsort has a discrete image and zero gradients almost everywhere, making it unsuitable for gradient-based optimization.The paper defines argsort as returning the permutation that orders a real-valued vector in decreasing order.Source: §1 · p. 1
  • Rather than relax argsort directly, the paper relaxes its permutation-matrix representation, Pargsort, analogously to how softmax relaxes argmax.The argsort permutation can be recovered from Pargsort by matrix multiplication.Source: §1 · p. 1
  • Applied to vectors, SoftSort produces differentiable convex combinations that approximate the hard permutation ordering by scores.The operation is performed by multiplying the vectors on the left by the SoftSort matrix.Source: Figure 1 · p. 2, §1 · p. 2
  • Each SoftSort row is a softmax over negative distances to the corresponding ranked score, with temperature controlling the approximation.The paper predominantly uses the L1 distance, while its theoretical results hold for any qualifying distance function.Source: §1 · p. 2
  • SoftSort achieves state-of-the-art performance on multiple reordering benchmarks and is faster than NeuralSort while requiring the same gradient steps to converge.The paper describes NeuralSort as the fastest competing approach.Source: §1 · p. 2
  • SoftSort shares NeuralSort’s row-stochasticity, convergence to Pargsort as temperature tends to zero, and row-wise argmax projection to a permutation matrix.The paper says SoftSort is significantly easier to reason about mathematically, simplifying proofs of these properties.Source: §1 · p. 2

2. Related Work

Prior relaxed sorting work includes approaches to ranks, bipartite matching, optimal transport, and NeuralSort; the paper presents these as context for its simpler Pargsort relaxation.

2.1. Relaxed Rank Operators

  • SoftRankτ and bπτ directly relax the rank operator rather than Prank, and the paper says they provide no easy way to obtain a relaxed argsort or Pargsort operator from them.SoftRankτ and its gradients compute exactly in O(n^3); bπτ computes in O(n^2), faster than Taylor et al.’s O(n^3) approach.Source: §2.1 · pp. 2–3

2.2. Sorting via Bipartite Maximum Matchings

  • Mena et al. relax a bipartite matching operator with Sinkhorn iterations, mapping matrices toward doubly stochastic matrices.The reported complexity is O(Ln^2), where L is the number of Sinkhorn iterations.Source: §2.2 · p. 3

2.3. Sorting via Optimal Transport

  • Cuturi et al. derive a Pargsort relaxation from entropy-regularized optimal transport, computed with Sinkhorn iterations at O(Ln^2) complexity.Their relaxation maps to the Birkhoff polytope and converges to Pargsort as regularization tends to zero.Source: §2.3 · p. 3

2.4. Sorting via a sum-of-top-k elements identity

  • NeuralSort replaces an argmax in a top-k identity with row-wise softmax, yielding a Pargsort relaxation with O(n^2) computation.Its temperature controls approximation, and NeuralSort converges to Pargsort as temperature tends to zero.Source: §2.4 · p. 3

3. SoftSort: A simple relaxed sorting operator

SoftSort is a family of continuous relaxations for Pargsort, defined using a row-wise softmax of negative pairwise distances and a temperature parameter.

3. SoftSort: A simple relaxed sorting operator

  • SoftSort defines a continuous relaxation of Pargsort using a row-wise softmax, a temperature, and an almost-everywhere differentiable semi-metric.A semi-metric has the properties of a metric except the triangle inequality; positive powers of absolute value are examples in R.Source: §3 · p. 4

2. Row Affinity: Pn

SoftSort is a unimodal row-stochastic relaxation whose row-wise argmax recovers argsort, and it converges to the discrete permutation as temperature approaches zero when scores are distinct. Its rows represent distributions over which input occupies each rank.

2. Row Affinity: Pn

  • URS matrices require distinct row-wise argmaxes but not column-stochasticity; their rows can be projected onto the probability simplex.The authors consider this structure more natural for relaxing argsort than bistochastic matrices, which are used for bipartite matching or optimal transport.Source: §2 · p. 4
  • For any τ > 0, SoftSort is unimodal and row-stochastic, and its row-wise argmax yields argsort(s).For inputs with distinct elements, it converges to the discrete argsort permutation as τ → 0; this also holds almost surely for samples from absolutely continuous distributions.Source: §2 · pp. 4–5
  • SoftSort’s unimodality proof is straightforward, unlike NeuralSort’s proof, which requires a technical lemma and corollary.The cited NeuralSort results are Grover et al. (2019, Lemma 2, Corollary 3).Source: §2 · p. 5
  • Row r represents a distribution over which input value occupies rank r; with d = |·|, the first row is softmax(s).More generally, each row is the softmax of negative distances to the score at that rank.Source: §2 · p. 5
  • The authors found absolute distance marginally better than squared distance in experiments and use d = |·| unless stated otherwise.They leave learning d and exploring a larger family of functions for future work.Source: §2 · p. 5

4. Comparing SoftSort to NeuralSort

SoftSort’s score-pair logits avoid NeuralSort’s dependence on intermediate scores, and the two operators’ forms for sorted inputs extend to arbitrary inputs by relabeling. Its TensorFlow implementation is simpler and uses the sort operator, while the authors report it is faster than NeuralSort.

4.1. Mathematical Simplicity

  • For sorted scores, SoftSort’s logits are −|si−sj|, while NeuralSort’s depend on all intermediate scores between si and sj.The authors attribute NeuralSort’s more complex form to its coupling with the complex identity used in its derivation.Source: §4.1 · p. 5
  • The operator forms for sorted scores extend to arbitrary inputs by sorting the scores, applying the operator, and permuting columns with the inverse sorting permutation.The figures for sorted scores therefore apply to other inputs up to score relabeling and column permutation.Source: §4.1 · pp. 5–6

4.2. Code Simplicity

  • SoftSort’s TensorFlow implementation is simpler than NeuralSort’s, follows Eq. 4, and relies on the sort operator.The authors note that sort is continuous and almost everywhere differentiable; they also report SoftSort is faster than NeuralSort and that CPU/GPU sorting improvements benefit its speed.Source: §4.2 · p. 6

5. Experiments

Across sorting, quantile-regression, and differentiable-kNN experiments, SoftSort performs comparably to NeuralSort, while the operator-speed experiment finds SoftSort faster and reports the same number of gradient steps to convergence.

5.1. Models

  • The experiments compare deterministic operators with stochastic variants formed by relaxing sampled Plackett–Luce permutations.The stochastic variants use Gumbel reparameterization.Source: §5.1 · p. 6
  • Sorting-task results are comparable with the state of the art from n=7 onward, and quantile-regression results are also reported as comparable.Table 1 reports the proportion of correct permutations averaged over 10 runs; Table 2 reports median-prediction mean squared error and Spearman's R2.Source: Table 1 · p. 7, Table 2 · p. 7

5.2. Sorting Handwritten Numbers

  • On the first metric, SoftSort and NeuralSort perform identically for every n and both outperform the method of at n=9 and n=15.The authors attribute NeuralSort’s improved results over those reported by to better hyperparameters; both operators’ learning curves converge at the same speed.Source: §5.2 · p. 7

5.3. Quantile Regression

  • SoftSort and NeuralSort perform similarly on median prediction, with each sometimes outperforming the other.The authors report that both methods' learning curves converge at the same speed.Source: §5.3 · p. 7

5.5. Speed Comparison

  • At N=4000, SoftSort is about 6 times faster than NeuralSort (Grover et al., 2019) on both CPU and GPU; after performance improvements, NeuralSort remains 80% slower on CPU and 40% slower on GPU.Each epoch batches twenty vectors of size n.Source: Figure 6 · p. 8, §5.5 · p. 9

5.4. Differentiable kNN

  • On every evaluated dataset, the differentiable kNN model achieves accuracy comparable to NeuralSort, and results are comparable with or slightly above the state of the art.Larger k values consistently yielded better results in these experiments.Source: §5.4 · pp. 8–9

5.5. Speed Comparison

  • Reversely sorted inputs caused no significant slowdown, and shuffling inputs can avoid this concern in applications where it matters.Source: §5.5 · p. 9
  • Sparse SoftSort is left for future work because building and applying its n × n matrix accounts for most computation, while sorting itself is sub-quadratic.Source: §5.5 · p. 9

6. Conclusion

SoftSort is a simple relaxation whose properties resemble NeuralSort’s, while its mathematical simplicity supports easier proofs and implementation. The appendix reports similar learning curves for several tasks, a slower curve for one distance choice, and speed and memory results for PyTorch implementations.

6. Conclusion

  • The r-th SoftSort row is the softmax of negative distances to the r-th largest element.The conclusion presents this as the operator’s defining construction.Source: §6 · p. 9
  • SoftSort has NeuralSort-like properties but is simpler, faster, more modular, and just as easy to optimize empirically.The conclusion also says its proofs are effortless and it is trivial to implement.Source: §6 · p. 9

A.2.4. LEARNING CURVES

  • In quantile regression with N = 15, deterministic SoftSort and NeuralSort had almost identical learning curves and could essentially be exchanged.The curves average results over 10 runs.Source: §A.2.4 · p. 12, Figure 7 · p. 12

A.4.4. LEARNING CURVES

  • For N = 4000, SoftSort with d = |·|^2 and NeuralSort had almost identical learning curves, while SoftSort with d = |·| converged more slowly.The plotted metric is Spearman correlation against epoch.Source: §A.4.4 · p. 13, Figure 8 · p. 13

A.4.5. NEURALSORT PERFORMANCE IMPROVEMENT

  • The NeuralSort implementations had O(n^3) complexity, while reassociating the computation reduced it to O(n^2) and sped up the provided implementations.The passage attributes the cubic complexity to the original matrix multiplication arrangement.Source: §A.4.5 · pp. 13–14

A.4.6. PYTORCH RESULTS

  • In PyTorch, NeuralSort ran out of memory on CPU at n = 3600 and GPU at n = 3900, while SoftSort ran out on CPU at n = 3700.The passage says the PyTorch speed results were otherwise similar to TensorFlow results.Source: §A.4.6 · p. 14

A.4.7. HARDWARE SPECIFICATION

The speed-comparison experiments used a Google Cloud CPU instance and a V100 GPU. The authors report closely reproducing the GPU and CPU results on separate Amazon EC2 instances.

A.4.7. HARDWARE SPECIFICATION

  • The Google Cloud CPU tests used an n1-highmem-2 instance with 2 vCPUs and 13 GB memory.The GPU tests used a V100.Source: §A.4.7 · p. 14
  • The V100 GPU results were closely reproduced on an Amazon EC2 p2.xlarge instance with a Tesla K80 GPU.The alternative instance had 4 vCPUs and 61 GB memory.Source: §A.4.7 · p. 14
  • The CPU results were closely reproduced on an Amazon EC2 c5.2xlarge instance with 8 vCPUs and 16 GB memory.Source: §A.4.7 · p. 14

B. Proof of Proposition 1

The proof uses permutation-matrix identities and row-wise softmax behavior to establish Proposition 1 for SoftSort and NeuralSort. For SoftSort, the argument also uses that sorting a sorted vector leaves it unchanged.

B. Proof of Proposition 1

  • Lemma 1 shows that applying a pointwise function commutes with right multiplication by a permutation matrix.The proof explains that right multiplication permutes columns, preserving the pointwise operation’s correspondence.Source: B. Proof of Proposition 1 · p. 14
  • Lemma 2 establishes the row-wise softmax identity used to move a permutation matrix into the softmax operation.The proof section invokes this identity in both operator proofs.Source: B. Proof of Proposition 1 · p. 14
  • For SoftSort, the proof uses sort(sort(s)) = sort(s) and the lemmas to derive Proposition 1; the NeuralSort proof applies the same permutation and softmax identities.The NeuralSort derivation also simplifies using sort(s)^T P_argsort(s) = s^T and 1^T P_argsort(s) = 1^T.Source: B. Proof of Proposition 1 · pp. 15–16

C. Proof of Proposition 2

For the k = 1 differentiable kNN proposition, the proof uses only SoftSort’s first row. With normalized embeddings, it rewrites the distance-based softmax as one proportional to exponentiated inner products before multiplying by the label.

C. Proof of Proposition 2

  • When k = 1, the differentiable kNN result uses only the first row of the SoftSort matrix.The proposition specifies the SoftSort distance choice and a norm-1 embedding function.Source: C. Proof of Proposition 2 · p. 16
  • Because s[1] is at least every s_i, the proof removes the negative absolute-value terms and cancels s[1] using softmax invariance to additive constants.Source: C. Proof of Proposition 2 · p. 16
  • Normalized embeddings let the proof express s_i through inner products; at temperature 2, the softmax values are proportional to e^(Φ(x_i)·Φ(ẋ)), and multiplying by the label yields the stated identity.Source: C. Proof of Proposition 2 · p. 16

D. Magnitude of Matrix Entries

On the less stringent metric of individual permutation elements correctly identified, SoftSort and NeuralSort perform identically across all n and outperform Cuturi et al. for n = 9, 15. The section also characterizes matrix-entry shapes, their conditions, and the relative simplicity of the SoftSort implementation.

D. Magnitude of Matrix Entries

  • SoftSort|·| rows are proportional to Laplace densities evaluated at s_j, while NeuralSort rows are Gaussian-shaped only when the s_i are equally spaced.Under equal spacing, the NeuralSort Gaussian variance is aτ; SoftSort|·|2 has Gaussian variance τ and does not generally recover NeuralSort.Source: D. Magnitude of Matrix Entries · pp. 16–17
  • SoftSort|·|2 rows are proportional to Gaussian densities evaluated at s_j, and the operators perform comparably on benchmarks except for the speed comparison task.On that task, NeuralSort and SoftSort|·|2 seem to outperform SoftSort|·|.Source: D. Magnitude of Matrix Entries · pp. 16–17, Figure 10 · p. 17
  • The comparison between NeuralSort and SoftSort|·|2 is limited to equally spaced inputs; changing the spacing requires changing SoftSort|·|2's temperature to match NeuralSort.Even in the equally spaced case, the operators' gradients need not agree; their gradient comparisons remain under-explored, with only empirical learning-curve insights reported.Source: D. Magnitude of Matrix Entries · p. 17

F. PyTorch Implementation

E. Sorting Task - Proportion of Individual Permutation Elements Correctly Identified

  • SoftSort and NeuralSort perform identically for all n and outperform Cuturi et al. (2019) for n = 9, 15 on the proportion of individual permutation elements correctly identified.Results are means and standard deviations over 10 runs; this metric is less stringent than the first metric.Source: E. Sorting Task - Proportion of… · p. 17, Table 5 · p. 18

Related papers

Loading 2006.16038v1…