GRAPH TRAVERSAL WITH TENSOR FUNCTIONALS: A META-ALGORITHM FOR SCALABLE LEARNING

Abstract

Graph Representation Learning (GRL) methods have impacted fields from chemistry to social science. However, their algorithmic implementations are specialized to specific use-cases e.g. message passing methods are run differently from node embedding ones. Despite their apparent differences, all these methods utilize the graph structure, and therefore, their learning can be approximated with stochastic graph traversals. We propose Graph Traversal via Tensor Functionals (GTTF), a unifying meta-algorithm framework for easing the implementation of diverse graph algorithms and enabling transparent and efficient scaling to large graphs. GTTF is founded upon a data structure (stored as a sparse tensor) and a stochastic graph traversal algorithm (described using tensor operations). The algorithm is a functional that accept two functions, and can be specialized to obtain a variety of GRL models and objectives, simply by changing those two functions. We show for a wide class of methods, our algorithm learns in an unbiased fashion and, in expectation, approximates the learning as if the specialized implementations were run directly. With these capabilities, we scale otherwise non-scalable methods to set state-of-the-art on large graph datasets while being more efficient than existing GRL libraries -with only a handful of lines of code for each method specialization. GTTF and its various GRL implementations are on: https://github.

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; Veličković et al. ( 2018 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 tensorfoot_0 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. 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 0 1 2 3 4 (a) Example graph G      1 1 1 1 1 1 1 1 1 1      (b) Adjacency matrix for graph G      1 0 2 3 4 1 1 4 1 3           1 4 1 2 2      (c) CompactAdj for G with sparse A ∈ Z n×n and dense δ ∈ Z n . We store IDs of adjacent nodes in A 0 1 1 2 4 2 0 4 1 3 0 2 1 4 (d) Walk Forest. GTTF in- vokes AC C U M U L A T EFN once per (green) instance. 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. 



To 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



); 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.

Figure 1: (c)&(d) Depict our data structure & traversal algorithm on a toy graph in (a)&(b).

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.

