---
title: Stochastic Optimization of Sorting Networks
url: https://www.emergentmind.com/topics/stochastic-optimization-of-sorting-networks
type: topic
---

# Stochastic Optimization of Sorting Networks

Stochastic optimization of sorting networks addresses the fundamental challenge of making sorting operations differentiable and amenable to gradient-based learning. Classical sorting is non-differentiable, which traditionally prohibits its integration within end-to-end trainable neural architectures. "NeuralSort: Stochastic Optimization of Sorting Networks via Continuous Relaxations" introduces NeuralSort, a differentiable surrogate for sorting based on continuous relaxations of permutation matrices. This framework enables stochastic optimization over permutations by combining NeuralSort with a reparameterized estimator for the Plackett–Luce distribution using Gumbel perturbations, thereby making sorting networks directly tractable within deep learning pipelines [1903.08850].

## 1. Non-Differentiability of Sorting Operations

The standard sort function takes $s\in\mathbb{R}^n$ and produces a permutation $\pi$ that orders the elements, typically represented by a permutation matrix $P\in\{0,1\}^{n\times n}$, with exactly one "1" per row and column. The sorting operator is piecewise constant: infinitesimal changes in $s$ rarely alter the ranking unless a tie occurs, resulting in a Jacobian $\partial P/\partial s$ that is zero almost everywhere and undefined at ties. This non-differentiability means that embedding a sort operation inside a computational graph will zero out or remove all gradient information with respect to inputs, presenting a fundamental barrier to direct gradient-based optimization for objectives dependent on the output ordering [1903.08850].

## 2. NeuralSort: Continuous Relaxation of Permutations

### 2.1 Unimodal Row-Stochastic Matrix Relaxation

Permutation matrices $P$ are relaxed to unimodal row-stochastic matrices $U$ with the following properties:

1. $U_{ij}\geq 0$ for all $i,j$
2. $\sum_{j=1}^n U_{ij}=1$ for all $i$
3. Each row $i$ has a unique arg max at column $u_i$, and $\{u_1,\ldots,u_n\}$ forms a permutation of $\{1,\ldots,n\}$.

The pairwise-difference matrix $A(s)_{ij} = |s_i - s_j|$ enables the exact construction of $P$:
$$
P_{sort(s)}[i,j] = 1 \quad \text{if} \quad j = \operatorname{argmax}_j \left\{ (n+1-2i)s_j - (A(s)\mathbf{1}_n)_j \right\}
$$
otherwise, $P_{ij} = 0$ [1903.08850].

### 2.2 Softmax-Based Relaxation

NeuralSort replaces row-wise hard arg max with a softmax, yielding the continuous relaxation:
$$
\widehat{P}_{ij}(s;\tau) = \text{softmax}_j \left(\frac{(n+1-2i) s_j - \sum_{k=1}^{n} |s_j - s_k|}{\tau}\right)
$$
with temperature parameter $\tau>0$. An alternative softmax-based formulation is:
$$
P_{ij}(s;\tau) = 
\frac{\exp(-\max(s_i - s_j, 0) / \tau)}{\sum_{k=1}^n \exp(-\max(s_i - s_k, 0) / \tau)}
$$
These relaxations yield differentiable matrices for any $\tau>0$, recovering hard permutation matrices in the $\tau\to 0$ limit (in absence of ties).

## 3. Stochastic Gradient Estimation via Plackett–Luce and Gumbel Tricks

### 3.1 Plackett–Luce Permutation Distribution

The Plackett–Luce (PL) distribution models random permutations $\pi\in S_n$ parameterized by positive scores $s_1,\ldots,s_n$:
$$
p(\pi|s) = \prod_{i=1}^n \frac{e^{s_{\pi_i}}}{\sum_{j=i}^n e^{s_{\pi_j}}}
$$
This reflects a sequential draw without replacement, with probabilities proportional to exponentiated scores.

### 3.2 Gumbel Reparameterization for Sampling and Gradients

Sampling from PL$(s)$ is enabled using the Gumbel-max trick with i.i.d. Gumbel$(0,1)$ noise $g_i$:
\begin{align*}
\tilde{s}_i &= \log s_i + g_i \\
\pi &= \operatorname{sort\_indices}(\tilde s)
\end{align*}
This renders permutation sampling as a deterministic (but non-differentiable) function of $s$ and $g$. The expectation of interest is:
$$
L(s, \theta) = \mathbb{E}_{\pi\sim PL(s)}[f(P_\pi; \theta)]
$$
Approximating the discrete $P_{sort}$ with NeuralSort yields:
$$
\widehat{L}(s, \theta) = \mathbb{E}_{g}[f(\widehat{P}(\log s + g; \tau); \theta)]
$$
Gradients w.r.t. $s$ are then given by:
$$
\nabla_{s} \widehat{L} = \mathbb{E}_{g}[\nabla_s f(\widehat{P}(\log s + g; \tau); \theta)]
$$
which can be efficiently approximated using Monte Carlo sampling.

## 4. Stochastic Optimization Workflow

The NeuralSort stochastic optimization loop can be implemented as follows:

```python
# Pseudocode for stochastic optimization in NeuralSort

Given: data D, model parameters θ, PL-scores s(x;θ), temperature τ
for each minibatch x in D:
    s ← network_forward(x; θ)
    # Stochastic sampling loop
    for m in 1..M:
        g^(m) ← sample i.i.d. Gumbel(0,1) for each element in s
        ŝ^(m) ← log s + g^(m)
        P̂^(m) ← NeuralSort(ŝ^(m); τ)  # Soft permutation matrix
    L̂ ← (1/M) ∑_{m=1}^M f(P̂^(m); θ)  # Surrogate loss
    θ ← θ − η ∇_θ L̂                 # Back-propagation
end
```
In the limit $\tau\to 0$, row-wise arg max can be applied to recover hard permutations, supporting straight-through optimization [1903.08850].

## 5. Complexity Analysis and Computational Characteristics

The construction of the pairwise-difference matrix $A(s)$ requires $O(n^2)$ operations, fully parallelizable on GPUs. Each softmax operation per row costs $O(n)$, yielding $O(n^2)$ complexity per forward pass with no iterative normalization. Memory requirements are $O(n^2)$ per relaxed permutation per sample. 

Comparatively:

| Approach                      | Forward Pass Complexity | GPU Parallelism | Differentiability          |
|-------------------------------|------------------------|-----------------|---------------------------|
| Standard sorting              | $O(n \log n)$          | Limited         | No                        |
| Sinkhorn-based relaxations    | $O(n^2)$ per iteration | Good            | Yes (doubly-stochastic)   |
| NeuralSort                    | $O(n^2)$ one-shot      | High            | Yes (unimodal stochastic) |

NeuralSort’s single-pass $O(n^2)$ computation is competitive and often faster than iterative Sinkhorn methods for practical $n$ [1903.08850].

## 6. Empirical Performance and Impact

NeuralSort and its stochastic extension (via PL reparameterization) were evaluated on several tasks:

- **Semantic Sorting (large-MNIST, $n=5$):** Deterministic NeuralSort achieves approximately 84% exact permutation accuracy, outperforming Sinkhorn baselines (3–9%) and a naïve row-stochastic predictor (9%). Individual-rank accuracy (correctly placed elements) improves from 60% (Sinkhorn) to 92% (NeuralSort).
- **Quantile Regression (median estimation):** Mean squared error decreases from $344\times 10^{-4}$ (Sinkhorn) to $34\times 10^{-4}$ (NeuralSort), with $R^2$ improving from 0.25 to 0.94 for $n=5$.
- **Differentiable k-Nearest Neighbors ($n=100$, top $k$ selection):**
  - MNIST: 99.5% (NeuralSort) vs. 97.2% (standard kNN), 99.4% (CNN)
  - Fashion-MNIST: 93.5% (NeuralSort) vs. 85.8% (kNN), 93.4% (CNN)
  - CIFAR-10: 90.7% (NeuralSort) vs. 35.4% (kNN), 95.1% (CNN)

Across all tasks, stochastic NeuralSort offers comparable accuracy to its deterministic version while enabling principled uncertainty estimation for permutations. This framework supports a one-shot, differentiable surrogate for sorting, hard permutation projection for metrics, and a reparameterized estimator for optimizing over permutation distributions in deep learning pipelines [1903.08850].

Source: https://www.emergentmind.com/topics/stochastic-optimization-of-sorting-networks