Package: differentiable-sorting
Differentiable parallel approximate sorting networksA Python implementation of a fully-differentiableapproximatesorting function for power-of-2 length vectors. Uses Numpy, PyTorch or Tensorflow (or autograd, JAX or cupy), but trivial to use in other backends. Works on GPU.fromdifferentiable_sortingimportbitonic_matrices,diff_sort,diff_argsort# sort 8 element vectorssort_matrices=bitonic_matrices(8)x=[5.0,-1.0,9.5,13.2,16.2,20.5,42.0,18.0]print(diff_sort(sort_matrices,x))>>>[-1.0074.9969.43913.21215.94818.2120.60242.]# we can relax towards averaging by plugging in another# softmax function to the networkfromdifferentiable_sortingimportsoftmax_smoothprint(diff_sort(sort_matrices,x,lambdaa,b:softmax_smooth(a,b,0.0)))# as above>>>[-1.0074.9969.43913.21215.94818.2120.60242.]print(diff_sort(sort_matrices,x,lambdaa,b:softmax_smooth(a,b,0.05)))# smoothed>>>[1.2425.3339.60712.44616.84518.99520.93237.999]print(diff_sort(sort_matrices,x,lambdaa,b:softmax_smooth(a,b,1.0)))# relax completely to mean>>>[15.42515.42515.42515.42515.42515.42515.42515.425]print(np.mean(x))>>>15.425###### Ranking# We can rank as wellx=[1,2,3,4,8,7,6,4]print(diff_argsort(sort_matrices,x))>>>[0.1.2.3.7.6.5.3.]# smoothed ranking functionprint(diff_argsort(sort_matrices,x,sigma=0.25))>>>[0.131.092.3.116.996.5.3.11]# using autograd to differentiate smooth argsortfromautogradimportjacobianjac_rank=jacobian(diff_argsort,argnum=1)print(jac_rank(sort_matrices,np.array(x),0.25))>>>[[2.162-1.059-0.523-0.287-0.01-0.018-0.056-0.21][-0.0660.562-0.186-0.155-0.005-0.011-0.035-0.105][-0.012-0.0130.041-0.005-0.-0.001-0.002-0.008][-0.012-0.025-0.1080.564-0.05-0.086-0.141-0.14][-0.001-0.001-0.003-0.0050.104-0.058-0.028-0.008][-0.-0.001-0.002-0.004-0.0010.028-0.012-0.007][-0.-0.-0.001-0.002-0.016-0.0180.038-0.001][-0.012-0.025-0.108-0.209-0.05-0.086-0.1410.633]]Caveats:May not be very efficient (!), requiring approximatelylog_2(n)^2matrix multiplies of sizen x n. These are just permutations, so can also be efficiently implemented via indexing, requiring2nindexing operations per layer.Dynamic range is limited, especially withfloat32. Very large or very small values will cause trouble. Values distributed between 1 and ~300 work reasonably withfloat64(and similarly for negative values). Values of magnitude 0.0-1.0 are troublesome. Inputs should be pre-normalised (e.g. batch-norm followed by a constant scale by 100)The networks aretheoreticallydifferentiable, but gradients may be very small for larger networks.This appears to be the idea presented inNeuroSortby Grover et al.Grover, A., Wang, E., Zweig, A., & Ermon, S. (2019). Stochastic optimization of sorting networks via continuous relaxations. arXiv preprint arXiv:1903.08850.Vector sortingWe can also sort vector valued entries using a particular "key" function, assuming this function is also differentiable (e.g. sort a matrix by a particular column) usingvector_sort:vector_sort(bitonic_matrices(4),np.array([[1,5],[30,30],[6,9],[80,-2]]),lambdax:(x@[1,0]),# sort by column 1alpha=1.0)>>>array([[80.,-2.],[30.,30.],[5.97,8.97],[1.03,5.03]])vector_sort(bitonic_matrices(4),np.array([[1,5],[30,30],[6,9],[80,-2]]),lambdax:(x@[0,1]),# sort by column 2)>>>array([[30.,30.],[5.91,8.93],[1.16,5.07],[79.93,-1.99]])Bitonic sortingBitonic sortscreate sorting networks with a sequence of fixed conditional swapping operations executed on annelement vector in parallel where,n=2^k.Image: from Wikipedia, by user Bitonic, CC0The sorting network forn=2^kelements hask(k-1)/2"layers" where parallel compare-and-swap operations are used to rearrange anelement vector into sorted order. The sequence of operations is independent of the data.Differentiable compare-and-swapIf we define thesoftmax(a,b)function (not the traditional "softmax" used for classification, but rather log-sum-exp!) as thecontinuous approximation to themax(a,b)function,softmax(a,b) = log(exp(a) + exp(b)).thensoftmin(a,b) = a + b - softmax(a,b).[Alternativelysoftmin(a,b) = -log(exp(-a) + exp(-b))but this adds numerical instability]Notice that we now have a differentiable compare-and-swap operation:softcswap(a,b) = (softmin(a,b), softmax(a,b))We can alternatively use:smoothmax(a,b,alpha)=(a*exp(a*alpha)+b*exp(b*alpha))/(exp(a*alpha)+exp(b*alpha))which has a configurablealphaterm, allowing interpolation between a hard maximum (alpha -> infinity) and mean averaging (alpha -> 0).Softmax/logsumexp/quasimax across the range [-5, 5] for various alpha.Smoothmax across the range [-5, 5] for various alpha.Differentiable sortingFor each layer in the sorting network, we can split all of the pairwise comparison-and-swaps into left-hand and right-hand sides. We can any write function that selects the relevant elements of the vector as a multiply with a binary matrix.For each layer, we can derive two binary matricesleftandrightwhich select the elements to be compared for the left and right hands respectively. This will result in the comparison between twon/2length vectors. We can also derive two matricesl_invandr_invwhich put the results of the compare-and-swap operation back into the right positions in the original vector.The entire sorting network can then be written in terms of matrix multiplies and thesoftcswap(a, b)operation.defsoftmax(a,b):returnnp.log(np.exp(a)+np.exp(b))defsoftmin(a,b):returna+b-softmax(a,b)defsoftcswap(a,b):returnsoftmin(a,b),softmax(a,b)defdiff_bisort(matrices,x):forl,r,l_inv,r_invinmatrices:a,b=softcswap(l@x,r@x)x=l_inv@a+r_inv@breturnxThe rest of the code is simply computing thel, r, l_inv, r_invmatrices, which are fixed for a givenn.If we're willing to include a split and join operation, we can reduce this to a singlen x nmultiply for each layer (plus one extra multiply at the start) by fusing consecutive permutations:defdiff_sort_weave(fused,x,softmax=np.maximum):""" Given a set of bitonic sort matrices generated bybitonic_woven_matrices(n), sort a sequence x of length n. """split=len(x)//2x=fused[0]@x# initial permutationformatinfused[1:]:a,b=x[:split],x[split:]mx,mn=softcswap(a,b)x=mat@np.concatenate([mn,mx])returnxLibrariesThe base code works with NumPy. If you want to useautogradjaxorcupythen install theautoraypackage.pip install autorayThe code should then automatically work with whatever backend you are using if you pass arrays from the right backend.PyTorch and TensorflowPyTorch is supported:importtorchfromdifferentiable_sorting.torchimportbitonic_matrices,diff_sortfromtorch.autogradimportVariablex=[5.0,-1.0,9.5,13.2,16.2,20.5,42.0,18.0]matrices=bitonic_matrices(8)torch_input=Variable(torch.from_numpy(np.array(x)).float(),requires_grad=True)result=diff_sort(matrices,torch_input)print(result)>>>tensor([-1.0075,4.9958,9.4394,13.2117,15.9480,18.2103,20.6023,42.0000],grad_fn=<AddBackward0>)print(torch.autograd.grad(result[0],torch_input)[0])>>>tensor([7.3447e-03,9.9257e-01,8.1266e-05,3.9275e-06,3.4427e-08,1.4447e-09,0.0000e+00,9.5952e-09])Tensorflow is also supported:importtensorflowastffromdifferentiable_sorting.tensorflowimportbitonic_matrices,diff_sort,diff_argsorttf_input=tf.reshape(tf.convert_to_tensor([5.0,-1.0,9.5,13.2,16.2,20.5,42.0,18.0],dtype=tf.float64),(-1,1))tf_output=tf.reshape(diff_sort(tf_matrices,tf_input),(-1,))withtf.Session()ass:print(s.run((tf_output)))>>>[-1.0074.9969.43913.21215.94818.2120.60242.]This implementation was inspired bythis tweetby @francoisfleuret:FranÃ§ois Fleuret @francoisfleuret Jun 14Discussion with Ronan Collober reminded me that (max, +) is a semi-ring, and made me realize that the same is true for (softmax, +) wheresoftmax(a, b) = log(exp(a) + exp(b))All this coolness at the heart of his paperhttps://arxiv[dot] org/abs/1902.06022Error analysisThe plot below shows the relative RMS (RMS error divided by the maximum range of the input vector) between the softmax sorted array and the ground truth sorted array, for vectors of lengthn=2throughn=512, with test values distributed randomly uniformly in ranges from [0, 2^-5] to [0, 2^10]. The main factor affecting precision is the numerical range. Small values will be corrupted, but values > ~300 will overflow (infloat64). The dotted line is shown at 10% relative error.Softmaxerror curve.Smoothmaxerror curve.There is also a function to pretty-print bitonic networks:pretty_bitonic_network(8)0  1  2  3  4  5  6  7 
    â•­â”€â”€â•¯  â”‚  â”‚  â”‚  â”‚  â”‚  â”‚ 
    â”‚  â”‚  â•°â”€â”€â•®  â”‚  â”‚  â”‚  â”‚ 
    â”‚  â”‚  â”‚  â”‚  â•­â”€â”€â•¯  â”‚  â”‚ 
    â”‚  â”‚  â”‚  â”‚  â”‚  â”‚  â•°â”€â”€â•® 
    â•­â”€â”€â”€â”€â”€â•¯  â”‚  â”‚  â”‚  â”‚  â”‚ 
    â”‚  â•­â”€â”€â”€â”€â”€â•¯  â”‚  â”‚  â”‚  â”‚ 
    â”‚  â”‚  â”‚  â”‚  â•°â”€â”€â”€â”€â”€â•®  â”‚ 
    â”‚  â”‚  â”‚  â”‚  â”‚  â•°â”€â”€â”€â”€â”€â•® 
    â•­â”€â”€â•¯  â”‚  â”‚  â”‚  â”‚  â”‚  â”‚ 
    â”‚  â”‚  â•­â”€â”€â•¯  â”‚  â”‚  â”‚  â”‚ 
    â”‚  â”‚  â”‚  â”‚  â•°â”€â”€â•®  â”‚  â”‚ 
    â”‚  â”‚  â”‚  â”‚  â”‚  â”‚  â•°â”€â”€â•® 
    â•­â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â•¯  â”‚  â”‚  â”‚ 
    â”‚  â•­â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â•¯  â”‚  â”‚ 
    â”‚  â”‚  â•­â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â•¯  â”‚ 
    â”‚  â”‚  â”‚  â•­â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â”€â•¯ 
    â•­â”€â”€â”€â”€â”€â•¯  â”‚  â”‚  â”‚  â”‚  â”‚ 
    â”‚  â•­â”€â”€â”€â”€â”€â•¯  â”‚  â”‚  â”‚  â”‚ 
    â”‚  â”‚  â”‚  â”‚  â•­â”€â”€â”€â”€â”€â•¯  â”‚ 
    â”‚  â”‚  â”‚  â”‚  â”‚  â•­â”€â”€â”€â”€â”€â•¯ 
    â•­â”€â”€â•¯  â”‚  â”‚  â”‚  â”‚  â”‚  â”‚ 
    â”‚  â”‚  â•­â”€â”€â•¯  â”‚  â”‚  â”‚  â”‚ 
    â”‚  â”‚  â”‚  â”‚  â•­â”€â”€â•¯  â”‚  â”‚ 
    â”‚  â”‚  â”‚  â”‚  â”‚  â”‚  â•­â”€â”€â•¯