Cold Brew: Distilling Graph Node Representations with Incomplete or Missing Neighborhoods

Graph Neural Networks ( GNNs ) have achieved state of the art performance in node classification , regression , and recommendation tasks . GNNs work well when high-quality and rich connectivity structure is available . However , this requirement is not satisfied in many real world graphs where the node degrees have power-law distributions as many nodes have either fewer or noisy connections . The extreme case of this situation is a node may have no neighbors at all , called Strict Cold Start ( SCS ) scenario . This forces the prediction models to rely completely on the node ’ s input features . We propose Cold Brew to address the SCS and noisy neighbor setting compared to pointwise and other graph-based models via a distillation approach . We introduce feature-contribution ratio ( FCR ) , a metric to study the viability of using inductive GNNs to solve the SCS problem and to select the best architecture for SCS generalization . We experimentally show FCR disentangles the contributions of various components of graph datasets and demonstrate the superior performance of Cold Brew on several public benchmarks and proprietary e-commerce datasets . Our codes will be fully open-sourced upon acceptance . 1 INTRODUCTION . Graph Neural Networks ( GNNs ) achieve state-of-the-art results across a wide range of tasks such as graph classification , node-classification , -regression , link prediction , and recommendation ( Wu et al. , 2020 ; Goyal & Ferrara , 2018 ; Kherad & Bidgoly , 2020 ; Shaikh et al. , 2017 ; Silva et al. , 2010 ; Zhang et al. , 2019 ) . GNNs rely on the principle of message passing to aggregate node features across multi-hop neighborhoods in order to learn aggregated representations . The success of modern GNN models relies on the presence of dense connections and high-quality neighborhoods . State-of-the-art GNNs ( Kipf & Welling , 2016 ; Veličković et al. , 2017 ; Hamilton et al. , 2017 ; Xu et al. , 2018a ; Wu et al. , 2019 ; Klicpera et al. , 2018 ) benefit from efficient message-passing frameworks developed to best leverage ( deep ) neighborhood information to learn node representations . Other methods like the spectral diffusion ( Joachims , 2003 ; Chaudhuri et al. , 2012 ; Zhang & Rohe , 2018 ; Wang et al. , 2021 ) and label propagation ( Huang et al. , 2020a ; Zhu , 2005 ) also require a non-empty and high-quality neighborhood per node to perform message passing . Even the inductive GNNs ( e.g. , Hamilton et al . ( 2017 ) ) learn a function of the node feature and the node neighborhood , requiring the neighborhood to be present during the inference time . Most of the large-scale real-world graphs are power-law in nature , with a majority of the nodes having very few connections ( Hao et al. , 2021 ; Ding et al. , 2021 ; Lam et al. , 2008 ; Lu et al. , 2020 ) . In Figure 1 we conceptually show the long-tail distribution of such a graph and also empirically show this characteristics on several public datasets . In many information retrieval and recommender systems applications , there is another ubiquitous challenge : the Strict Cold Start ( SCS ) ( Li et al. , 2019b ; Ding et al. , 2021 ) , wherein some “ nodes ” will have no edges in the graph . For instance , in most of the public and proprietary datasets we used in the experimentation we have anywhere between 3 % to 6 % isolated nodes with no neighbors . In these cases , existing GNNs fail to perform due to sparsity or absence of neighborhood . In this paper , we develop GNN models that can achieve truly inductive capabilities : one can learn effective node embeddings for “ orphaned ” nodes in a graph . This capability is important to fully realize the potential of large-scale GNN models on modern , industrialsized datasets with very long tails and nodes with no neighbors in the graph . To this end , Normalized Number of Nodes 101 102 103 104 No de D eg re es cora citeseer pubmed arxiv chameleon actor squirrel wisconsin cornell texas tures , we develop a metric called Feature-Contribution Ratio ( FCR ) that disentangles the graph data into node features and the neighborhood structure . We then build submodules for each part and search for the best architecture . Besides selecting the optimal model architecture , we leverage FCR to provide a measure of importance of the node feature compared to the neighborhood structure , quantifying the difficulty of learning truly inductive GNNs . We summarize our key contributions as follows : • To generalize better to tail and SCS nodes , we design the Cold Brew framework to distill a GNN teacher into an MLP student . We improve over GNN models by adding the node-wise Structural Embedding ( SE ) to the Cold Brew ’ s teacher GNN to strengthen the expressiveness of the teacher GNN model . We designed a novel mechanism for the MLP student to discover the “ latent/virtual neighborhoods ” when they are missing , and perform message passing as its teacher model does . • We propose Feature-Contribution Ratio ( FCR ) : a new metric for graph datasets that quantifies the contribution of node features w.r.t . the adjacency structure in the dataset for a specific downstream task . FCR indicates the difficulty in generalizing to tail and cold-start nodes in a given dataset . We also leverage FCR as a principled “ screen process ” to select the best model architecture for both the GNN teacher and the MLP student in Cold Brew . • Through extensive experiments on several public datasets as well as a proprietary ecommerce graph dataset , we validate the effectiveness of Cold Brew in cold-start generalization . We further uncover the contribution ratio of node features in node prediction tasks with FCR . 1.1 PROBLEM SETUP . GNNs effectively learn the representations of two components in graph data : they process node features by distributed node-wise transformations and process adjacency structure by localized neighborhood aggregations . For the first component , GNNs apply shared feature transformations to all nodes regardless of the neighborhoods , while for the second component , GNNs use permutationinvariant aggregators to collect neighborhood information . We take the node classification problem in this paper for ease of exposition , and the proposed method can be trivially adapted to semi-supervised/unsupervised settings . We denote the graph data of interest by G with node set V , |V| = N . Each node possesses a din−dimensional feature and a dout−dimensional label ( either dout classes or a continuous vector in the case of regression ) . Let X0 ∈ RN×din and Y ∈ RN×dout be the matrices of node features and labels , respectively . Let Ni be the neighborhood of node i ∈ [ N ] . As is common with modern , large-scale graphs , |Ni| is small for several nodes in the graph , which we refer to as tail nodes . For a number of nodes , |Ni| = 0 , and we refer to these nodes as cold-start nodes . A traditional GNN learns representations for the ith node at the lth layer as a function of the node representation itself and its neighborhoods ’ representations , at the ( l − 1 ) th layer : xli : = f ( { xl−1i } , { x l−1 j } j∈Ni ) ( 1 ) where f ( · ) is a general function that applies node-wise transformation on node xl−1i and aggregates information of its neighborhood { xl−1j } j∈Ni to obtain the final node representation . Given i ’ s input features x0i and its neighborhood Ni , one can use ( 1 ) to obtain its representation and predict yi , making these models inductive . We are interested in improving the performance of these GNNs on a set of tail and cold-start nodes , whereNi for node i is either unreliable1 or absent . In these cases , applying ( 1 ) will yield a suboptimal node representation , since { xl−1j } j∈Ni will be unreliable or empty at inference time . 2 RELATED WORK . GNNs learn by aggregating neighborhood information to learn node representations ( Kipf & Welling , 2016 ; Veličković et al. , 2017 ; Hamilton et al. , 2017 ; Xu et al. , 2018a ; Wu et al. , 2019 ; Klicpera et al. , 2018 ) . Inductive variants of GNNs such as GraphSAGE ( Hamilton et al. , 2017 ) require initial node features as well as the neighborhood information of each node to learn the representation . Most works on improving GNNs have focused on learning better aggregation functions , and methods that can work when the neighborhood is absent or noisy have not been sufficiently exploited , except two very recent or concurrent works ( Hu et al. , 2021 ; Zhang et al. , 2021 ) . In the context of cold start , ( Hao et al. , 2021 ) and ( Ding et al. , 2021 ) employ a transfer learning approach . ( Yang et al. , 2021 ) proposes a knowledge distillation for GNN , while ( Chen et al. , 2020b ) proposes a self-distillation approach . In all the above cases , the models need full knowledge of the neighbors of the cold-start nodes in question and do not address the case of noisy or missing neighborhoods . Another possible solution is to directly train an MLP that only take node features . ( Hu et al. , 2021 ) proposes to learn graph embeddings with only node-wise MLP , while using contrastive loss to regularize the graph structure . Some previous works have studied the relation between node feature similarity and edge connections and studied their influence on the selection of appropriate graph models . ( Pei et al. , 2020 ) proposed the homophily metric that categorizes the graphs into assortative and disassortative classes . ( Wang et al. , 2021 ) dissected the feature propagation steps of linear GCNs from a perspective of continuous graph diffusion and analyzed why linear GCNs fail to benefit from more propagation steps . ( Liu et al. , 2020a ) further studied the influence of homophily on model selection and proposed a non-local GNN . 3 STRICT COLD START GENERALIZATION . We now address the problem of generalization to the tail and cold-start nodes , where the neighborhood information is missing/noisy ( Section 1 ) . A naive baseline is to train an MLP to map node features to labels . However , such method disregards all graph information , and we show via our FeatureContribution Ratio and other experimental results that for most assortative graph datasets , the node-wise MLP approach is suboptimal . The key idea of our framework is the following : the GNN maps node features into a d-dimensional embedding space , and since the number of nodes N is usually much bigger than the embedding dimensionality d , we end up with an overcomplete set for this space using the embeddings as the basis . This implies the possibility that any node representation can be cast as a linear combination of K N existing node representations . Our aim will be to train a student model that can accurately discover the combination of the best K existing node embeddings of a target isolated node . We call 1For example , a user with only one movie watched or an item with too few purchases . this procedure latent/virtual neighborhood discovery , which is equivalent to using MLPs to “ mimic ” the node representations learned by the teacher GNN . We adopt the knowledge distillation procedure ( Yang et al. , 2021 ; Chen et al. , 2020b ) to improve the quality of the learned embeddings for tail and cold-start nodes . We use a teacher GNN model to embed the nodes onto a low-dimensional manifold by utilizing the graph structure . Then , the job of the student is to learn a mapping from the node features to this manifold without knowledge of the graph that the teacher has . We further aim to let the student model generalize to SCS cases where the teacher model fails , beyond just mimicking the teacher as standard knowledge distillation does .

Summary: The paper proposes Cold Brew to distill the knowledge of a GNN teacher into an MLP student to handle the tail and cold start generalization problem by using the head part of the graph to guide the discovery of the latent neighborhoods of tail and isolation nodes. The paper also proposes a new metric to measure the contribution ratio of node features w.r.t. the adjacency structure. The experiments on several public datasets and a proprietary e-commerce graph show the effectiveness of the proposed method.