SoftSort: Differentiable Argsort Relaxation

Updated 22 February 2026

SoftSort is a continuous, differentiable relaxation of the argsort operator that maps score vectors to soft permutation matrices via a softmax over pairwise distances.
It provides mathematical tractability with convergence to an exact permutation matrix as the temperature parameter approaches zero, enabling seamless integration into gradient-based models.
Empirical results show SoftSort’s efficiency and accuracy benefits in tasks such as ranking, algorithm unrolling, and differentiable pooling compared to traditional sorting methods.

SoftSort is a continuous, differentiable relaxation of the discrete argsort operator, designed to enable sorting-based and ranking-based objectives within gradient-based optimization frameworks. By mapping a vector of scores to a soft permutation matrix via a row-wise softmax over pairwise distances between sorted and unsorted entries, SoftSort provides both mathematical tractability and computational efficiency, with applications in learning-to-rank, algorithm unrolling, differentiable pooling, and large-scale permutation learning. As the temperature parameter approaches zero, SoftSort converges to the exact permutation matrix, providing a tunable interpolation between smooth differentiability and hard assignment.

1. Formal Definition and Mathematical Properties

For $s=(s_1,\dots,s_n)\in\mathbb{R}^n$ , the goal is to approximate the permutation matrix $P_{\argsort(s)}$ encoding the argsort (indices that sort $s$ in descending order). SoftSort replaces this with a differentiable operator $P_{\mathrm{SoftSort}}(s)\in [0,1]^{n\times n}$ , defined as

$P_{ij} = \frac{\exp(-|u_i-s_j|/\tau)}{\sum_{k=1}^n \exp(-|u_i-s_k|/\tau)}$

where $u_i$ are the order statistics of $s$ (i.e., $u = \mathrm{sort}(s)$ in decreasing order), and $\tau > 0$ is a temperature hyperparameter controlling sharpness. As $\tau \to 0$ , $P \to P_{\argsort(s)}$.

Key mathematical properties of SoftSort include:

Row-stochasticity and non-negativity: Each row of $P$ sums to 1 with non-negative entries.
Asymptotic exactness: $\lim_{\tau \rightarrow 0} P_{\mathrm{SoftSort}}(s) = P_{\argsort(s)}$.
Permutation equivariance: $\mathrm{SoftSort}_\tau(s) = \mathrm{SoftSort}_\tau(\mathrm{sort}(s))\,P_{\argsort(s)}$.
Lipschitz continuity: Each row mapping is $(\sqrt{n}+1)/\tau$ -Lipschitz in $s$ (Mohammad-Taheri et al., 21 May 2025).
Closed-form differentiability: The Jacobian of the mapping with respect to $s$ follows directly from the softmax gradient.

The typical computational complexity is $O(n^2)$ per forward pass (dominated by pairwise distance matrix and softmax), with a minimal code footprint—a three-line implementation suffices (Prillo et al., 2020).

2. Motivations and Theoretical Basis

Classical sorting and ranking operations are non-differentiable, with the discrete permutation matrix exhibiting zero gradients almost everywhere. This renders them incompatible with gradient-based learning paradigms central to modern machine learning. Existing relaxations, including NeuralSort and Sinkhorn sorts, either incur higher complexity or lack the directness of SoftSort.

SoftSort addresses this by providing:

A direct relaxation closely tied to order statistics—each softmax row corresponds to a sorted position.
Avoidance of iterative optimization (unlike Sinkhorn sorts that require numerous row/column normalizations).
Theoretical guarantees on convergence, row-unimodality, and differentiability almost everywhere.

The operator’s design allows its use as a replacement for hard argsort in end-to-end differentiable pipelines, without introducing spurious gradients or cumbersome computational overhead (Prillo et al., 2020).

3. Algorithmic Integration and Practical Implementations

SoftSort’s relaxations enable algorithmic integration in a diverse array of models:

Ranking and Structured Prediction: Neural networks can be trained with ranking losses or other custom objectives that require permutation invariance, with SoftSort bridging the gap between hard sorting and backpropagation (Petersen et al., 2024).
Greedy Algorithm Unrolling: In unrolled greedy sparse recovery algorithms such as OMP and IHT, the non-differentiable argsort operator is replaced by SoftSort, enabling differentiable analogues such as Soft-OMP and Soft-IHT. These can be directly unrolled into trainable networks (e.g., OMP-Net, IHT-Net) (Mohammad-Taheri et al., 21 May 2025).
Vision-LLM Fusion: In TS-VLM, SoftSort underpins the Text-Guided SoftSort Pooling (TGSSP) for multi-view aggregation in real-time driving reasoning, offering query-driven feature fusion in place of heavy cross-attention (Chen et al., 19 May 2025).
Permutation Learning and Scalability: Extensions such as ShuffleSoftSort leverage the $O(N)$ parameterization of SoftSort for large-scale layouts (e.g., self-organizing Gaussian splatting), broadening applicability to large $N$ without quadratic storage (Barthel et al., 17 Mar 2025).

Minimal pseudocode for 1D SoftSort is:

def softsort(scores, tau):
    sorted_scores = sort(scores, descending=True)
    logits = -abs(sorted_scores[:, :, None] - scores[:, None, :]) / tau
    P_hat = softmax(logits, axis=-1)
    return P_hat

Extensions to multidimensional settings may involve iterative shuffling and axis-wise application, as in ShuffleSoftSort (Barthel et al., 17 Mar 2025).

4. Extensions: Newton Losses and Beyond

Optimization with SoftSort loss functions can be challenging due to non-convexity and unstable gradients (vanishing/exploding). Newton Losses introduce a second-order loss reshaping mechanism:

Compute the SoftSort cross-entropy loss $\ell_{\mathrm{SS}}(y)$ .
Obtain first-order (gradient $g = \nabla_y \ell_{\mathrm{SS}}$ ) and second-order (Hessian $H$ or empirical Fisher $F$ ) information.
Update using a Newton-type step: $z^\star = y - (C + \lambda I)^{-1} g$
Optimize the network to regress toward $z^\star$ via squared error

$\ell^*(y) = \frac{1}{2} \|y - z^\star\|^2$

This procedure accelerates training and stabilizes convergence, especially in batched or high-dimensional regimes. Empirically, it improves fully correct ranking rates (e.g., for $n=10$ MNIST sorting, baseline SoftSort $27.46\%$ , SoftSort+NL $\sim 54\%$ ) and per-element accuracies, with consistent gains across both Hessian and Fisher-based variants (Petersen et al., 2024).

5. Applications and Empirical Results

SoftSort’s differentiable relaxation is central to a range of empirical achievements:

Sorting, Quantile Regression, $k$ -NN: Achieves or surpasses NeuralSort and OT-based methods for permutation recovery, quantile prediction, and differentiable nearest-neighbor classification—matching or exceeding accuracy with reduced runtime (SoftSort 40–80% faster, up to 6 $\times$ at high $n$ ) (Prillo et al., 2020).
Greedy Sparse Recovery Networks: OMP-Net and IHT-Net with SoftSort-based differentiable selection/thresholding outperform traditional greedy algorithms by an order of magnitude in noise-limited recovery, with empirically validated trade-offs in the temperature parameter $\tau$ for balancing smoothness and accuracy (Mohammad-Taheri et al., 21 May 2025).
Vision-Language Aggregation: TS-VLM’s TGSSP (SoftSort with Sinkhorn projection) outperforms cross-attention and pooling alternatives across BLEU-4, METEOR, ROUGE-L, and CIDEr on the DriveLM benchmark, at a fraction of the computational and memory overhead (e.g., SoftSort $184$ MFLOPs vs. full SinkhornSort $33216$ MFLOPs) (Chen et al., 19 May 2025).
Permutation Learning at Scale: ShuffleSoftSort enables high-quality, low-memory permutation learning for large grids (e.g., $N=1024$ ; $0.854$ DPQ, half the runtime of plain SoftSort; versus $0.913$ DPQ for Gumbel-Sinkhorn at $O(N^2)$ storage), and is suitable for millions of elements in applications such as Self-Organizing Gaussians (Barthel et al., 17 Mar 2025).

6. Limitations, Hyperparameterization, and Variants

Several critical factors influence SoftSort's performance:

Temperature Parameter ( $\tau$ ): Controls the trade-off between smoothness and accuracy. Lower $\tau$ increases accuracy but can induce steep gradients and instability; larger $\tau$ enhances smoothness at the cost of permutation precision (Mohammad-Taheri et al., 21 May 2025).
Computational Bottlenecks: Forward and backward passes require $O(n^2)$ time for $n$ items. While parameterization is minimal, scaling to extreme $N$ demands further algorithmic advances (Barthel et al., 17 Mar 2025).
Memory Considerations: Standard SoftSort stores only $O(N)$ parameters, while some OT/Sinkhorn approaches require $O(N^2)$ .
Gradient Instability: Small $\tau$ can cause gradient explosion or vanishing; Newton Losses mitigate this via curvature-aware corrections (Petersen et al., 2024).
Multidimensional and Complex Permutations: SoftSort is inherently 1D; ShuffleSoftSort or iterative axis-wise application are necessary for high-quality multidimensional permutation learning (Barthel et al., 17 Mar 2025).
Row/Column Bistochasticity: The basic formulation is row-stochastic, not doubly stochastic, unless further projections (e.g., Sinkhorn) are applied (Chen et al., 19 May 2025).

Typical hyperparameters include $\tau\in\{0.001, 0.01, 0.1, 1\}$ , regularization weights for entropic/transport variants ( $\epsilon$ ), number of Sinkhorn iterations ( $T$ ), and damping parameters for Newton Losses ( $\lambda$ ). Empirical ablations consistently show SoftSort attaining optimal efficiency-accuracy trade-offs among contemporary pooling, sorting, and permutation-learning methods (Chen et al., 19 May 2025).

SoftSort is part of a family of differentiable sorting operators:

NeuralSort: Pairwise differences, more complex combinatorics, similar asymptotic and complexity profile but less concise implementation (Prillo et al., 2020).
Sinkhorn/Optimal Transport: Doubly-stochastic relaxations with additional entropic smoothing, extended via iterative projections (Chen et al., 19 May 2025).
Gumbel-Sinkhorn and Low-Rank Factorizations: Permit injection of noise or low-rank representations for regularized or scalable permutation learning, at higher memory cost (Barthel et al., 17 Mar 2025).

Convergence proofs, Lipschitz estimates, and error bounds for SoftSort-based approximations are well characterized (see Propositions 3.1–3.2 and error theorems in (Mohammad-Taheri et al., 21 May 2025)). The operator’s intuitive structure—a softmax over absolute distance to order statistics—distills the permutation learning problem to its essential mathematical core.

SoftSort is thus established as an efficient, theoretically principled, and broadly applicable differentiable relaxation of the argsort operator—enabling seamless integration of sorting and ranking operations within modern machine learning systems.