diffsort - Differentiable Sorting Networks

Official implementation for our ICML 2021 Paper "Differentiable Sorting Networks for Scalable Sorting and Ranking Supervision". In this work, we leverage classic sorting networks and relax them to propose a new differentiable sorting function: diffsort. This allows propagating gradients through (an approximation of) the sorting / ranking function / operation. Herein, diffsort outperforms existing differentiable sorting functions on the four-digit MNIST and the SVHN sorting tasks. In this repo, we present the PyTorch implementation of our ICML 2021 paper on differentiable sorting networks. Paper @ ArXiv, Video @ Youtube.

Installation

diffsort can be installed via pip from PyPI with

pip install diffsort

Or from source, e.g., in a virtual environment like

virtualenv -p python3 .env1
. .env1/bin/activate
pip install .

Usage

import torch
from diffsort import DiffSortNet

vector_length = 2**4
vectors = torch.randperm(vector_length, dtype=torch.float32, device='cpu', requires_grad=True).view(1, -1)
vectors = vectors - 5.

# sort using a bitonic-sorting-network
sorter = DiffSortNet('bitonic', vector_length, steepness=5)
sorted_vectors, permutation_matrices = sorter(vectors)
print(sorted_vectors)

Experiments

Will be published soon.

Citing

@inproceedings{Petersen2021-diffsort,
  title={Differentiable Sorting Networks for Scalable Sorting and Ranking Supervision},
  author={Petersen, Felix and Borgelt, Christian and Kuehne, Hilde and Deussen, Oliver},
  booktitle={International Conference on Machine Learning (ICML)},
  year={2021}
}

License

diffsort is released under the MIT license. See LICENSE for additional details about it.