Permutation Learning with Only N Parameters: From SoftSort to Self-Organizing Gaussians

TL;DR

This work tackles the memory bottlenecks of permutation learning by introducing ShuffleSoftSort, a differentiable method that learns permutations with only $N$ parameters, unlike Gumbel-Sinkhorn's $O(N^2)$ memory. By iteratively shuffling indices and applying SoftSort for $R$ steps under a temperature schedule $\tau$, the approach preserves previous ordering while enabling more flexible, multidimensional sorting. The method incorporates a row-wise computation and a loss combining neighborhood, stochastic, and standard-deviation terms to converge to a valid permutation, achieving high-quality results with substantially reduced memory. This enables scalable permutation learning for large-scale tasks such as grid-based image sorting and self-organizing Gaussian representations in 3D scene reconstruction, with practical storage reductions and end-to-end differentiability.

Abstract

Sorting and permutation learning are key concepts in optimization and machine learning, especially when organizing high-dimensional data into meaningful spatial layouts. The Gumbel-Sinkhorn method, while effective, requires N*N parameters to determine a full permutation matrix, making it computationally expensive for large datasets. Low-rank matrix factorization approximations reduce memory requirements to 2NM (with M << N), but they still struggle with very large problems. SoftSort, by providing a continuous relaxation of the argsort operator, allows differentiable 1D sorting, but it faces challenges with multidimensional data and complex permutations. In this paper, we present a novel method for learning permutations using only N parameters, which dramatically reduces storage costs. Our method extends SoftSort by iteratively shuffling the N indices of the elements and applying a few SoftSort optimization steps per iteration. This modification significantly improves sorting quality, especially for multidimensional data and complex optimization criteria, and outperforms pure SoftSort. Our method offers improved memory efficiency and scalability compared to existing approaches, while maintaining high-quality permutation learning. Its dramatically reduced memory requirements make it particularly well-suited for large-scale optimization tasks, such as "Self-Organizing Gaussians", where efficient and scalable permutation learning is critical.

Figures (6)

• Figure 1: Example of grid-based sorting for 1024 random RGB colors sorted by SoftSort (left) and the new proposed approach using the newly proposed ShuffleSoftSort (right). The loss function minimizes the average color distance of neighboring grid cells.
• Figure 2: Permutation learning optimizes the differentiable permutation matrix $P_{\text{soft}}$ by adjusting the weights $w$ based on a given loss function. The input vectors $x$ are reordered into $x_{\text{sort}}$ using $P_{\text{hard}}$, the binarized form of $P_{\text{soft}}$.
• Figure 3: 1D color arrangement highlighting SoftSort's challenges (see text).
• Figure 4: ShuffleSoftSort improves sorting by iteratively applying SoftSort to shuffled elements of the vector $x$. The $N$ elements are first randomly shuffled and then sorted using SoftSort over a few iterations, with the loss computed on the reverse-shuffled output. Repeating this process progressively refines the permutation and helps overcome the inherent limitations of SoftSort.
• Figure 5: Sorting example of a dataset of e-commerce images, simplifying navigation, browsing, and retrieval of large image databases.