Reproduce “Stochastic Optimization of Sorting Networks via Continuous Relaxations” on OpenResearch

About this paper

Aditya Grover, Eric Wang, Aaron Zweig, Stefano Ermon

NeuralSort makes sorting trainable inside neural computation graphs by replacing hard permutation matrices with temperature-controlled, unimodal row-stochastic relaxations. Combined with Gumbel perturbations, it also provides a low-variance, reparameterized approximation for optimizing over Plackett–Luce permutations; experiments show strong gains over Sinkhorn-style baselines in semantic sorting, quantile regression, and differentiable k-nearest-neighbor learning.

Key insights

• Matching the relaxation geometry to sorting is important: unimodal row-stochastic matrices preserve rank-to-item structure without imposing the column-sum constraints of the Birkhoff polytope.
• Temperature controls the tradeoff between smooth, lower-variance optimization and fidelity to a discrete permutation; as temperature approaches zero, the relaxation converges to hard sorting under distinct scores.
• The method generalizes beyond an isolated sorting benchmark by allowing ordering operations to shape learned representations, although its quadratic cost and relaxation bias limit scalability and exactness.

Summary by alphaXiv.