Learning to Pool in Graph Neural Networks for Extrapolation

1 INTRODUCTION . Many real-world data , such as relationships between people in social networks or chemical bonds between atoms , can naturally be represented as graphs . Finding models with proper inductive biases to better describe such graph data has been a common goal for many researchers , and Graph Neural Networks ( GNNs ) ( Scarselli et al. , 2009 ; Kipf & Welling , 2017 ; Hamilton et al. , 2017 ; Veličković et al. , 2018 ; Xu et al. , 2019 ; Maron et al. , 2019 ; Xu et al. , 2020 ) are considered to be the most successful model . They have proved effective for a variety of tasks , including recommendation ( Ying et al. , 2018a ) , drug discovery ( Stokes et al. , 2020 ) , and chip design ( Mirhoseini et al. , 2020 ) . An important design choice for a GNN often overlooked is the specification of pooling functions , the functions used for the aggregation or readout operation in GNNs . They are usually required to be invariant w.r.t . the permutation of nodes in a graph , and common choices are element-wise summation ( sum ) , maximum ( max ) , minimum ( min ) , or average ( mean ) . Some recent works also proposed to use parametric models and learn them from data as well ( Ying et al. , 2018b ; Lee et al. , 2019b ; Gao & Ji , 2019 ; Yuan & Ji , 2020 ) . While most of the previous works on this line focused on improving predictive performance for their own tasks , recently , Xu et al . ( 2021 ) studied the impact of the choice of pooling functions on the ability of a neural network to extrapolate . Specifically , Xu et al . ( 2021 ) highlighted the importance of the choice of pooling functions in order to make GNNs generalize over the data lying outside of the support of the training data distribution , and they argued that the specification of the pooling functions acts as an important inductive bias that can make GNNs either completely fail to extrapolate or gracefully generalize to out-of-distribution data . As a motivating example , consider the problem of counting the number of nodes in a graph . If we are to solve this problem with a single-layer GNN having one readout layer , probably the best pooling function would be sum , and the corresponding model will readily generalize to graphs with a much larger number of nodes than the ones seen during training . On the other hand , if we choose the pooling function as max instead , it may still fit the training data well but completely fail to predict the number of nodes in out-of-distribution graphs . The findings in Xu et al . ( 2021 ) raise a natural question ; which pooling functions should be used for a given problem in order to make GNNs constructed with them successfully extrapolate for out-of-distribution data ? Xu et al . ( 2021 ) did not present any guide but empirically showed that we do have the “ right ” pooling function for each problem tested , and when a pooling function is not properly selected , GNNs completely fails to extrapolate . The caveat here is that we do not know which pooling function is the right choice before actually training and validating the model . To this end , in this paper , we present a generic learning-based method to find proper pooling functions for a given arbitrary problem . Our method , entitled Generalized Norm-based Pooling ( GNP ) , formulates the pooling functions as a generic Lp norm-like function ( including negative p as well ) , and learns the parameters inside the pooling functions in an end-to-end fashion . Unlike previous learning-based pooling methods that are usually tailored for specific tasks or focused on improving predictive performances , GNP can be applied to arbitrary tasks , and it improves the extrapolation ability of GNNs constructed with it . Also , GNP includes most of the pooling functions being used for GNNs as special cases . Despite the enhanced flexibility , GNP incurs minimal overhead in GNN in terms of the model complexity . A naı̈ve application of GNP to GNNs is likely to fail because of some difficulty in training , so we propose a simple remedy to this . Using nine graph-level , node-level , and set-related tasks , we demonstrate that GNNs with GNP trained by our training scheme extrapolate for out-of-distribution data comparably and sometimes even better than those with pooling functions that are carefully chosen among of widely-used ones . In addition , we demonstrate the effectiveness of GNP on four real-world tasks ( graph classification , influence maximization , node classification , and graph regression ) using six additional GNN architectures ( GCN ( Kipf & Welling , 2017 ) , GAT ( Veličković et al. , 2018 ) , PNA ( Corso et al. , 2020 ) , hierarchical SAGPool ( Lee et al. , 2019b ) , ASAPool ( Ranjan et al. , 2020 ) , and MONSTOR ( Ko et al. , 2020 ) ) and eight real-world graphs . We summarize our contributions as follows : • Generalized pooling function : We propose GNP , a simple yet flexible pooling function that can readily be applied to arbitrary tasks involving GNNs , with minimal parameter overhead . • Effective training methods : We propose effective training methods for GNP . • Extensive experiments : We empirically demonstrate that GNNs with GNP generalize to out-of- distribution data on nine extrapolation tasks . We also show successful application of GNP to six GNN architectures on four real-world tasks in eight real-world graphs with up to one million edges . 2 RELATED WORK . Aggregation functions Various aggregation functions have been appeared to enhance the performance of GNNs . Hamilton et al . ( 2017 ) proposed GraphSAGE with four different aggregation methods ; max , mean , GCN ( Kipf & Welling , 2017 ) , and LSTM ( Hochreiter & Schmidhuber , 1997 ) . Veličković et al . ( 2018 ) proposed Graph Attention neTworks ( GATs ) including attention-based aggregation functions ( Vaswani et al. , 2017 ) . Xu et al . ( 2019 ) proposed Graph Isomorphism Networks ( GINs ) and proved that GNN can satisfy the 1-Weisfeiler-Lehman ( WL ) condition only with sum pooling function as aggregation function . Recently , Li et al . ( 2020 ) proposed a trainable softmax and power-mean aggregation function that generalizes basic operators . Compared to these methods designed to improve interpolation performance on specific tasks , ours can improve extrapolation performance for generic tasks . Readout functions Zhang et al . ( 2018 ) suggested SortPooling that chooses top-k values from the sorted list of the node features to construct outputs . Another popular idea is hierarchical pooling , where outputs are obtained by iteratively coarsening nodes in graphs in a hierarchical fashion ( Ying et al. , 2018b ; Gao & Ji , 2019 ; Lee et al. , 2019b ; Yuan & Ji , 2020 ) . Although demonstrated to be effective for the tasks they have been designed for , most of these methods require heavy computation and it is not straightforward to extend them for aggregation functions . On the other hand , our GNP can be applied to both aggregation and readout functions with minimal overhead . Pooling functions in generic context Vinyals et al . ( 2015 ) proposed Set2Set to get a representation of set-structured data with a LSTM-based pooling function . Lee et al . ( 2019a ) proposed to use an attention-based pooling function to get summaries of set data . For convolutional neural networks , there were some approaches to generalize average pooling and max pooling widely used for many neural network architectures . Gulcehre et al . ( 2014 ) proposed a normalized learnable Lp norm function that generalizes average pooling and max pooling . Lee et al . ( 2016 ) further extended those pooling functions with learnable tree-structured pooling filters . Norm-based pooling functions There have been several works to employ norm-based pooling functions . Gulcehre et al . ( 2014 ) proposed a learnable Lp norm function of the form f ( v ) = ( 1 |v| |v|∑ i=1 |vi|p ) 1/p to substitute max pooling or average pooling used in convolutional neural networks . Similar normbased pooling functions were used for acoustic modeling ( Swietojanski & Renals , 2016 ) and text representation ( Wu et al. , 2020 ) . Compared to GNP , these pooling methods can not express the sum pooling . Li et al . ( 2020 ) further generalized this by multiplying |v|q to include sum pooling as well , but not considered the case where p is positive and the case where p is negative at the same time . GNP is the most generic norm-based pooling function , compared to all aforementioned approaches , and more importantly , no other works studied their usefulness in the context of learning to extrapolate . Extrapolation Trask et al . ( 2018 ) pointed out that most of the feed-forward neural networks fail to extrapolate even for the simplest possible identity mapping , and suggested using alternative computation units mimicking the behavior of arithmetic logic units . The ability to extrapolate is also important in the GNN context , for instance , many combinatorial optimization problems involving graphs often require extrapolation . Selsam et al . ( 2019 ) ; Prates et al . ( 2019 ) tackled the extrapolation problem by performing large iterations of message passing . Using various classical graph algorithms , Veličković et al . ( 2020 ) showed that the extrapolation performance of GNNs depends heavily on the choice of the aggregation function . Similarly , Xu et al . ( 2021 ) demonstrated that choosing the right non-linear function for both MLPs and GNNs is crucial for the extrapolation . 3 MAIN CONTRIBUTION : GENERALIZED NORM-BASED POOLING . In this section , we present our Generalized Norm-based Pooling ( GNP ) and discuss its expressiveness . Then , we describe some difficulties in training GNP and our remedy . Lastly , we present a task on which a GNN with GNP can extrapolate , while that equipped with the basic pooling functions can not . 3.1 GENERALIZATION OF BASIC POOLING FUNCTIONS . While GNP is motivated by the Lp-norm function , which includes the sum and max functions as special cases , further ingredients are added to make GNP more flexible than the Lp-norm function . Specifically , we allow p to be negative to let GNP express a wider class of functions than the previous norm-based or learning-based pooling functions . Let V = { vi } ni=1 be a set of node features with vi ∈ Rd for i = 1 , . . . , n. We define GNP to be an element-wise function where the output for each jth element is GNPj ( V ) = 1 nq ( n∑ i=1 |vi , j |p ) 1/p , where p ∈ R \ { 0 } and q ∈ R are learnable parameters . GNP includes the basic pooling functions ( sum , mean , max , and min ) as special cases . Proposition 1 . Suppose all the entries of v are non-negative in equation 3.1 . Then , GNP includes sum , max , mean as special cases . If we further restrict v to be positive , GNP includes min . Proof . GNPj ( V ) is equivalent to elementwise sum when ( p , q ) = ( 1 , 0 ) and elementwise mean when ( p , q ) = ( 1 , 1 ) . When q = 0 , we have lim p→∞ GNPj ( V ) = max i vi lim p→∞ ( n∑ i=1 ( |vi , j | maxi |vi , j | ) p ) 1/p = max i |vi , j | · 1 = max i |vi , j | , ( 1 ) so GNP converges to max . Similarly , we can obtain min as a limit for p→ −∞ . 3.2 HANDLING OF NEGATIVE p The GNP function in equation 3.1 is not continuous and even not defined at p = 0 . Hence , directly learning GNP in the original form as in equation 3.1 ( even with p = 0 ignored ) can cause instability , especially when an algorithm is trying to move from a positive p value to a negative p value . Instead , we suggest splitting the GNP function into two parts , GNP+ with positive p and GNP− with negative p , and let the model choose the right balance between them . Specifically , define GNP+j ( V ) = 1 nq+ ( n∑ i=1 |vi , j |p + ) 1/p+ , GNP−j ( V ) = 1 nq− ( n∑ i=1 |vi , j |−p − ) −1/p− , ( 2 ) where p+ > 0 , q+ , p− > 0 , and q− are learnable parameters . Given a set of node features V , we first split the feature dimension into two , and compute the output from GNP+ for the first half and from GNP− for the second half . Then we mix two outputs with a single linear layer to get the final output . y = [ GNP+1 ( V ) . . . GNP + bd/2c ( V ) GNP − bd/2c+1 ( V ) . . . GNP − d ( V ) ] , ( 3 ) GNP ( V ) = Wy + b , ( 4 ) where b·c is the floor function , W ∈ Rd×d and b ∈ Rd are learnable parameters . Note that widelyused GNN layers have a linear layer or MLP after message-passing between nodes . Instead of using an additional linear layer , GNP concatenates the outputs of GNP+ and GNP− and passes them to the linear layer or MLP . Therefore , we have only four extra parameters ( p+ , p− , q+ and q− ) for each GNN layer of typical GNN architectures . With this design , GNP can easily switch between positive p and negative p , choosing proper values according to tasks .

Summary: This paper focus on the extrapolation ability of graph neural networks and propose a new pooling function for graph-level readout based on vector norm. The proposed method can be applied to replace the commonly used pooling function like max/mean/sum in GNNs and is proved able for extrapolation in a simple example. Experiments on several synthetic tasks and a few real-world datasets are conducted to verify the effectiveness of the proposed method.