Implementing Differentiable Optimal Transport: A Case Study

in this video i'll present a case study on implementing differentiable optimal transport in pi torch we will look at two different methods the first method unrolling the forward pass optimization algorithm is easy to implement but slow and memory intensive in the backward pass the second method uses implicit differentiation based on deep declarative networks it requires more work to implement including some fancy linear algebra tricks but is fast and very memory efficient let's start by looking at the optimal transport problem there are plenty of papers and online resources that discuss the motivation theory and applications of optimal transport for us it will suffice to simply examine the resulting optimal transport optimization problem given a real matrix m optimal transport finds the positive matrix p that minimizes the inner product between m and p and an entropy regularization term we also require p to satisfy row sum and column sum constraints here the scalar quantity gamma controls the strength of the regularization term and vectors r and c encode the row sum and column sum constraints they must each sum to 1. when r and c are uniform vectors the matrix p will be doubly stochastic so optimal transport performs a mapping from an m by n real matrix denoted capital m to an m by n positive matrix denoted capital p whose entry satisfy constraints on their row and column sums it can be shown that the solution to the entropy regularized optimal transport problem takes the following form moreover there is an efficient iterative algorithm for finding the solution the sinkhole knob algorithm starts with a positive matrix and alternately rescales its rows and columns until convergence here we show the complete algorithm for finding p and here is a pi torch implementation of the algorithm that operates on a batch of inputs the first few lines just extract the problem sizes and set parameters r and c to uniform distributions if not provided if r and c are provided they must sum to one next we initialize p to the positive matrix obtained by exponentiating scaled entries of the input m the scaling factor is minus gamma next we perform row scaling and then column scaling we also check for convergence and finally we return p given that we have a simple iterative algorithm for computing the solution one way to back propagate gradients is to simply unroll the algorithm and apply the chain rule backwards through the iterations this is exactly what automatic differentiation will do in pytorch calling our sinkhorn function on a matrix m with required underscore grad equals true does the trick a downside of this approach is that all the intermediate calculations need to be remembered in the forward pass which may be prohibitive for large problems nevertheless the approach serves as a very effective baseline implementation especially given that autograd does all the heavy lifting of calculating the gradients and applying the chain rule for us an alternative approach is to consider the problem from the perspective of deep declarative networks this allows us to frame layers or processing nodes in an end-to-end learnable model such as a neural network as optimization problems and provides the tools for back propagating through these problems by implicit differentiation of the optimality conditions the key operation that we need to perform during back propagation is the calculation of the derivative of the loss function with respect to the input m which by the chain rule can be written as the product of the derivative of the loss function with respect to the output p multiplied by the derivative of the output with respect to the input the deep declarative network's paper gives this result for that calculation where the input to the backward pass dj dp is written here as v transpose and the derivative of the output with respect to the input is given by the remainder of the expression now formally dj dm and djdp are m by n matrices indeed this is how they are propagated through the network i.e as tensors the same size as the corresponding data in the forward pass this makes dpdm a four-dimensional tensor to apply the deep declarative network result we need to conceptually vectorize the inputs and outputs you can think of this as flattening the data matrices row wise to produce a vector of length m n as illustrated here as we will see later this is only for thinking through the mathematical operations in the actual implementation we can leave the inputs and outputs in their native matrix or tensor forms let us consider the objective function again written out this time with explicit sums the i j k l f entry within the b matrix is the mixed partial derivative of the objective function with respect to p i j and m k l we can see that this is one if i j equals k l and zero otherwise so b is simply a very large identity matrix since multiplying by an identity matrix leaves the multiplicand unchanged we can treat it as a no op and simply drop b from our calculations moving on to h the i j k off entry of h is the second derivative of the objective function with respect to p i j and p k l since the objective decomposes over the entries of p the resulting matrix h is diagonal with diagonal entries 1 over gamma times p i j the inverse is therefore trivial to compute moreover multiplication by h inverse simply scales corresponding entries by gamma p i j which we write succinctly as gamma times p arrow computing v transpose h inverse in vectorized form would look as shown however we can perform the same operation keeping our representation in matrix form by calculating an element-wise product so we have a very efficient way of multiplying by h inverse now let's turn our attention to the constraints which define the matrix a we have m plus n constraints in total it turns out however that if any m plus n minus 1 of the constraints are satisfied then the last constraint will automatically be satisfied this is because the sum of the row sums must equal the sum of the column sums we can therefore remove one of the constraints in fact for the matrix a to be full rank a requirement of corollary 4.9 from the deep declarative network's paper we must remove one constraint we choose to remove the first one we can now write out the matrix a explicitly the top m minus n rows encode the first set of constraints that is those on the row sums of p and the bottom n rows encode the second set of constraints that is those on the column sums of p the matrix a has m times n columns which is the number of entries in the matrix p let us consider a concrete small 3x3 example in this example it is easy to see that the first row is a linear combination of the other rows specifically it can be constructed by adding the last three rows and subtracting rows two and three remembering that r and c sum to one re-examining our expression for the backward going gradient it appears that we now have everything we need v transpose is the incoming gradient of the loss with respect to the output h inverse is a diagonal matrix with diagonal entries gamma p i j b is an mn by mn identity matrix and a is a matrix of coefficients encoding our constraints blindly forming these matrices and evaluating this expression would be computationally intractable let's dig into the different parts in some more detail consider the expression a h inverse a transpose we can readily see that this is a square matrix of dimension m plus n minus 1 by m plus n minus 1. each entry can be written as the double summation shown the resulting matrix has a very simple structure composed of four submatrices diagonal matrices for the 1 1 and 2 2 blocks and the matrix p and p transpose for the 1 2 and 2 1 blocks respectively here the first row and column are removed consistent with the first constraint being removed from a as discussed earlier we need to compute the inverse of this matrix which can be done efficiently using a block matrix inversion formula here we show the explicit formulae for each block putting this all together we can write our previous expression for the backward going gradient dj dm as shown let's look at a pi torch implementation for this expression that operates on batch inputs as our synchorn function did before first we compute the last term minus gamma v transpose p this is what the gradient would be if we didn't run any sinkhorn iterations it's a good approximation if the constraints are nearly satisfied next we compute gamma v transpose pa or actually it's negative reusing the calculation we just made we separate the calculation into two parts based on how the top and bottom blocks of matrix a operate on the multiplicand we then compute the inverse of a h inverse a transpose using block matrix inversion as discussed earlier here instead of computing the inverse directly we decompose the one one block into a lower triangular matrix multiplied by its transpose via celeski factorization we can then compute the one-two block using celeski solve which makes the operation very efficient with the blocks computed we can perform matrix multiplication with the previously computed gamma v transpose pa term once again we make use of cholesky solve rather than directly inverting the 1 1 block and once again we separate the calculation into two parts finally we post multiply by a p and subtract from our initial calculation of gamma v transpose p to get the exact backward going derivative of the loss function with respect to the input we can compare the running time of unrolling the synchorn iterates using automatic differentiation with implicit differentiation here we show the running time for computing the full inverse of a h inverse a transpose as well as the block inverse variant we also showed the running time for the approximate gradient the implicit differentiation method with block inverse is slightly faster than automatic differentiation of synchorn and significantly faster than taking the full inverse a similar trend can be seen on the gpu where all methods can make use of parallelization for larger batch sizes more impressively the implicit differentiation method does not require storing intermediate calculations in the forward pass making it much more memory efficient in addition to being slightly faster here we show memory used for a single forward and backward pass as a function of the number of synchorn iterations and a function of problem size for fixed number 10 of synchorn iterations our hard work has been worth it you can check out full code and examples in the deep declarative networks repository visit http deepdeclarativenetworks.com thank you for watching