PaperlayerView exact revision on arXiv

Source-linked AI summary

Stochastic Optimization of Sorting Networks via Continuous Relaxations

Aditya Grover, Eric Wang, Aaron Zweig, Stefano Ermon

arXiv:1903.08850v2stat.MLcs.LGcs.NE

TL;DR

Sorting is non-differentiable with respect to its inputs, limiting end-to-end gradient optimization in pipelines that sort complex objects. The paper introduces NeuralSort and a Plackett-Luce reparameterized gradient estimator, then evaluates them on three semantic-ordering tasks and reports improved performance over prior work.

  • Problem

    Sorting operators produce permutation-valued outputs but are non-differentiable with respect to their inputs, obstructing gradient-based optimization in models that learn orderings.

  • Method

    NeuralSort continuously relaxes permutation matrices to unimodal row-stochastic matrices, while Gumbel perturbations provide a reparameterized sampler for Plackett-Luce permutations.

  • Results

    Across three illustrative tasks, including a fully differentiable k-nearest neighbors method, the proposed relaxations outperform prior work in end-to-end learning of semantic orderings.

  • Takeaways & Limitations

    The framework supports gradient estimation for computation graphs involving sorting and stochastic optimization over permutations.

  • Takeaways & Limitations

    The analysis assumes unique maximizers for arg max, with tie handling requiring additional bookkeeping.

Abstract

from arXiv · show

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.

1 INTRODUCTION

NeuralSort addresses the non-differentiability of sorting in machine-learning pipelines by providing a continuous relaxation and a stochastic optimization method for permutations. The framework is evaluated on three tasks involving learned semantic orderings of high-dimensional objects.

  • Motivation: The method targets applications including top-k classification, information retrieval, and multi-object tracking, where learned representations must precede sorting.The motivating example is k-nearest neighbors, whose canonical pixel-basis ordering can be suboptimal for image classification.
  • NeuralSort: NeuralSort replaces permutation-matrix outputs with unimodal row-stochastic matrices that support differentiable optimization and discrete projection.Its temperature controls approximation quality, while row-wise arg max recovers the sorting permutation.
  • Stochastic optimization: NeuralSort also enables stochastic optimization over latent permutations by addressing gradient estimation for permutation-valued distributions.The paper focuses on the Plackett-Luce family and its reparameterized sampling procedure.
  • Experiments: The framework is demonstrated on handwritten-digit sorting, quantile regression, and an end-to-end learned-basis extension of k-nearest neighbors.These tasks require learning semantic orderings rather than relying on fixed input representations.

2 PRELIMINARIES

The preliminaries formalize permutations, sorting, Plackett-Luce distributions, and stochastic computation graphs. They establish that sorting produces permutation-valued outputs and that gradients through such operators are problematic.

  • Permutations: An n-dimensional permutation is a list of unique indices from {1, 2, . . . , n}, associated with a permutation matrix.The matrix representation encodes the selected ordering using binary entries.
  • Sorting: The sort operator maps n real-valued scores to the permutation representing their descending order, with ties resolved by appearance order.The sorted vector can be obtained by multiplying the corresponding permutation matrix by the input vector.
  • Plackett-Luce distributions: The Plackett-Luce distribution generates permutations by selecting items sequentially without replacement according to score-proportional choice probabilities.Its support contains all n! permutations, and the score vector parameterizes the distribution.
  • Stochastic computation graphs: A stochastic computation graph represents inputs, deterministic functions, and conditionally distributed stochastic nodes in a directed acyclic graph.The abstraction specifies both forward values and backward gradient computations.
  • Differentiability: Gradient-based optimization fails at the sort node because gradients with respect to its scores are not defined.The paper therefore extends computation graphs with relaxed sorting and permutation-distribution nodes.

3 NEURALSORT: THE RELAXED SORTING OPERATOR

NeuralSort relaxes sorting into the space of unimodal row-stochastic matrices, preserving differentiability while allowing efficient recovery of discrete permutations. This construction is designed for surrogate objectives and straight-through optimization.

  • Objective: The proposed surrogate replaces the permutation matrix in f(P_z; θ) with an approximation that is differentiable with respect to the scores.The original objective uses z = sort(s), where s contains real-valued scores and f is differentiable in θ and z.
  • Relaxation properties: A useful relaxation must be continuous and almost-everywhere differentiable, while also admitting an efficient projection back to the discrete output space.These properties support automatic differentiation and evaluation with losses requiring discrete outputs.
  • Straight-through optimization: Straight-through optimization can use the exact sorting operator in the forward pass and the relaxed operator for backward gradient estimation.The discrete output is retained for objectives that require exact permutations.
  • Assumption: The construction assumes a unique maximizer for arg max, with additional bookkeeping available to resolve ties.The stated results extend to non-unique inputs when tie handling is included.
  • Unimodal row-stochastic matrices: NeuralSort maps scores to unimodal row-stochastic matrices, a relaxation motivated by the inclusion relationships among permutation, doubly-stochastic, and row-stochastic matrices.Unimodal row-stochastic matrices retain row normalization and distinct row maxima without requiring every column to sum to one.

2. Row Affinity: Pn

NeuralSort relaxes sorting into a continuous unimodal row-stochastic matrix whose row-wise maxima recover the sorted permutation. The relaxation is differentiable almost everywhere, converges to the permutation matrix as temperature vanishes, and supports efficient hard projection.

  • Computation: Computing the pairwise-difference matrix requires O(n^2) operations, but the computation is highly parallelizable and suitable for GPU implementation.This contrasts with the O(n log n) complexity of the best known sorting algorithms.
  • NeuralSort: NeuralSort replaces the non-differentiable arg max with soft max to produce a continuous relaxation of the sorting permutation matrix.The relaxed matrix depends on a temperature τ > 0 and is differentiable almost everywhere with respect to the input scores.
  • NeuralSort: For every τ > 0, the relaxed matrix is unimodal and row-stochastic, with each row summing to one and having a distinct maximum.These properties define the target matrix family and permit row-wise extraction of a permutation.
  • Projection: Row-wise arg max of the relaxed matrix returns the same permutation as sorting the input scores.This provides a direct projection from the relaxed representation to the hard permutation matrix.
  • Limiting behavior: As τ approaches 0 from above, each relaxed row converges almost surely to the corresponding row of the sorting permutation matrix.The almost-sure statement assumes independently drawn scores from an absolutely continuous distribution.

4 STOCHASTIC OPTIMIZATION OVER PERMUTATIONS

The paper uses NeuralSort to optimize objectives involving random permutations, focusing on Plackett–Luce distributions and reparameterized Gumbel-based sampling. Replacing the sampler’s sorting step with NeuralSort yields differentiable surrogate gradients that can be estimated efficiently by Monte Carlo.

  • Stochastic objectives: Permutation-valued stochastic nodes represent uncertainty in latent orderings within stochastic computation graphs.The resulting objective averages a differentiable function of a permutation matrix over a parameterized distribution on permutations.
  • Gradient estimation: Exact evaluation of permutation expectations is generally intractable because it requires summing over a combinatorially large set, while Monte Carlo provides efficient estimates.Gradients with respect to distribution parameters cannot be obtained directly because the sampling distribution depends on those parameters.
  • Plackett–Luce sampling: The Plackett–Luce distribution admits reparameterized sampling by adding independent Gumbel perturbations to log-scores and sorting the perturbed values.This expresses a PL sample as a deterministic function of the scores and fixed randomness.
  • Differentiable surrogate: The reparameterized sampler remains non-differentiable because it contains a sorting operator, so NeuralSort supplies a differentiable surrogate objective.The surrogate replaces the hard sort while retaining the reparameterized computation structure.
  • Gradient estimation: The resulting reparameterized gradient estimates can be computed efficiently with Monte Carlo because the expectation is over a distribution independent of the score parameters.This avoids the direct parameter dependence that complicates score-function gradient estimation.

5 DISCUSSION AND RELATED WORK

NeuralSort differs from prior permutation relaxations centered on the Birkhoff polytope and addresses the factorial growth of permutation spaces. Its stochastic construction leverages tractable Plackett–Luce modeling and a distribution-specific reparameterization result.

  • Permutation relaxations: Prior methods relax permutation matrices to the Birkhoff polytope, the convex hull of permutation matrices also called the set of doubly-stochastic matrices.These approaches include Sinkhorn-based mappings and related distributions over doubly-stochastic matrices.
  • NeuralSort: NeuralSort instead maps permutation matrices to unimodal row-stochastic matrices, while its stochastic framework uses Plackett–Luce distributions over permutations.The PL distribution supports efficient sampling and exact, tractable density estimation.
  • Stochastic optimization: The proposed PL reparameterization also uses Gumbel noise but relies on a result stated as unique to the Plackett–Luce distribution.This distinguishes it from general categorical relaxations despite the shared use of Gumbel perturbations.
  • Scaling: The number of n-dimensional permutations is n!, so treating permutations as categorical outcomes creates a combinatorially large state space.This factorial growth limits direct application of categorical Gumbel-softmax methods to permutations.

6 EXPERIMENTS

The experiments evaluate NeuralSort on semantic sorting, median estimation, and differentiable kNN tasks involving high-dimensional inputs. NeuralSort variants outperform competing methods across the reported tasks and metrics.

  • Sorting handwritten numbers: The experiments use large-MNIST sequences to learn permutations sorting images by their unobserved numerical labels.The task requires learning image semantics before sorting, with weak supervision from ground-truth permutations.
  • Sorting handwritten numbers: NeuralSort approaches significantly outperform baseline approaches in correctly recovering complete permutations for every tested sequence length.Deterministic and stochastic variants have comparable performance, while vanilla RS ranks individual elements better than it recovers whole permutation matrices.
  • Quantile regression: Quantile regression extends sorting to predicting the 50-th percentile while also mapping the selected median image to its scalar value.The task uses sequences of large-MNIST images and targets the median label, requiring both semantic sorting and a secondary regression function.
  • Quantile regression: NeuralSort approaches outperform competing methods on both MSE and R2 for quantile regression.Stochastic NeuralSort is consistently best on MSE, whereas deterministic NeuralSort is slightly better on R2.
  • End-to-end, differentiable k-nearest neighbors: The differentiable kNN classifier learns representations before selecting neighbors and outperforms all baseline kNN variants on the reported benchmark datasets.Its performance is much closer to that of a standard CNN than to the baseline kNN methods.

7 CONCLUSION

The paper presents NeuralSort as a continuous sorting relaxation and derives a reparameterized gradient estimator for stochastic optimization over permutations. Across three illustrative high-dimensional ordering tasks, the proposed relaxations outperform prior work.

  • Conclusion: NeuralSort relaxes sorting to unimodal row-stochastic matrices, enabling gradient estimation in computation graphs involving sort operators.The paper also derives a reparameterized gradient estimator for the Plackett-Luce distribution.
  • Conclusion: On three illustrative tasks, NeuralSort relaxations outperform prior work for end-to-end learning of semantic orderings of high-dimensional objects.The applications include a fully differentiable k-nearest neighbors algorithm.
  • Conclusion: Both deterministic and stochastic NeuralSort are described as easy to implement, with reference implementations provided in TensorFlow and PyTorch.The paper identifies exploring alternate sorting relaxations and extensions such as beam search as future work.

APPENDICES

The appendices provide TensorFlow and PyTorch implementations of deterministic and stochastic NeuralSort. The code constructs relaxed sorting matrices and samples perturbed scores for stochastic optimization.

  • Deterministic NeuralSort: The deterministic implementation accepts batched score vectors and a scalar temperature to compute a relaxed sorting matrix.Its computation forms pairwise absolute score differences, a rank-dependent scaling term, and a softmax relaxation.
  • Stochastic NeuralSort: The stochastic implementation samples Gumbel perturbations and combines them with the score parameters before applying the sorting relaxation.The sampler supports a specified number of samples and temperature, with implementations shown in TensorFlow and PyTorch.
  • Deterministic NeuralSort: The deterministic implementation applies softmax to rank-specific scores after subtracting the pairwise-distance term, scaled by temperature.The resulting matrix is formed as P_hat = softmax(P_max / tau).
  • Stochastic NeuralSort: The PyTorch sampler generates Gumbel noise from uniform random variables using a numerically stabilized logarithmic transformation.The implementation uses eps = 1e-10 in both logarithmic terms.

B PROOFS OF THEORETICAL RESULTS

The proofs establish that NeuralSort produces row-stochastic matrices whose row argmaxes recover sorted-score indices, including a recursive tie-handling rule. They also connect the softmax relaxation and Gumbel perturbations to the intended sorting and permutation-distribution properties.

  • NeuralSort properties: Softmax makes every NeuralSort row positive and row-stochastic, while monotonicity preserves the row argmax of its score expression.The proof uses softmax normalization for positivity and unit row sums, then uses monotonicity to identify argmax indices.
  • NeuralSort properties: For distinct scores, row i has its maximum at the index [i] of the i-th largest score, so the row argmax vector is a permutation.The argument relies on the established score expression and the definition of sort(s) as the ordered index vector.
  • NeuralSort properties: As temperature τ approaches zero, the softmax relaxation converges to an argmax because the score elements are distinct almost surely.The proof invokes the variational characterization of softmax and the distributional assumption that prevents ties almost surely.
  • Tie handling: With tied scores, a recursive argmax rule selects the smallest previously unused maximizing index, otherwise the smallest maximizing index.This tie-breaking protocol is efficiently computable with additional bookkeeping.
  • Tie handling: The tie proofs show that these assignments remain distinct across rows, yielding a valid permutation even when multiple entries share a maximum.The argument groups rows and columns associated with equal scores and applies a pigeon-hole argument to guarantee an unused maximizer.

D EXPERIMENTAL DETAILS AND ANALYSIS

The experiments compare NeuralSort and Sinkhorn-based relaxations under matched architectures across sorting, quantile regression, and differentiable k-nearest-neighbors settings. Temperature affects gradient variance, while quantile-regression analysis shows lower absolute MSE but slightly declining R2 as the sequence length grows.

  • Experimental setup: The sorting experiments use matched convolutional architectures and row-wise cross-entropy against the true permutation matrix for all methods.NeuralSort maps image representations to scalar scores, whereas Sinkhorn methods map them to n-dimensional vectors before forming relaxed matrices.
  • Temperature analysis: Higher temperatures τ produce lower variance in gradient estimates during training for NeuralSort-based stochastic optimization.The study reports running log-variance as a function of τ and separately measures approximation error against projected hard permutation matrices.
  • Quantile regression: Quantile regression predicts the median of labels in image sequences while learning both a high-dimensional sorting representation and a scalar mapping from the inferred median.The experiment uses four-digit sequences because of resource constraints and tunes τ using validation loss.
  • Quantile regression: As sequence length n increases, absolute median-prediction MSE decreases because true medians concentrate, while R2 shows a slight dip.The paper interprets this pattern as an easier absolute-scale prediction problem but a harder semantic-sorting problem over more elements.
  • Differentiable k-nearest neighbors: The differentiable k-nearest-neighbors experiments evaluate NeuralSort on MNIST, Fashion-MNIST, and CIFAR-10 using learned encodings for distance computation.The reported comparison includes deterministic and stochastic NeuralSort accuracies broken down by k.
Loading 1903.08850v2…