Fast differentiable sorting and ranking

[BatyLeo]

Following this discourse post, here is a Julia implementation of Fast differentiable sorting and ranking (see here for the original Python implementation).

This includes:

• soft_rank and soft_sort differentiable methods
• options for choosing regularization: l2 or entropic
• Two new custom optimization layers:SoftRank <: AbstractRegularized and SoftSort <: AbstractOptimizationLayer compatible with the InferOpt ecosystem

---

Reproducing figures from the paper:

using InferOpt, Plots, Zygote

plot(x -> sort([0.0, x, 1.0, 2.0]; rev=true)[2], label="Sort")
plot!(x -> soft_sort_l2([0.0, x, 1.0, 2.0]; rev=true)[2], label="Soft sort l2")
plot!(y -> gradient(x -> soft_sort_l2([0.0, x, 1.0, 2.0]; ε=1.0, rev=true)[2], y)[1], label="Soft sort l2 derivative")

plot(x -> sort([0.0, x, 1.0, 2.0]; rev=true)[2], label="Sort")
plot!(x -> soft_sort_kl([0.0, x, 1.0, 2.0]; rev=true)[2], label="Soft sort kl")
plot!(y -> gradient(x -> soft_sort_kl([0.0, x, 1.0, 2.0]; ε=1.0, rev=true)[2], y)[1], label="Soft sort kl derivative")

plot(x -> ranking([x, 3.0, 1.0, 2.0]; rev=true)[1], label="Ranking")
plot!(x -> soft_rank_l2([x, 3.0, 1.0, 2.0]; rev=true)[1], label="Soft rank l2")
plot!(y -> gradient(x -> soft_rank_l2([x, 3.0, 1.0, 2.0]; ε=1.0, rev=true)[1], y)[1], label="Soft rank l2 derivative")

plot(x -> ranking([x, 3.0, 1.0, 2.0]; rev=true)[1], label="Ranking")
plot!(x -> soft_rank_kl([x, 3.0, 1.0, 2.0]; rev=true)[1], label="Soft rank kl")
plot!(y -> gradient(x -> soft_rank_kl([x, 3.0, 1.0, 2.0]; ε=1.0, rev=true)[1], y)[1], label="Soft rank kl derivative")

[codecov-commenter]

Codecov Report

Patch coverage: 91.66% and project coverage change: +1.95% 🎉

Comparison is base (f9d8dab) 85.34% compared to head (ff85e4a) 87.29%.

❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #88      +/-   ##
==========================================
+ Coverage   85.34%   87.29%   +1.95%     
==========================================
  Files          18       23       +5     
  Lines         389      559     +170     
==========================================
+ Hits          332      488     +156     
- Misses         57       71      +14     

Files Changed	Coverage Δ
src/InferOpt.jl	100.00% <ø> (ø)
src/regularized/soft_rank.jl	60.00% <60.00%> (ø)
src/utils/isotonic_regression/isotonic_kl.jl	100.00% <100.00%> (ø)
src/utils/isotonic_regression/isotonic_l2.jl	100.00% <100.00%> (ø)
src/utils/isotonic_regression/projection.jl	100.00% <100.00%> (ø)

... and 1 file with indirect coverage changes

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.