Scalable Deep Generative Modeling for Sparse Graphs

The main scalability bottleneck is the large output space. One way to reduce the output space size is to use a hierarchical recursive decomposition, inspired by the classic random graph model R-MAT (Chakrabarti et al., 2004). In R-MAT, each edge is put into the adjacency matrix by dividing the 2D matrix into 4 equally sized quadrants, and recursively descending into one of the quadrants, until reaching a single entry of the matrix. In this generation process, the complexity of generating each edge is only $O(\log n)$.

We adopt the recursive decomposition design of R-MAT, and further simplify and neuralize it, to make our model efficient and expressive. In our model, we generate edges following the node ordering row by row, so that each edge only needs to be put into the right place in a single row, reducing the process to 1D, as illustrated in Figure 1. For any edge $(u,v)$, the process of picking node $v$ as one of $u$’s neighbors starts by dividing the node index interval $[1,n]$ in half, then recursively descending into one half until reaching a single entry. Each $v$ corresponds to a unique sequence of decisions $x^{v}_{1},...,x^{v}_{d}$, where $x^{v}_{i}\in\{\text{left},\text{right}\}$ is the $i$-th decision in the sequence, and $d=\lceil\log_{2}n\rceil$ is the maximum number of required decisions to specify $v$.

The probability of $p(v|u)$ can then be formulated as

where each $p(x_{i}=x^{v}_{i})$ is the probability of following the decision that leads to $v$ at step $i$.

Let us use $E_{u}=\left\{(u,v)\in E\right\}$ to denote the set of edges incident to node $u$, and $\mathcal{N}_{u}=\{v|(u,v)\in E_{u}\}$. Generating only a single edge is similar to hierarchical softmax (Mnih & Hinton, 2009), and applying the above procedure repeatedly can generate all of $|\mathcal{N}_{u}|$ edges in $O(|\mathcal{N}_{u}|\log n)$ time. But we can do better than that when generating all these edges.

Further improvement using binary trees. As illustrated in the left half of Figure 2, the process of jointly generating all of $E_{u}$ is equivalent to building up a binary tree ${\mathcal{T}}_{u}$, where each tree node $t\in{\mathcal{T}}_{u}$ corresponds to a graph node index interval $[v_{l},v_{r}]$, and for each $v\in\mathcal{N}_{u}$ the process starts from the root $[1,n]$ and ends in a leaf $[v,v]$.

Taking this perspective, we propose a more efficient generation process for $E_{u}$, which generates the tree directly instead of repeatedly generating each leaf through a path from the root. We propose a recursive process that builds up the tree following a depth-first or in-order traversal order, where we start at the root, and recursively for each tree node $t$: (1) decide if $t$ has a left child denoted as $\text{lch}(t)$, and (2) if so recurse into $\text{lch}(t)$ and generate the left sub-tree, and then (3) decide if $t$ has a right child denoted as $\text{rch}(t)$, (4) if so recurse into $\text{rch}(t)$ and generate the right sub-tree, and (5) return to $t$’s parent. This process is shown in Algorithm 1, which will be elaborated in next section.

Overloading the notation a bit, we use $p(\text{lch}(t))$ to denote the probability that tree node $t$ has a left child, when $\text{lch}(t)\neq\emptyset$, or does not have a left child, when $\text{lch}(t)=\emptyset$, under our model, and similarly define $p(\text{rch}(t))$. Then the probability of generating $E_{u}$ or equivalently tree ${\mathcal{T}}_{u}$ is

This new process generates each tree node exactly once, hence the time complexity is proportional to the tree size $O(|{\mathcal{T}}_{u}|)$, and it is clear that $|{\mathcal{T}}_{u}|\leq|\mathcal{N}_{u}|\log n$, since $|\mathcal{N}_{u}|$ is the number of leaf nodes in the tree and $\log n$ is the max depth of the tree. The time saving comes from removing the duplicated effort near the root of the tree. When $|\mathcal{N}_{u}|$ is large, i.e. as some fraction of $n$ when $u$ is one of the “hub” nodes in the graph, the tree ${\mathcal{T}}_{u}$ becomes dense and our new generation process will be significantly faster, as the time complexity becomes close to $O(n)$ while generating each leaf from the root would require $\Omega(n\log n)$ time.

In the following, we present our approach to make this model fully autoregressive, i.e. making $p(\text{lch}(t))$ and $p(\text{rch}(t))$ depend on all the decisions made so far in the process of generating the graph, and make this model neuralized so that all the probability values in the model come from expressive deep neural networks.

In this section we consider how to add autoregressive conditioning to $p(\text{lch}(t))$ and $p(\text{rch}(t))$ when generating ${\mathcal{T}}_{u}$.

In our generation process, the decision about whether $\text{lch}(t)$ exists for a particular tree node $t$ is made after $t$, all its ancestors, and all the left sub-trees of the ancestors are generated. We can use a top-down context vector $h^{top}_{u}(t)$ to summarize all these contexts, and modify $p(\text{lch}(t))$ to $p(\text{lch}(t)|h^{top}_{u}(t))$. Similarly, the decision about $\text{rch}(t)$ is made after generating $\text{lch}(t)$ and its dependencies, and $t$’s entire left-subtree (see Figure 2 right half for illustration). We therefore need both the top-down context $h^{top}_{u}(t)$, as well as the bottom-up context $h^{bot}_{u}(\text{lch}(t))$ that summarizes the sub-tree rooted at $\text{lch}(t)$, if any. The autoregressive model for $p({\mathcal{T}}_{u})$ therefore becomes

and we can recursively define

where $\text{TreeCell}^{bot}$ and $\text{TreeCell}^{top}$ are two TreeLSTM cells (Tai et al., 2015) that combine information from the incoming nodes into a single node state, and $\text{embed}(\text{left})$ and $\text{embed}(\text{right})$ represents the embedding vector for the binary values “left” and “right”. We initialize $h^{bot}_{u}(\emptyset)=\vec{0}$, and discuss $h^{top}_{u}(\text{root})$ in the next section.

The distributions can then be parameterized as

With the efficient recursive edge generation and autoregressive conditioning presented in Section 2.1 and Section 2.2 respectively, we are ready to present the full autoregressive model for generating the entire adjacency matrix $A$.

The full model will utilize the autoregressive model for $\mathcal{N}_{u}$ as building blocks. Specifically, we are going to generate the adjacency matrix $A$ row by row:

Let $g_{u}^{0}=h_{u}^{bot}(t_{1})$ be the embedding that summarizes ${\mathcal{T}}_{u}$, suppose we have an efficient mechanism to encode $\left[g_{1},g_{2},\ldots,g_{u}\right]$ into $h^{row}_{u}$, then we can effectively use $h^{row}_{u-1}$ to generate ${\mathcal{T}}_{u}$ and thus the entire process would become autoregressive. Again, since there are $n$ rows in total, using a chain structured LSTM would make the history length too long for large graphs. Therefore, we use an approach inspired by the Fenwick tree (Fenwick, 1994) which is a data structure that maintains the prefix sum efficiently. Given an array of numbers, the Fenwick tree allows calculating any prefix sum or updating any single entry in $O(\log L)$ for a sequence of length $L$. We build such a tree to maintain and update the prefix embeddings. We denote it as row-binary forest as such data structure is a forest of binary trees.

Figure 3 demonstrates one solution. Before generating the edge-binary tree ${\mathcal{T}}_{u}$, the embeddings that summarize each individual edge-binary tree $\mathcal{R}_{u}=\left\{{\mathcal{T}}_{u^{\prime}}:u^{\prime}<u\right\}$ will be organized into the row-binary forest $\mathcal{G}_{u}$. This forest is organized into $\left\lfloor\log(u-1)\right\rfloor+1$ levels, with the bottom $0$-th level as edge-binary tree embeddings. Let $g^{i}_{j}\in\mathcal{G}_{u}$ be the $j$-th node in the $i$-th level, then

where $0\leq i\leq\left\lfloor\log(u-1)\right\rfloor+1,1\leq j\leq\left\lfloor\frac{|\mathcal{R}_{u}|}{2^{i}}\right\rfloor$.

Algorithm 2 summarizes the entire procedure for sampling a graph from our model in a fully autoregressive manner.

A classical autoregressive model like LSTM is fully sequential, which allows no concurrency across steps. Thus training LSTM with a sequence of length $L$ takes $\Omega(L)$ of synchronized computation. In this section, we show how to exploit the characteristic of BiGG to increase concurrency during training.

Given a graph $G$, estimating the likelihood under the current model can be divided into four steps.

1. 1.

   Computing $h_{u}^{bot}(t),\forall u\in V,t\in{\mathcal{T}}_{u}$ : as the graph is given during training, the corresponding edge-binary trees $\left\{{\mathcal{T}}_{u}\right\}$ are also known. From Eq (7) we can see that the embeddings $h_{u}^{bot}(t)$ of all nodes at the same depth of tree can be computed concurrently without dependence, and the synchronization only happens between different depths. Thus $O(\log n)$ steps of synchronization is sufficient.

2. 2.

   Computing $g_{j}^{i}\in\mathcal{G}_{n}$: the row-binary forest grows monotonically, thus the forest $\mathcal{G}_{n}$ in the end contains all the embeddings needed for computing $\left\{h_{u}^{row}\right\}$. Similarly, computing $\left\{g_{j}^{i}\right\}$ synchronizes $O(\log n)$ steps.

3. 3.

   Computing $h_{u}^{row},\forall u\in V$: as each $h_{u}^{row}$ runs an LSTM independently, this stage simply runs LSTM on a batch of $n$ sequences with length $O(\log n)$.

4. 4.

   Computing $h_{u}^{top}$ and likelihood: the last step computes the likelihood using Eq (8) and  (9). This is similar to the first step, except that the computation happens in a top-down direction in each ${\mathcal{T}}_{u}$.

Figure 4 demonstrates this process. In summary, the four stages each take $O(\log n)$ steps of synchronization. This allows us to train large graphs much more efficiently than a simple sequential autoregressive model.

It is possible that during training the graph is too large to fit into memory. Thus to train on large graphs, model parallelism is more important than data parallelism.

To split the model, as well as intermediate computations, into different machines, we divide the adjacency matrix into multiple consecutive chunks of rows, where each machine is responsible for one chunk. It is easy to see that by doing so, Stage 1 and Stage 4 mentioned in Section 3.1 can be executed concurrently on all the machines without any synchronization, as the edge-binary trees can be processed independently once the conditioning states like $\left\{h_{u}^{row}\right\}$ are made ready by synchronization.

Figure 5 illustrates the situation when training a graph with 7 nodes using 2 GPUs. During Stage 1, GPU 1 and 2 work concurrently to compute $\left\{g_{u}^{0}\right\}_{u=1}^{3}$ and $\left\{g_{u}^{0}\right\}_{u=4}^{7}$, respectively. In Stage 2, the embeddings $g_{1}^{1}$ and $g_{3}^{0}$ are needed by GPU 2 when computing $g_{1}^{2}$ and $g_{2}^{1}$. We denote such embeddings as ‘g-messages’. Note that such ‘g-messages’ will transit in the opposite direction when doing a backward pass in the gradient calculation. Passing ‘g-messages’ introduces serial dependency across GPUs. However as the number of such embeddings is upper bounded by $O(\log n)$ depth of row-binary forest, the communication cost is manageable.

We divide rows as in Section 3.2. During the forward pass, only the ‘g-messages’ between chunks are kept. The only difference from the previous section is that the edge-binary tree embeddings are recomputed, due to the single GPU limitation. The memory cost will be:

Here $O(k\log n)$ accounts for the memory holding the ‘g-message’, and $O(\frac{m}{k})$ accounts for the memory of ${\mathcal{T}}_{u}$ in each chunk. The optimal $k$ is achieved when $k\log n=\frac{m}{k}$, hence $k=O(\sqrt{\frac{m}{\log n}})$ and the corresponding memory cost is $O(\sqrt{m\log n})$. Also note that such sublinear cost requires only one additional feedforward in Stage 1, so this will not hurt much of the training speed.

Note that to represent an interval $A[u,v_{l}:v_{r}]$ of length $b=v_{r}-v_{l}+1\leq L$, we use vector $\mathbf{v}\in\left\{-1,0,1\right\}^{L}$ where

That is to say, we use ternary bit vector to encode both the interval length and the binary adjacency information.

During generation of $\left\{{\mathcal{T}}_{u}\right\}$, each tree node $t$ of the edge-binary tree knows the span $[v_{l},v_{r}]$ which corresponds to the columns it will cover. One way is to augment $h^{top}_{u}(t)$ with the position encoding as:

where PE is the position encoding using sine and cosine functions of different frequencies as in Vaswani et al. (2017). Similarly, the $h_{u}^{row}$ in Eq (12) can be augmented by $\text{PE}(n-u)$ in a similar way. With such augmentation, the model will know more context into the future, and thus help improve the generative quality.

Please refer to our released open source code located at https://github.com/google-research/google-research/tree/master/bigg for more implementation and experimental details.

In this part, we compare the quality of our model with previous work on a set of benchmark datasets. We present results on median sized general graphs with number of nodes ranging in 0.1k to 5k in Section 4.1.1, and on large SAT graphs with up to 20k nodes in Section 4.1.2. In Section 4.3 we perform ablation studies of BiGG with different sparsity and node orders.

In this section, we will evaluate the scalability of BiGG regarding the time complexity, memory consumption and the quality of generated graphs with respect to the number of nodes in graphs.

In this section, we take a deeper look at the performance of BiGG with different node ordering in Section 4.3.1. We also show the effect of edge density to the generative performance in Section 4.3.2.