Sinkhorn-Sort: Differentiable Sorting

• Sinkhorn-Sort is a differentiable relaxation of sorting that uses entropic-regularized optimal transport and Sinkhorn iteration to approximate permutation matrices.
• It enables end-to-end backpropagation by providing smooth gradients, balancing between exact permutations and mean operations through temperature tuning.
• Applications include numeric sorting, quantile regression, policy learning, and Transformer attention mechanisms, offering improved accuracy and reduced computational complexity.

Sinkhorn-Sort is a class of algorithmically differentiable operators providing smooth relaxations of sorting and ranking, constructed via entropic-regularized optimal transport and implemented efficiently using Sinkhorn iteration. These methods underpin a spectrum of modern deep learning solutions requiring end-to-end learnable sorting, ranking, and permutation modules, enabling differentiability for learning tasks that traditionally rely on discrete, piecewise-constant operations.

1. Mathematical Foundations: Sorting as Optimal Transport

The central insight of Sinkhorn-Sort is that sorting can be cast as a one-dimensional optimal transport (OT) or Kantorovich assignment problem. Given an input vector $x = (x_1, ..., x_n) \in \mathbb{R}^n$, the goal is to map $x$ to its sorted version. This is formulated as a transport problem between the input values $x$ and a fixed target vector $t = (1, ..., n)$, both endowed with uniform discrete measures $a = b = (1/n, ..., 1/n)$.

The assignment cost is captured by the matrix $C_{ij} = |x_i - t_j|$, more generally $h(t_j - x_i)$ for convex $h$. The unregularized OT problem seeks the doubly-stochastic matrix $P^*$ minimizing $\langle P, C \rangle$ under row and column sum constraints. For strictly convex $x$0, $x$1 is $x$2 times a permutation matrix encoding the sorting permutation. The sorted vector and rank vector are recovered as:

• $x$3
• $x$4

This perspective connects discrete sorting directly to continuous optimization structures (Cuturi et al., 2019).

2. Entropic Regularization and the Sinkhorn Operator

To enable differentiability, a Shannon entropy penalty $x$5 is introduced, controlled by temperature $x$6. The regularized OT problem becomes: $x$7 The regularized problem admits a unique solution in the Birkhoff polytope, which is everywhere differentiable in $x$8 when $x$9 (Cuturi et al., 2019, Mena et al., 2018). As $x$0, $x$1 approaches the permutation matrix solution, but gradients vanish almost everywhere.

The Sinkhorn operator parameterizes $x$2 as $x$3 with $x$4, and $x$5 chosen by alternate matrix scaling (Sinkhorn-Knopp iteration):

• $x$6
• $x$7

This procedure converges exponentially fast to the regularized solution (Cuturi et al., 2019, Mena et al., 2018).

3. Sinkhorn-Sort and Differentiable Relaxations

The $x$8-smooth sort operator or S-sort is defined as $x$9, providing a differentiable approximation to sorting. Each coordinate

$t = (1, ..., n)$0

smoothly interpolates between mean (for large $t = (1, ..., n)$1) and exact sorting (as $t = (1, ..., n)$2).

Key properties:

• Differentiability: $t = (1, ..., n)$3 is everywhere differentiable for $t = (1, ..., n)$4.
• Gradient trade-off: Small $t = (1, ..., n)$5 approximates permutations but causes vanishing gradients; large $t = (1, ..., n)$6 increases smoothness but dilutes permutation structure.
• Computational complexity: Each Sinkhorn iteration requires $t = (1, ..., n)$7 time; typically $t = (1, ..., n)$8–$t = (1, ..., n)$9 iterations suffice for convergence.
• Automatic Differentiation: S-sort can be differentiated by unrolling iterations or via the implicit function theorem, ensuring nonzero gradients throughout (Cuturi et al., 2019, Mena et al., 2018).

Generalizations yield soft CDF (“K-rank”), soft quantiles (“K-quantile”), and empirical CDF approximations.

4. Neural and Policy Network Integration

Sinkhorn-Sort is integrated into neural architectures via two principal pipelines:

Neural Sinkhorn-Sort (Mena et al., 2018): A permutation-equivariant neural network $a = b = (1/n, ..., 1/n)$0 outputs a score matrix $a = b = (1/n, ..., 1/n)$1, processed by Gumbel noise (optional) and temperature scaling, exponentiation, and Sinkhorn normalization to yield a doubly-stochastic matrix $a = b = (1/n, ..., 1/n)$2. The soft-sorted output is $a = b = (1/n, ..., 1/n)$3. This formulation supports end-to-end backpropagation through sorting, enabling training on global or structured losses.

Sinkhorn Policy Gradient (SPG) (Emami et al., 2018): In reinforcement learning, a Sinkhorn layer is deployed as the actor in an actor-critic framework, producing soft permutation matrices from the input and score network, and using Hungarian rounding at inference to actualize discrete permutations. The critic receives both continuous and discrete assignments, with loss terms to align the critic’s value on relaxed and discrete actions. Training employs temperature annealing, regularization, and straight-through estimators for efficient permutation policy learning.

Comparison with related approaches:

Sinkhorn-Sort outperforms earlier sequence-to-sequence and NeuralSort methods in sorting and ranking accuracy for long sequences, and provides stronger generalization to unfamiliar input distributions (Cuturi et al., 2019, Mena et al., 2018).

5. Transformer and Attention Applications

Sinkhorn-Sort is leveraged in attention mechanisms to achieve quasi-global sparse attention with efficient memory characteristics. In the Sinkhorn Transformer (Tay et al., 2020):

• The input sequence is partitioned into blocks, pooled, and passed through a meta-sorting network to score and permute blocks via a Sinkhorn-approximated permutation matrix $a = b = (1/n, ..., 1/n)$4.
• The sorted blocks are “unpooled,” and attention is computed within and across these re-ordered blocks, yielding content-based locality.
• Variants such as causal Sinkhorn balancing are used to enforce autoregressive properties, and algorithms like SortCut dynamically truncate output sequences post sorting to meet computational budgets.

This structure reduces the complexity of attention from $a = b = (1/n, ..., 1/n)$5 to $a = b = (1/n, ..., 1/n)$6 ($a = b = (1/n, ..., 1/n)$7 = sequence length, $a = b = (1/n, ..., 1/n)$8 = block size), with further reductions via SortCut to $a = b = (1/n, ..., 1/n)$9 for selected top-$C_{ij} = |x_i - t_j|$0 blocks (Tay et al., 2020).

Empirical results: Sinkhorn Transformers match or improve upon vanilla Transformers and other efficient sparse models on language modeling, image generation, algorithmic sequence sorting, and classification tasks, consistently providing memory and compute savings (Tay et al., 2020).

6. Practical Implementation and Optimization

Effective deployment of Sinkhorn-Sort requires careful tuning of hyperparameters:

• Temperature: $C_{ij} = |x_i - t_j|$1 or $C_{ij} = |x_i - t_j|$2 balances permutation sharpness and gradient flow; best performance typically occurs for $C_{ij} = |x_i - t_j|$3.
• Iteration count: Sufficient (e.g., $C_{ij} = |x_i - t_j|$4=10–20) Sinkhorn steps ensure near-doubly-stochastic solutions.
• Numerical stability: Log-domain updates are advisable for small temperature regimes to avoid underflow/overflow.
• Regularization: Penalties enforcing doubly-stochasticity and entropy bonuses may improve stability.
• Parallelism: Batch and tensorized implementations accelerate training and inference.

A variety of strategies, including Gumbel noise for stochastic approximations and temperature annealing, are effective in balancing permutation discreteness with trainability (Cuturi et al., 2019, Mena et al., 2018, Emami et al., 2018, Tay et al., 2020).

7. Extensions, Applications, and Empirical Performance

Table: Representative Sinkhorn-Sort Applications and Results

Domain	Approach	Performance Highlights
Numeric sorting	Neural Sinkhorn-Sort (Mena et al., 2018)	Perfect sorting up to $C_{ij} = |x_i - t_j|$5, >99% at $C_{ij} = |x_i - t_j|$6
Policy learning	SPG (Emami et al., 2018)	KT=0.999 (N=20), 0.985 (N=50)
Quantile regression	S-quantile (Cuturi et al., 2019)	Competitive quantile/MSE error on UCI
Transformer attention	Sinkhorn Transformer (Tay et al., 2020)	Matches/outperforms full attention with less memory

Sinkhorn-Sort has been validated across tasks such as sorting integers, quantile regression, differentiable top-$C_{ij} = |x_i - t_j|$7 losses, attention for long sequence models, and policy learning for assignment and matching. In each, it delivers efficient, differentiable surrogates for discrete permutation operations, often outperforming or matching baselines while enabling scalable learning (Cuturi et al., 2019, Mena et al., 2018, Tay et al., 2020, Emami et al., 2018).

---

Sinkhorn-Sort bridges discrete combinatorial structure and differentiable computation by leveraging entropic OT, providing a foundation for permutation- and ranking-centric deep learning architectures. It enables a wide array of applications—from sorting operators for neural network modules to complexity-efficient sparse attention and policy learning—by harnessing the smooth interpolations within the Birkhoff polytope enabled by the Sinkhorn operator.