Graph Traversal with Tensor Functionals: A Meta-Algorithm for Scalable Learning

1 INTRODUCTION . Graph representation learning ( GRL ) has become an invaluable approach for a variety of tasks , such as node classification ( e.g. , in biological and citation networks ; Veličković et al . ( 2018 ) ; Kipf & Welling ( 2017 ) ; Hamilton et al . ( 2017 ) ; Xu et al . ( 2018 ) ) , edge classification ( e.g. , link prediction for social and protein networks ; Perozzi et al . ( 2014 ) ; Grover & Leskovec ( 2016 ) ) , entire graph classification ( e.g. , for chemistry and drug discovery Gilmer et al . ( 2017 ) ; Chen et al . ( 2018a ) ) , etc . In this work , we propose an algorithmic unification of various GRL methods that allows us to re-implement existing GRL methods and introduce new ones , in merely a handful of code lines per method . Our algorithm ( abbreviated GTTF , Section 3.2 ) , receives graphs as input , traverses them using efficient tensor1 operations , and invokes specializable functions during the traversal . We show function specializations for recovering popular GRL methods ( Section 3.3 ) . Moreover , since GTTF is stochastic , these specializations automatically scale to arbitrarily large graphs , without careful derivation per method . Importantly , such specializations , in expectation , recover unbiased gradient estimates of the objective w.r.t . model parameters . 1To disambiguate : by tensors , we refer to multi-dimensional arrays , as used in Deep Learning literature ; and by operations , we refer to routines such as matrix multiplication , advanced indexing , etc GTTF uses a data structure Â ( Compact Adjacency , Section 3.1 ) : a sparse encoding of the adjacency matrix . Node v contains its neighbors in row Â [ v ] , Âv , notably , in the first degree ( v ) columns of Â [ v ] . This encoding allows stochastic graph traversals using standard tensor operations . GTTF is a functional , as it accepts functions ACCUMULATEFN and BIASFN , respectively , to be provided by each GRL specialization to accumulate necessary information for computing the objective , and optionally to parametrize sampling procedure p ( v ’ s neighbors | v ) . The traversal internally constructs a walk forest as part of the computation graph . Figure 1 depicts the data structure and the computation . From a generalization perspective , GTTF shares similarities with Dropout ( Srivastava et al. , 2014 ) . Our contributions are : ( i ) A stochastic graph traversal algorithm ( GTTF ) based on tensor operations that inherits the benefits of vectorized computation and libraries such as PyTorch and Tensorflow . ( ii ) We list specialization functions , allowing GTTF to approximately recover the learning of a broad class of popular GRL methods . ( iii ) We prove that this learning is unbiased , with controllable variance . Wor this class of methods , ( iv ) we show that GTTF can scale previously-unscalable GRL algorithms , setting the state-of-the-art on a range of datasets . Finally , ( v ) we open-source GTTF along with new stochastic traversal versions of several algorithms , to aid practitioners from various fields in applying and designing state-of-the-art GRL methods for large graphs . 2 RELATED WORK . We take a broad standpoint in summarizing related work to motivate our contribution . Method Fa m ily Sc al e L ea rn in g Models GCN , GAT MP 7 exact node2vec NE 3 approx WYS NE 7 exact Stochastic Sampling Methods SAGE MP 3 approx FastGCN MP 3 approx LADIES MP 3 approx GraphSAINT MP 3 approx CluterGCN MP 3 heuristic Software Frameworks PyG Both inherits / reDGL Both implements Algorithmic Abstraction ( ours ) GTTF Both 3 approx Models for GRL have been proposed , including message passing ( MP ) algorithms , such as Graph Convolutional Network ( GCN ) ( Kipf & Welling , 2017 ) , Graph Attention ( GAT ) ( Veličković et al. , 2018 ) ; as well as node embedding ( NE ) algorithms , including node2vec ( Grover & Leskovec , 2016 ) , WYS ( Abu-El-Haija et al. , 2018 ) ; among many others ( Xu et al. , 2018 ; Wu et al. , 2019 ; Perozzi et al. , 2014 ) . The full-batch GCN of Kipf & Welling ( 2017 ) , which drew recent attention and has motivated many MP algorithms , was not initially scalable to large graphs , as it processes all graph nodes at every training step . To scale MP methods to large graphs , researchers proposed Stochastic Sampling Methods that , at each training step , assemble a batch constituting subgraph ( s ) of the ( large ) input graph . Some of these sampling methods yield unbiased gradient estimates ( with some variance ) including SAGE ( Hamilton et al. , 2017 ) , FastGCN ( Chen et al. , 2018b ) , LADIES ( Zou et al. , 2019 ) , and GraphSAINT ( Zeng et al. , 2020 ) . On the other hand , ClusterGCN ( Chiang et al. , 2019 ) is a heuristic in the sense that , despite its good performance , it provides no guarantee of unbiased gradient estimates of the full-batch learning . Gilmer et al . ( 2017 ) and Chami et al . ( 2021 ) generalized many GRL models into Message Passing and Auto-Encoder frameworks . These frameworks prompt bundling of GRL methods under Software Libraries , like PyG ( Fey & Lenssen , 2019 ) and DGL ( Wang et al. , 2019 ) , offering consistent interfaces on data formats . We now position our contribution relative to the above . Unlike generalized message passing ( Gilmer et al. , 2017 ) , rather than abstracting the model computation , we abstract the learning algorithm . As a result , GTTF can be specialized to recover the learning of MP as well as NE methods . Morever , unlike Software Frameworks , which are re-implementations of many algorithms and therefore inherit the scale and learning of the copied algorithms , we re-write the algorithms themselves , giving them new properties ( memory and computation complexity ) , while maintaining ( in expectation ) the original algorithm outcomes . Further , while the listed Stochastic Sampling Methods target MP algorithms ( such as GCN , GAT , alike ) , as their initial construction could not scale to large graphs , our learning algorithm applies to a wider class of GRL methods , additionally encapsulating NE methods . Finally , while some NE methods such as node2vec ( Grover & Leskovec , 2016 ) and DeepWalk ( Perozzi et al. , 2014 ) are scalable in their original form , their scalability stems from their multi-step process : sample many ( short ) random walks , save them to desk , and then learn node embeddings using positional embedding methods ( e.g. , word2vec , Mikolov et al . ( 2013 ) ) – they are sub-optimal in the sense that their first step ( walk sampling ) takes considerable time ( before training even starts ) and also places an artificial limit on the number of training samples ( number of simulated walks ) , whereas our algorithm conducts walks on-the-fly whilst training . 3 GRAPH TRAVERSAL VIA TENSOR FUNCTIONALS ( GTTF ) . At its core , GTTF is a stochastic algorithm that recursively conducts graph traversals to build representations of the graph . We describe the data structure and traversal algorithm below , using the following notation . G = ( V , E ) is an unweighted graph with n = |V | nodes and m = |E| edges , described as a sparse adjacency matrix A ∈ { 0 , 1 } n×n . Without loss of generality , let the nodes be zero-based numbered i.e . V = { 0 , . . . , n− 1 } . We denote the out-degree vector δ ∈ Zn – it can be calculated by summing over rows of A as δu = ∑ v∈V A [ u , v ] . We assume δu > 0 for all u ∈ V : pre-processing can add self-connections to orphan nodes . B denotes a batch of nodes . 3.1 DATA STRUCTURE . Internally , GTTF relies on a reformulation of the adjacency matrix , which we term CompactAdj ( for `` Compact Adjacency '' , Figure 1c ) . It consists of two tensors : 1. δ ∈ Zn , a dense out-degree vector ( figure 1c , right ) 2 . Â ∈ Zn×n , a sparse edge-list matrix in which the row u contains left-aligned δu non-zero values . The consecutive entries { Â [ u , 0 ] , Â [ u , 1 ] , . . . , Â [ u , δu − 1 ] } contain IDs of nodes receiving an edge from node u . The remaining |V | − δu are left unset , therefore , Â only occupies O ( m ) memory when stored as a sparse matrix ( Figure 1c , left ) . CompactAdj allows us to concisely describe stochastic traversals using standard tensor operations . To uniformly sample a neighbor to node u ∈ V , one can draw r ∼ U [ 0 .. ( δu − 1 ) ] , then get the neighbor ID with Â [ u , r ] . In vectorized form , given node batch B and access to continuous U [ 0 , 1 ) , we sample neighbors for each node in B as : R ∼ U [ 0 , 1 ) b , where b = |B| , then B′ = Â [ B , bR ◦ δ [ B ] c ] is a b-sized vector , with B′u containing a neighbor of Bu , floor operation b.c is applied element-wise , and ◦ is Hadamard product . 3.2 STOCHASTIC TRAVERSAL FUNCTIONAL ALGORITHM . Our traversal algorithm starts from a batch of nodes . It expands from each into a tree , resulting in a walk forest rooted at the nodes in the batch , as depicted in Figure 1d . In particular , given a node batch B , the algorithm instantiates |B| seed walkers , placing one at every node in B. Iteratively , each walker first replicates itself a fanout ( f ) number of times . Each replica then samples and transitions to a neighbor . This process repeats a depth ( h ) number of times . Therefore , each seed walker becomes the ancestor of a f -ary tree with height h. Setting f = 1 recovers traditional random walk . In practice , we provide flexibility by allowing a custom fanout value per depth . Algorithm 1 : Stochastic Traverse Functional , parametrized by ACCUMULATEFN and BIASFN . input : u ( current node ) ; T ← [ ] ( path leading to u , starts empty ) ; F ( list of fanouts ) ; ACCUMULATEFN ( function : with side-effects and no return . It is model-specific and records information for computing model and/or objective , see text ) ; BIASFN ← U ( function mapping u to distribution on u ’ s neighbors , defaults to uniform ) 1 def Traverse ( T , u , F , ACCUMULATEFN , BIASFN ) : 2 if F.size ( ) = 0 then return # Base case . Traversed up-to requested depth 3 f ← F.pop ( ) # fanout duplication factor ( i.e . breadth ) at this depth . 4 sample_bias ← BIASFN ( T , u ) 5 if sample_bias.sum ( ) = 0 then return # Special case . No sampling from zero mass 6 sample_bias ← sample_bias / sample_bias.sum ( ) # valid distribution 7 K ← Sample ( Â [ u , : δu ] ] , sample_bias , f ) # Sample f nodes from u ’ s neighbors 8 for k ← 0 to f − 1 do 9 Tnext ← concatenate ( T , [ u ] ) 10 ACCUMULATEFN ( Tnext , K [ k ] , f ) 11 Traverse ( Tnext , K [ k ] , f , ACCUMULATEFN , BIASFN ) # Recursion 12 13 def Sample ( N , W , f ) : 14 C ← tf.cumsum ( W ) # Cumulative sum . Last entry must = 1 . 15 coin_flips← tf.random.uniform ( ( f , ) , 0 , 1 ) 16 indices← tf.searchsorted ( C , coin_flips ) 17 return N [ indices ] Functional Traverse is listed in Algorithm 1 . It accepts : a batch of nodes2 ; a list of fanout values F ( e.g . to F = [ 3 , 5 ] samples 3 neighbors per u ∈ B , then 5 neighbors for each of those ) ; and more notably , two functions : ACCUMULATEFN and BIASFN . These functions will be called by the functional on every node visited along the traversal , and will be passed relevant information ( e.g . the path taken from root seed node ) . Custom settings of these functions allow recovering wide classes of graph learning methods . At a high-level , our functional can be used in the following manner : 1 . Construct model & initialize parameters ( e.g . to random ) . Define ACCUMULATEFN and BIASFN . 2 . Repeat ( many rounds ) : i. Reset accumulation information ( from previous round ) and then sample batch B ⊂ V . ii . Invoke Traverse on ( B , ACCUMULATEFN , BIASFN ) , which invokes the FN ’ s , allowing the first to accumulate information sufficient for running the model and estimating an objective . iii . Use accumulated information to : run model , estimate objective , apply learning rule ( e.g . SGD ) . ACCUMULATEFN is a function that is used to track necessary information for computing the model and the objective function . For instance , an implementation of DeepWalk ( Perozzi et al. , 2014 ) on top of GTTF , specializes ACCUMULATEFN to measure an estimate of the sampled softmax likelihood of nodes ’ positional distribution , modeled as a dot-prodct of node embeddings . On the other hand , GCN ( Kipf & Welling , 2017 ) on top of GTTF uses it to accumulate a sampled adjacency matrix , which it passes to the underlying model ( e.g . 2-layer GCN ) as if this were the full adjacency matrix . BIASFN is a function that customizes the sampling procedure for the stochastic transitions . If provided , it must yield a probability distribution over nodes , given the current node and the path that lead to it . If not provided , it defaults to U , transitioning to any neighbor with equal probability . It can be defined to read edge weights , if they denote importance , or more intricately , used to parameterize a second order Markov Chain ( Grover & Leskovec , 2016 ) , or use neighborhood attention to guide sampling ( Veličković et al. , 2018 ) , as discussed in the Appendix . 2Our pseudo-code displays the traversal starting from one node rather than a batch only for clarity , as our actual implementation is vectorized e.g . u would be a vector of nodes , T would be a 2D matrix with each row containing transition path preceeding the corresponding entry in u , ... etc . Refer to Appendix and code . 3.3 SOME SPECIALIZATIONS OF ACCUMULATEFN & BIASFN

Summary: The authors propose a "meta-algorithm" for approximating various graph representation learning schemes: generate batches of random trees with fixed fanout (and possibly biased probabilities of selecting different edges), and use them to accumulate information to approximate operations on the graph.  The idea is beautifully simple, and generalizes running independent random walkers, an approach that is used in deriving many related algorithms.  The biasing and accumulation functions are user provided, and the authors show how different choices of these functions can be used to approximate a number of graph representation learning schemes.  The authors also provide a software framework, though it was inaccessible at review due to anonymization.  Experiments show the approach is much more scalable than competing approaches (though, to be fair, some of the competition was not targeting scalability).