FSPool: Learning Set Representations with Featurewise Sort Pooling

1 INTRODUCTION . Consider the following task : you have a dataset wherein each datapoint is a set of 2-d points that form the vertices of a regular polygon , and the goal is to learn an auto-encoder on this dataset . The only variable is the rotation of this polygon around the origin , with the number of points , size , and centre of it fixed . Because the inputs and outputs are sets , this problem has some unique challenges . Encoder : This turns the set of points into a latent space . The order of the elements in the set is irrelevant , so the feature vector the encoder produces should be invariant to permutations of the elements in the set . While there has been recent progress on learning such functions ( Zaheer et al. , 2017 ; Qi et al. , 2017 ) , they compress a set of any size down to a single feature vector in one step . This can be a significant bottleneck in what these functions can represent efficiently , particularly when relations between elements of the set need to be modeled ( Murphy et al. , 2019 ; Zhang et al. , 2019b ) . Decoder : This turns the latent space back into a set . The elements in the target set have an arbitrary order , so a standard reconstruction loss can not be used naïvely – the decoder would have to somehow output the elements in the same arbitrary order . Methods like those in Achlioptas et al . ( 2018 ) therefore use an assignment mechanism to match up elements ( section 2 ) , after which a usual reconstruction loss can be computed . Surprisingly , their model is still unable to solve the polygon reconstruction task with close-to-zero reconstruction error , despite the apparent simplicity of the dataset . In this paper , we introduce a set pooling method for neural networks that addresses both the encoding bottleneck issue and the decoding failure issue . We make the following contributions : 1 . We identify the responsibility problem ( section 3 ) . This is a fundamental issue with existing set prediction models that has not been considered in the literature before , explaining why these models struggle to model even the simple polygon dataset . 2 . We introduce FSPOOL : a differentiable , sorting-based pooling method for variable-size sets ( section 4 ) . By using our pooling in the encoder of a set auto-encoder and inverting the sorting in the decoder , we can train it with the usual MSE loss for reconstruction without the need for an assignment-based loss . This avoids the responsibility problem . 3 . We show that our auto-encoder can learn polygon reconstructions with close-to-zero error , which is not possible with existing set auto-encoders ( subsection 6.1 ) . This benefit transfers over to a set version of MNIST , where the quality of reconstruction and learned representation is improved ( subsection 6.2 ) . In further classification experiments on CLEVR ( subsection 6.3 ) and several graph classification datasets ( subsection 6.4 ) , using FSPool in a set encoder improves over many non-trivial baselines . Lastly , we show that combining FSPool with Relation Networks significantly improves over standard Relation Networks in a model that heavily relies on the quality of the representation ( subsection 6.5 ) . 2 BACKGROUND . The problem with predicting sets is that the output order of the elements is arbitrary , so computing an elementwise mean squared error does not make sense ; there is no guarantee that the elements in the target set happen to be in the same order as they were generated . The existing solution around this problem is an assignment-based loss , which assigns each predicted element to its “ closest ” neighbour in the target set first , after which a traditional pairwise loss can be computed . We have a predicted set Ŷ with feature vectors as elements and a ground-truth set Y , and we want to measure how different the two sets are . These sets can be represented as matrices with the feature vectors placed in the columns in some arbitrary order , so Ŷ = [ ŷ ( 1 ) , . . . , ŷ ( n ) ] and Y = [ y ( 1 ) , . . . , y ( n ) ] with n as the set size ( columns ) and d as the number of features per element ( rows ) . In this work , we assume that these two sets have the same size . The usual way to produce Ŷ is with a multi-layer perceptron ( MLP ) that has d× n outputs . Linear assignment One way to do this assignment is to find a linear assignment that minimises the total loss , which can be solved with the Hungarian algorithm in O ( n3 ) time . With Π as the space of all n-length permutations : LH ( Ŷ , Y ) = min π∈Π n∑ i ||ŷ ( i ) − y ( π ( i ) ) ||2 ( 1 ) Chamfer loss Alternatively , we can assign each element directly to the closest element in the target set . To ensure that all points in the target set are covered , a term is added to the loss wherein each element in the target set is also assigned to the closest element in the predicted set . This has O ( n2 ) time complexity and can be run efficiently on GPUs . LC ( Ŷ , Y ) = ∑ i min j ||ŷ ( i ) − y ( j ) ||2 + ∑ j min i ||ŷ ( i ) − y ( j ) ||2 ( 2 ) Both of these losses are examples of permutation-invariant functions : the loss is the same regardless of how the columns of Y and Ŷ are permuted . 3 RESPONSIBILITY PROBLEM . It turns out that standard neural networks struggle with modeling symmetries that arise because there are n ! different list representations of the same set , which we highlight here with an example . Suppose we want to train an auto-encoder on our polygon dataset and have a square ( so a set of 4 points with the x-y coordinates as features ) with some arbitrary initial rotation ( see Figure 1 ) . Each pair in the 8 outputs of the MLP decoder is responsible for producing one of the points in this square . We mark each such pair with a different colour in the figure . If we rotate the square ( top left in figure ) by 90 degrees ( top right in figure ) , we simply permute the elements within the set . They are the same set , so they also encode to the same latent representation and decode to the same list representation . This means that each output is still responsible for producing the point at the same position after the rotation , i.e . the dark red output is still responsible for the top left point , the light red output is responsible for the top right point , etc . However , this also means that at some point during that 90 degree rotation ( bottom path in figure ) , there must exist a discontinuous jump ( red arrow in figure ) in how the outputs are assigned . We know that the 90 degree rotation must start and end with the top left point being produced by the dark red output . Thus , we know that there is a rotation where all the outputs must simultaneously change which point they are responsible for , so that completing the rotation results in the top left point being produced by the dark red output . Even though we change the set continuously , the list representation ( MLP or RNN outputs ) must change discontinuously . This is a challenge for neural networks to learn , since they can typically only model functions without discontinuous jumps . As we increase the number of vertices in the polygon ( number of set elements ) , it must learn an increasing frequency of situations where all the outputs must discontinuously change at once , which becomes very difficult to model . Our experiment in subsection 6.1 confirms this . This example highlights a more general issue : whenever there are at least two set elements that can be smoothly interchanged , these discontinuities arise . We show this more formally in Appendix A . For example , the set of bounding boxes in object detection can be interchanged in much the same way as the points of our square here . An MLP or RNN that tries to generate these ( like in Rezatofighi et al . ( 2018 ) ; Stewart & Andriluka ( 2016 ) ) must handle which of its outputs is responsible for what element in a discontinuous way . Note that traditional object detectors like Faster R-CNN do not have this responsibility problem , because they do not treat object detection as a proper set prediction task with their anchor-based approach . When the set elements come from a finite domain ( often a set of labels ) and not Rd , it does not make sense to interpolate set elements . Thus , the responsibility problem does not apply to methods for such problems , for example Welleck et al . ( 2018 ) ; Rezatofighi . et al . ( 2017 ) . 4 FEATUREWISE SORT POOLING . The main idea behind our pooling method is simple : sorting each feature across the elements of the set and performing a weighted sum . The numerical sorting ensures the property of permutationinvariance . The difficulty lies in how to determine the weights for the weighted sum in a way that works for variable-sized sets . A key insight for auto-encoding is that we can store the permutation that the sorting applies in the encoder and apply the inverse of that permutation in the decoder . This allows the model to restore the arbitrary order of the set element so that it no longer needs an assignment-based loss for training . This avoids the problem in Figure 1 , because rotating the square by 90◦ also permutes the outputs of the network accordingly . Thus , there is no longer a discontinuity in the outputs during this rotation . In other words , we make the auto-encoder permutation-equivariant : permuting the input set also permutes the neural network ’ s output in the same way . We describe the model for the simplest case of encoding fixed-size sets in subsection 4.1 , extend it to variable-sized sets in subsection 4.2 , then discuss how to use this in an auto-encoder in subsection 4.3 . 4.1 FIXED-SIZE SETS . We are given a set of n feature vectors X = [ x ( 1 ) , . . . , x ( n ) ] where each x ( i ) is a column vector of dimension d placed in some arbitrary order in the columns of X ∈ Rd×n . From this , the goal is to produce a single feature vector in a way that is invariant to permutation of the columns in the matrix . We first sort each of the d features across the elements of the set by numerically sorting within the rows of X to obtain the matrix of sorted features ~X : ~Xi , j = SORT ( Xi , : ) j ( 3 ) where Xi , : is the ith row of X and SORT ( · ) sorts a vector in descending order . While this may appear strange since the columns of ~X no longer correspond to individual elements of the set , there are good reasons for this . A transformation ( such as with an MLP ) prior to the pooling can ensure that the features being sorted are mostly independent so that little information is lost by treating the features independently . Also , if we were to sort whole elements by one feature , there would be discontinuities whenever two elements swap order . This problem is avoided by our featurewise sorting . Efficient parallel implementations of SORT are available in Deep Learning frameworks such as PyTorch , which uses a bitonic sort ( O ( log2 n ) parallel time , O ( n log2 n ) comparisons ) . While the permutation that the sorting applies is not differentiable , gradients can still be propagated pathwise according to this permutation in a similar way as for max pooling . Then , we apply a learnable weight matrix W ∈ Rd×n to ~X by elementwise multiplying and summing over the columns ( row-wise dot products ) . yi = n∑ j Wi , j ~Xi , j ( 4 ) y ∈ Rd is the final pooled representation of ~X . The weight vector allows different weightings of different ranks and is similar in spirit to the parametric version of the gather step in Gather-Excite ( Hu et al. , 2018 ) . This is a generalisation of both max and sum pooling , since max pooling can be obtained with the weight vector [ 1 , 0 , . . . , 0 ] and sum pooling can be obtained with the 1 vector . Thus , it is also a maximally powerful pooling method for multi-sets ( Xu et al. , 2019 ) while being potentially more flexible ( Murphy et al. , 2019 ) in what it can represent .

Summary: This paper presents a new approach to representing sets of inputs and a strategy for set auto-encoders. Multiple experiments demonstrate that the new approach performs better than baseline models from recent literature around set representations. The forward encoding strategy is simple enough for a practitioner to implement, and likely to perform better (in terms of training time/gradient flow; if not also test metrics) than existing sum, min, max, mean pooling strategies. It is not substantially more expensive.