Learning Representations of Sets through Optimized Permutations

1 INTRODUCTION . Consider a task where each input sample is a set of feature vectors with each feature vector describing an object in an image ( for example : person , table , cat ) . Because there is no a priori ordering of these objects , it is important that the model is invariant to the order that the elements appear in the set . However , this puts restrictions on what can be learned efficiently . The typical approach is to compose elementwise operations with permutation-invariant reduction operations , such as summing ( Zaheer et al. , 2017 ) or taking the maximum ( Qi et al. , 2017 ) over the whole set . Since the reduction operator compresses a set of any size down to a single descriptor , this can be a significant bottleneck in what information about the set can be represented efficiently ( Qi et al. , 2017 ; Le & Duan , 2018 ; Murphy et al. , 2019 ) . We take an alternative approach based on an idea explored in Vinyals et al . ( 2015a ) , where they find that some permutations of sets allow for easier learning on a task than others . They do this by ordering the set elements in some predetermined way and feeding the resulting sequence into a recurrent neural network . For instance , it makes sense that if the task is to output the top-n numbers from a set of numbers , it is useful if the input is already sorted in descending order before being fed into an RNN . This approach leverages the representational capabilities of traditional sequential models such as LSTMs , but requires some prior knowledge of what order might be useful . Our idea is to learn such a permutation purely from data without requiring a priori knowledge ( section 2 ) . The key aspect is to turn a set into a sequence in a way that is both permutation-invariant , as well as differentiable so that it is learnable . Our main contribution is a Permutation-Optimisation ( PO ) module that satisfies these requirements : it optimises a permutation in the forward pass of a neural network using pairwise comparisons . By feeding the resulting sequence into a traditional model such as an LSTM , we can learn a flexible , permutation-invariant representation of the set while avoiding the bottleneck that a simple reduction operator would introduce . Techniques used in our model may also be applicable to other set problems where permutation-invariance is desired , building on the literature of approaches to dealing with permutation-invariance ( section 3 ) . In four different experiments , we show improvements over existing methods ( section 4 ) . The former two tasks measure the ability to learn a particular permutation as target : number sorting and image mosaics . We achieve state-of-the-art performance with our model , which shows that our method is suitable for representing permutations in general . The latter two tasks test whether a model can learn to solve a task that requires it to come up with a suitable permutation implicitly : classification from image mosaics and visual question answering . We provide no supervision of what the permutation should be ; the model has to learn by itself what permutation is most useful for the task at hand . Here , our model also beats the existing models and we improve the performance of a state-of-the-art model in VQA with it . This shows that our PO module is able to learn good permutation-invariant representations of sets using our approach . 2 PERMUTATION-OPTIMISATION MODULE . We will now describe a differentiable , and thus learnable model to turn an unordered set { xi } N with feature vectors as elements into an ordered sequence of these feature vectors . An overview of the algorithm is shown in Figure 1 and pseudo-code is available in Appendix A . The input set is represented as a matrix X = [ x1 , . . . , xN ] T with the feature vectors xi as rows in some arbitrary order . In the algorithm , it is important to not rely on the arbitrary order so that X is correctly treated as a set . The goal is then to learn a permutation matrix P such that when permuting the rows of the input through Y = PX , the output is ordered correctly according to the task at hand . When an entry Pik takes the value 1 , it can be understood as assigning the ith element to the kth position in the output . Our main idea is to first relate pairs of elements through an ordering cost , parametrised with a neural network . This pairwise cost tells us whether an element i should preferably be placed before or after element j in the output sequence . Using this , we can define a total cost that measures how good a given permutation is ( subsection 2.1 ) . The second idea is to optimise this total cost in each forward pass of the module ( subsection 2.2 ) . By minimising the total cost of a permutation , we improve the quality of a permutation with respect to the current ordering costs . Crucially , the ordering cost function – and thus also the total cost function – is learned . In doing so , the module is able to learn how to generate a permutation as is desired . In order for this to work , it is important that the optimisation process itself is differentiable so that the ordering cost is learnable . Because permutations are inherently discrete objects , a continuous relaxation of permutations is necessary . For optimisation , we perform gradient descent on the total cost for a fixed number of steps and unroll the iteration , similar to how recurrent neural networks are unrolled to perform backpropagation-through-time . Because the inner gradient ( total cost differentiated with respect to permutation ) is itself differentiable with respect to the ordering cost , the whole model is kept differentiable and we can train it with a standard supervised learning loss . Note that as long as the ordering cost is computed appropriately ( subsection 2.3 ) , all operations used turn out to be permutation-invariant . Thus , we have a model that respects the symmetries of sets while producing an output without those symmetries : a sequence . This can be naturally extended to outputs where the target is not a sequence , but grids and lattices ( subsection 2.4 ) . 2.1 TOTAL COST FUNCTION . The total cost function measures the quality of a given permutation and should be lower for better permutations . Because this is the function that will be optimised , it is important to understand what it expresses precisely . The main ingredient for the total cost of a permutation is the pairwise ordering cost ( details in subsection 2.3 ) . By computing it for all pairs , we obtain a cost matrix C where the entry Cij represents the ordering cost between i and j : the cost of placing element i anywhere before j in the output sequence . An important constraint that we put on C is that Cij = −Cji . In other words , if one ordering of i and j is “ good ” ( negative cost ) , then the opposite ordering obtained by swapping them is “ bad ” ( positive cost ) . Additionally , this constraint means that Cii = 0 . This makes sure that two very similar feature vectors in the input will be similarly ordered in the output because their pairwise cost goes to 0 . In this paper we use a straightforward definition of the total cost function : a sum of the ordering costs over all pairs of elements i and j . When considering the pair i and j , if the permutation maps i to be before j in the output sequence , this cost is simply Cij . Vice versa , if the permutation maps i to be after j in the output sequence , the cost has to be flipped to Cji . To express this idea , we define the total cost c : RN×N 7→ R of a permutation P as : c ( P ) = ∑ ij Cij ∑ k Pik ( ∑ k′ > k Pjk′ − ∑ k′ < k Pjk′ ) ( 1 ) This can be understood as follows : If the permutation assigns element i to position u ( so Piu = 1 ) and element j to position v ( so Pjv = 1 ) , the sums over k and k′ simplify to 1 when v > u and −1 when v < u ; permutation matrices are binary and only have one 1 in any row and column , so all other terms in the sums are 0 . That means that the term for each i and j becomes Cij when v > u and −Cij = Cji when v < u , which matches what we described previously . 2.2 OPTIMISATION PROBLEM . Now that we can compute the total cost of a permutation , we want to optimise this cost with respect to a permutation . After including the constraints to enforce that P is a valid permutation matrix , we obtain the following optimisation problem : minimize P c ( P ) subject to ∀i , k : Pik ∈ { 0 , 1 } , ∀i : ∑ k Pik = 1 , ∑ k Pki = 1 ( 2 ) Optimisation over P directly is difficult due to the discrete and combinatorial nature of permutations . To make optimisation feasible , a common relaxation is to replace the constraint that Pik ∈ { 0 , 1 } with Pik ∈ [ 0 , 1 ] ( Fogel et al. , 2013 ) . With this change , the feasible set for P expands to the set of doublystochastic matrices , known as the Birkhoff or assignment polytope . Rather than hard permutations , we now have soft assignments of elements to positions , analogous to the latent assignments when fitting a mixture of Gaussians model using Expectation-Maximisation . Note that we do not need to change our total cost function after this relaxation . Instead of discretely flipping the sign of Cij depending on whether element i comes before j or not , the sums over k and k′ give us a weight for each Cij that is based on how strongly i and j are assigned to positions . This weight is positive when i is on average assigned to earlier positions than j and negative vice versa . In order to perform optimisation of the cost under our constraints , we reparametrise P with the Sinkhorn operator S from Adams & Zemel ( 2011 ) ( defined in Appendix B ) so that the constraints are always satisfied . We found this to lead to better solutions than projected gradient descent in initial experiments . After first exponentiating all entries of a matrix , S repeatedly normalises all rows , then all columns of the matrix to sum to 1 , which converges to a doubly-stochastic matrix in the limit . P = S ( P̃ ) ( 3 ) This ensures that P is always approximately a doubly-stochastic matrix . P̃ can be thought of as the unnormalised permutation while P is the normalised permutation . By changing our optimisation to minimise P̃ instead of P directly , all constraints are always satisfied and we can simplify the optimisation problem to minP̃ c ( P ) without any constraints . It is now straightforward to optimise P̃ with standard gradient descent . First , we compute the gradient : ∂c ( P ) ∂Ppq = 2 ∑ j Cpj ( ∑ k′ > q Pjk′ − ∑ k′ < q Pjk′ ) ( 4 ) ∂c ( P ) ∂P̃pq = ∂P ∂P̃pq · ∂c ( P ) ∂P ( 5 ) From equation 4 , it becomes clear that this gradient is itself differentiable with respect to the ordering cost Cij , which allows it to be learned . In practice , both ∂c ( P ) /∂P̃ as well as ∂ [ ∂c ( P ) /∂P̃ ] /∂C can be computed with automatic differentiation . However , some implementations of automatic differentiation require the computation of c ( P ) which we do not use . In this case , implementing ∂c ( P ) /∂P̃ explicitly can be more efficient . Also notice that if we define Bjq = ∑ k′ > q Pjk′ −∑ k′ < q Pjk′ , equation 4 is just the matrix multiplication CB and is thus efficiently computable . For optimisation , P has to be initialised in a permutation-invariant way to preserve permutationinvariance of the algorithm . In this paper , we consider a uniform initialisation so that all Pik = 1/N ( PO-U model , left ) and an initialisation that linearly assigns ( Mena et al. , 2018 ) each element to each position ( PO-LA model , right ) . P̃ ( 0 ) ik = 0 or P̃ ( 0 ) ik = wkxi ( 6 ) where wk is a different weight vector for each position k. Then , we perform gradient descent for a fixed number of steps T . The iterative update using the gradient and a ( learnable ) step size η converges to the optimised permutation P ( T ) : P̃ ( t+1 ) = P̃ ( t ) − η ∂c ( P ( t ) ) ∂P ( t ) ( 7 ) One peculiarity of this is that we update P̃ with the gradient of the normalised permutation P , not of the unnormalised permutation P̃ as normal . In other words , we do gradient descent on P̃ but in equation 5 we set ∂Puv/∂P̃pq = 1 when u = p , v = q , and 0 everywhere else . We found that this results in significantly better permutations experimentally ; we believe that this is because ∂P /∂P̃ vanishes too quickly from the Sinkhorn normalisation , which biases P away from good permutation matrices wherein all entries are close to 0 and 1 ( Appendix D ) . The runtime of this algorithm is dominated by the computation of gradients of c ( P ) , which involves a matrix multiplication of two N ×N matrices . In total , the time complexity of this algorithm is T times the complexity of this matrix multiplication , which is Θ ( N3 ) in practice . We found that typically , small values for T such as 4 are enough to get good permutations .

Summary: This paper proposed an interesting idea of learning representations of sets by permutation optimizations. Through learning a permutation of the elements of a set, the proposed algorithm can learn a permutation-invariant representation of that set. To deal with the underlying difficult combinatorial optimization problem, the authors proposed to relax the optimization constraints and instead optimize over the set of doubly-stochastic matrices with reparameterization using the Sinkhorn operator. The cost function of this optimization is related to a pairwise ordering cost, which compares the order for each pair of the elements.