Janossy Pooling: Learning Deep Permutation-Invariant Functions for Variable-Size Inputs

1 INTRODUCTION . Pooling is a fundamental operation in deep learning architectures ( LeCun et al. , 2015 ) . The role of pooling is to merge a collection of related features into a single , possibly vector-valued , summary feature . A prototypical example is in convolutional neural networks ( CNNs ) ( LeCun et al. , 1995 ) , where linear activations of features in neighborhoods of image locations are pooled together to construct more abstract features . A more modern example is in neural networks for graphs , where each layer pools together embeddings of neighbors of a vertex to form a new embedding for that vertex , see for instance , ( Kipf & Welling , 2016 ; Atwood & Towsley , 2016 ; Hamilton et al. , 2017 ; Velickovic et al. , 2017 ; Monti et al. , 2017 ; Xu et al. , 2018 ; Liu et al. , 2018 ; Liben-Nowell & Kleinberg , 2007 ; van den Berg et al. , 2017 ; Duvenaud et al. , 2015 ; Gilmer et al. , 2017 ; Ying et al. , 2018 ; Xu et al. , 2019 ) . A common requirement of a pooling operator is invariance to the ordering of the input features . In CNNs for images , pooling allows invariance to translations and rotations , while for graphs , it allows invariance to graph isomorphisms . Existing pooling operators are mostly limited to predefined heuristics such as max-pool , min-pool , sum , or average . Another desirable characteristic of pooling layers is the ability to take variable-size inputs . This is less important in images , where neighborhoods are usually fixed a priori . However in applications involving graphs , the number of neighbors of different vertices can vary widely . Our goal is to design flexible and learnable pooling operators satisfying these two desiderata . Abstractly , we will view pooling as a permutation-invariant ( or symmetric ) function acting on finite but arbitrary length sequences h. All elements hi of the sequences are features lying in some space H ( which itself could be a high-dimensional Euclidean space Rd or some subset thereof ) . The sequences h are themselves elements of the union of products of the H-space : h ∈ ⋃∞ j=0 Hj ≡ H∪ . Throughout the paper , we will use Πn to represent the set of all permutations of the integers 1 to n , where n will often be clear from the context . In addition , hπ , π ∈ Π|h| , will represent a reordering of the elements of a sequence h according to π , where |h| is the length of the sequence h. We will use the double bar superscript f to indicate that a function is permutation-invariant , returning the same value no matter the order of its arguments : f ( h ) = f ( hπ ) , ∀π ∈ Π|h| . We will use the arrow superscript f ⇀ to indicate general functions on sequences h which may or may not be permutationinvariant1 . Functions f without any markers are ‘ simple ’ functions , acting on elements in H , scalars or any other argument that is not a sequence of elements in H. Our goal in this paper is to model and learn permutation-sensitive functions f ⇀ that can be used to construct flexible and learnable permutation-invariant neural networks . A recent step in this direction is work on DeepSets by Zaheer et al . ( 2017 ) , who argued for learning permutation-invariant functions through the following composition : y ( x ; θ ( ρ ) , θ ( f ) , θ ( h ) ) = ρ ( f ( |h| , h ; θ ( f ) ) ; θ ( ρ ) ) , where ( 1 ) f ( |h| , h ; θ ( f ) ) = |h|∑ j=1 f ( hj ; θ ( f ) ) and h ≡ h ( x ; θ ( h ) ) . ( 2 ) Here , ( a ) x ∈ X is one observation in the training data ( X itself may contain variable-length sequences ) , h ∈ H is the embedding ( output ) of the data given by the lower layers h : X×Ra → H∪ , a > 0 with parameters θ ( h ) ∈ Ra ; ( b ) f : H× Rb → F is a middle-layer embedding function with parameters θ ( f ) ∈ Rb , b > 0 , and F is the embedding space of f ; and ( c ) ρ : F×Rc → Y is a neural network with parameters θ ( ρ ) ∈ Rc , c > 0 , that maps to the final output space Y . Typically H and F are high-dimensional real-valued spaces ; Y is often Rd in d-dimensional regression problems or the simplex in classification problems . Effectively , the neural network f learns an embedding for each element in H , and given a sequence h , its component embeddings are added together before a second neural network transformation ρ is applied . Note that the function h may be the identity mapping h ( x ; · ) = x that makes f act directly on the input data . Zaheer et al . ( 2017 ) argue that if ρ is a universal function approximator , the above architecture is capable of approximating any symmetric function on h-sequences , which justifies the widespread use of average ( sum ) pooling to make neural networks permutation-invariant in Duvenaud et al . ( 2015 ) , Hamilton et al . ( 2017 ) , Kipf & Welling ( 2016 ) , Atwood & Towsley ( 2016 ) , among other works . We note that Zaheer et al . ( 2017 ) focus on functions of sets but the work was extended to functions of multisets by Xu et al . ( 2019 ) and that Janossy pooling can be used to represent multiset functions . Y Z tion hπ , ⇀ f ( |h| , hπ ; θ ( f ) ) is computed . These are summed and passed to a second function ρ ( · ; θ ( ρ ) ) which gives the final permutation-invariant output y ( x ; θ ( ρ ) , θ ( f ) , θ ( h ) ) ; the gray rectangle represents Janossy pooling . We discuss how this can be made computationally tractable . In practice , there is a gap between flexibility and learnability . While the architecture of equations 1 and 2 is a universal approximator to permutationinvariant functions , it does not easily encode structural knowledge about y . Consider trying to learn the permutation-invariant function y ( x ) = maxi , j≤|x| |xi − xj | . With higherorder interactions between the elements of h , the functions f of equation 2 can not capture any useful intermediate representations towards the final output , with the burden shifted entirely to the function ρ . Learning ρ means learning to undo mixing performed by the summation layer f ( |h| , h ; θ ( f ) ) =∑|h| j=1 f ( hj ; θ ( f ) ) . As we show in our experiments , in many applications this is too much to ask of ρ . 1LATEXcode for these markers is provided in the Supplementary Material . Contributions . We investigate a learnable permutation-invariant pooling layer for variable-size inputs inspired by the Janossy density framework , widely used in the theory of point processes ( Daley & Vere-Jones , 2003 , Chapter 7 ) . This approach , which we call Janossy pooling , directly allows the user to model what higher-order dependencies in h are relevant in the pooling . Figure 1 summarizes a neural network with a single Janossy pooling layer f ( detailed in Definition 2.1 below ) : given an input embedding h , we apply a learnable ( permutation-sensitive ) function f ⇀ to every permutation hπ of the input sequence h. These outputs are added together , and fed to the second function ρ . Examples of function f ⇀ include feedforward and recurrent neural networks ( RNNs ) . We call the operation used to construct f from f ⇀ the Janossy pooling . Definition 2.1 gives a more detailed description . We will detail three broad strategies for making this computation tractable and discuss how existing methods can be seen as tractability strategies under the Janossy pooling framework . Thus , we propose a framework and tractability strategies that unify and extend existing methods in the literature . We contribute the following analysis : ( a ) We show DeepSets ( Zaheer et al. , 2017 ) is a special case of Janossy pooling where the function f ⇀ depends only on the first element of the sequence hπ . In the most general form of Janossy pooling ( as described above ) , f ⇀ depends on its entire input sequence hπ . This naturally raises the possibility of intermediate choices of f ⇀ that allow practitioners to trade between flexibility and tractability . We will show that functions f ⇀ that depend on their first k arguments of hπ allow the Janossy pooling layer to capture up to k-ary dependencies in h. ( b ) We show Janossy pooling can be used to learn permutation-invariant neural networks y ( x ) by sampling a random permutation of h during training , and then modeling this permuted sequence using a sequence model such as a recurrent neural network ( LSTMs ( Hochreiter & Schmidhuber , 1997 ) , GRUs ( Cho et al. , 2014 ) ) or a vector model such as a feedforward network . We call this permutation-sampling learning algorithm π-SGD ( π-Stochastic Gradient Descent ) . Our analysis explains why this seemingly unsound procedure is theoretically justified , which sheds light on the recent puzzling success of permutation sampling and LSTMs in relational models ( Moore & Neville , 2017 ; Hamilton et al. , 2017 ) . We show that this property relates to randomized model ensemble techniques . ( c ) In Zaheer et al . ( 2017 ) , the authors describe a connection between DeepSets and infinite de Finetti exchangeabilty . We provide a probabilistic connection between Janossy pooling and finite de Finetti exchangeabilty ( Diaconis , 1977 ) . 2 JANOSSY POOLING . We first formalize the Janossy pooling function f . Start with a function f ⇀ , parameterized by θ ( f ) , which can take any variable-size sequence as input : a sequence of matrices ( such as images ) , a sequence of vectors ( such as a sequence of vector embeddings ) , or a variable-size sequence of features or embeddings representing the neighbors of a node in an attributed graph . In practice , we implement f ⇀ with a neural network . Formalizing Figure 1 from Section 1 , we use f ⇀ to define f : Definition 2.1 : [ Janossy pooling ] Consider a function f ⇀ : N × H∪ × Rb → F on variable-length but finite sequences h , parameterized by θ ( f ) ∈ Rb , b > 0 . A permutation-invariant function f : N×H∪ × Rb → F is the Janossy function associated with f ⇀ if f ( |h| , h ; θ ( f ) ) = 1 |h| ! ∑ π∈Π|h| f ⇀ ( |h| , hπ ; θ ( f ) ) , ( 3 ) where Π|h| is the set of all permutations of the integers 1 to |h| , and hπ represents a particular reordering of the elements of sequence h according to π ∈ Π|h| . We refer the operation used to construct f from f ⇀ as Janossy pooling . ♦ Definition 2.1 provides a conceptually simple approach for constructing permutation-invariant functions from arbitrary and powerful permutation-sensitive functions such as feedforward networks , recurrent neural networks , or convolutional neural networks . If f ⇀ is a vector-valued function , then so is f , and in practice , one might pass this vector output of f through a second function ρ ( e.g . a neural network parameterized by θ ( ρ ) ) : y ( x ; θ ( ρ ) , θ ( f ) , θ ( h ) ) = ρ 1 |h| ! ∑ π∈Π|h| f ⇀ ( |h| , hπ ; θ ( f ) ) ; θ ( ρ ) , where h ≡ h ( x ; θ ( h ) ) . ( 4 ) Equation 3 can capture any permutation-invariant function g for a flexible enough family of permutation-sensitive functions f ⇀ ( for instance , one could always set f ⇀ = g ) . Thus , at least theoretically , ρ in equation 4 provides no additional representational power . In practice , however , ρ can improve learnability by capturing common aspects across all terms in the summation . Furthermore , when we look at approximations to equation 3 or restrictions of f ⇀ to more tractable families , adding ρ can help recover some of the lost model capacity . Overall then , equation 4 represents one layer of Janossy pooling , forming a constituent part of a bigger neural network . Figure 1 summarizes this . Janossy pooling , as defined in equation 3 and 4 is intractable ; the computational cost of summing over all permutations ( for prediction ) , and backpropagating gradients ( for learning ) is likely prohibitive for most problems of interest . Nevertheless , it provides an overarching framework to unify existing methods , and to extend them . In what follows we present strategies for mitigating this , allowing novel and effective trade-offs between learnability and computational cost .

Summary: I really enjoyed this paper. It takes an idea which at first glance seems to be obviously bad (if you want permutation invariance, build a model that considers all permutations) and uses it to make the important point that the universal approximation results contained in Deep Sets [Zaheer et al. 2017] are not the last word on pooling. Janossy Pooling is intractable for most problems of interest (because it sums over all n! permutations of the input set) so the authors suggest 3 tractable alternatives: canonical orderings, k-ary dependencies and SGD / sampling-based approaches. Only the latter two are explored in detail, so I’ll focus on them: