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 presents Generalized Norm-based Pooling (GNP), a L^p norm like pooling function with the trainable p for Graph Neural Network to achieve extrapolation. The authors also propose a training method that use a single linear layer to address the gradient exploding issue. Equally splitting the GNP function, the model can learn both positive and negative p. Experiments were done to demonstrate GNP at the node-/graph-/set- levels.