Awesome

Code for "SoftSort: A Continuous Relaxation for the argsort Operator", ICML 2020.

This repository is a fork of ermongroup/neuralsort implementing the SoftSort operator and reproducing all the results reported in the paper "SoftSort: A Continuous Relaxation for the argsort Operator".

Requirements

The codebase is implemented in Python 3.7. To install the necessary requirements, run the following commands:

pip3 install -r requirements.txt

Sorting Handwritten Numbers Experiment

To reproduce the results in Table 1, just run:

cd tf
bash run_sort.sh
python3 run_sort_table_of_results.py

The first script (bash) will train all models. This takes a long time. You can inspect this script to see what parameters were used to train each model (which are the ones reported in the paper). The second script (python) will process the results from the models and print Table 1.

To train a single model directly, you can use the tf/run_sort.py script, with the following arguments:

  --M INT                 Minibatch size
  --n INT                 Number of elements to compare at a time
  --l INT                 Number of digits in each multi-mnist dataset element
  --tau FLOAT             Temperature (either of sinkhorn or neuralsort relaxation)
  --method STRING         One of 'deterministic_neuralsort', 'stochastic_neuralsort', 'deterministic_softsort', 'stochastic_softsort'
  --n_s INT               Number of samples for stochastic methods
  --num_epochs INT        Number of epochs to train
  --lr FLOAT              Initial learning rate

Quantile Regression Experiment

To reproduce the results in Table 2, just run:

cd tf
bash run_median.sh
python3 run_median_table_of_results.py

The first script (bash) will train all models. This takes a long time. You can inspect this script to see what parameters were used to train each model (which are the ones reported in the paper). The second script (python) will process the results from the models and print Table 2.

To train a single model directly, you can use the tf/run_median.py script, with the following arguments:

  --M INT                 Minibatch size
  --n INT                 Number of elements to compare at a time
  --l INT                 Number of digits in each multi-mnist dataset element
  --tau FLOAT             Temperature (either of sinkhorn or neuralsort relaxation)
  --method STRING         One of 'deterministic_neuralsort', 'stochastic_neuralsort', 'deterministic_softsort', 'stochastic_softsort'
  --n_s INT               Number of samples for stochastic methods
  --num_epochs INT        Number of epochs to train
  --lr FLOAT              Initial learning rate

Differentiable kNN Experiment

To reproduce the results in Table 3, run:

cd pytorch
bash run_dknn.sh
python3 run_dknn_table_of_results.py

The first script (bash) will train all the models. This takes about two days to sequentally test the different hyperparameter configurations. The seconds script iterates through logs and prints the best results.

To train a single model directly, you can use the pytorch/run_dknn.py script, with the following arguments:

  --simple                Whether to use our softsort, or the baseline neuralsort
  --k INT                 Number of nearest neighbors
  --tau FLOAT             Temperature of sorting operator
  --nloglr FLOAT          Negative log10 of learning rate
  --method STRING         One of 'deterministic', 'stochastic'
  --dataset STRING        One of 'mnist', 'fashion-mnist', 'cifar10'
  --num_train_queries INT Number of queries to evaluate during training.
  --num_train_neighbors INT Number of neighbors to consider during training.
  --num_samples INT       Number of samples for stochastic methods
  --num_epochs INT        Number of epochs to train

Speed Comparison Experiment

To reproduce the results in Figure 6, just run:

bash synthetic_experiment_speed_comparison.sh
python3 synthetic_experiment_speed_comparison_plot.py

The first script (bash) will train all models. This takes some time. The second script (python) will process the results and print the graphs in Figure 6 under the images/ directory.

Learning Curves

For the synthetic experiment learning curves, run:

bash synthetic_experiment_learning_curves.sh

Then, to generate the plot in Figure 8, run:

python3 synthetic_experiment_learning_curves_plot.py

To generate the run_sort and run_median learning curve plots (Figure 7), run:

cd tf
python3 run_sort_learning_curves.py
python3 run_median_learning_curves.py