# RSC: Accelerating Graph Neural Networks Training via Randomized Sparse
Computations

Zirui Liu<EMAIL_ADDRESS>Rice UniversityUnited States ,  Shengyuan Chen
<EMAIL_ADDRESS>Hong Kong Polytechnic UniversityHong Kong  ,
Kaixiong Zhou<EMAIL_ADDRESS>Rice UniversityUnited States ,  Daochen
Zha<EMAIL_ADDRESS>Rice UniversityUnited States ,  Xiao Huang
<EMAIL_ADDRESS>Hong Kong Polytechnic UniversityHong Kong  and
Xia Hu<EMAIL_ADDRESS>Rice UniversityUnited States

###### Abstract.

The training of graph neural networks (GNNs) is extremely time consuming
because sparse graph-based operations are hard to be accelerated by hardware.
Prior art explores trading off the computational precision to reduce the time
complexity via sampling-based approximation. Based on the idea, previous works
successfully accelerate the dense matrix based operations (e.g., convolution
and linear) with negligible accuracy drop. However, unlike dense matrices,
sparse matrices are stored in the irregular data format such that each
row/column may have different number of non-zero entries. Thus, compared to
the dense counterpart, approximating sparse operations has two unique
challenges (1) we cannot directly control the efficiency of approximated
sparse operation since the computation is only executed on non-zero entries;
(2) sub-sampling sparse matrices is much more inefficient due to the irregular
data format. To address the issues, our key idea is to control the accuracy-
efficiency trade off by optimizing computation resource allocation layer-
wisely and epoch-wisely. Specifically, for the first challenge, we customize
the computation resource to different sparse operations, while limit the total
used resource below a certain budget. For the second challenge, we cache
previous sampled sparse matrices to reduce the epoch-wise sampling overhead.
Finally, we propose a switching mechanisms to improve the generalization of
GNNs trained with approximated operations. To this end, we propose Randomized
Sparse Computation, which for the first time demonstrate the potential of
training GNNs with approximated operations. In practice, RSC can achieve up to
$11.6\times$ speedup for a single sparse operation and a $1.6\times$ end-to-
end wall-clock time speedup with negligible accuracy drop.

Graph Neural Networks, Efficient Training, Randomized Algorithm

††ccs: Computing methodologies Machine learning

##  1\. Introductions

Graph Neural Networks (GNNs) have achieved great success across different
graph-related tasks (Hamilton et al., 2017; Hu et al., 2020; Ying et al.,
2018; Jiang et al., 2022). However, despite its effectiveness, the training of
GNNs is very time consuming. Specifically, GNNs are characterized by an
interleaved execution that switches between the aggregation phase and the
update phase. Namely, in the aggregation phase, every node aggregates messages
from its neighborhoods at each layer, which is implemented based on sparse
matrix-based operations. (Fey and Lenssen, 2019; Wang et al., 2019) In the
update phase, each node will update its embedding based on the aggregated
messages, where the update function is implemented based on dense matrix-based
operations (Fey and Lenssen, 2019; Wang et al., 2019). In Figure 1, SpMM and
MatMul are the sparse operations and dense operations in the aggregation and
update phase, respectively. Through profiling, we found that the aggregation
phase may take more than $90\%$ running time for GNN training. This is because
the sparse matrix operations in the aggregation phase have many random memory
accesses and limited data reuse, which is hard to be accelerated by community
hardwares (e.g., CPUs and GPUs) (Duan et al., 2022a; Han et al., 2016; Duan et
al., 2022b). Thus, training GNNs with large graphs is often time-inefficient.

Figure 1. The time profiling of a two-layer GCNs on different datasets. SpMM
may take $70\%\sim 90\%$ of the total training time.

Most of the existing works towards this problem can be roughly divided into
three categories. First, some works propose the distributed GNNs training
system with extra hardware (Zheng et al., 2020; Ramezani et al., 2022; Wan et
al., 2022). Second, some works optimize the memory access pattern of sparse
operations via coalescing the memory access and fusing consecutive operations
(Zhang et al., 2022; Huang et al., 2020; Rahman et al., 2021). Third, some
works propose to accelerate the training process from the optimization aspect,
i.e., using fewer iterations to converge (Narayanan et al., 2022; Cong et al.,
2020).

In parallel, another orthogonal direction is to replace the expensive
operations with their faster approximated versions (Adelman et al., 2021;
Drineas et al., 2006a). The key idea is to sub-sample tensors onto low
dimensional spaces and perform the original operations here. For example, for
the linear operation between two matrices $\mathbf{A}\in\mathbb{R}^{n\times
m}$ and $\mathbf{B}\in\mathbb{R}^{m\times q}$, we first obtain
$\mathbf{A}^{\prime}\in\mathbb{R}^{n\times k}$ and
$\mathbf{B}^{\prime}\in\mathbb{R}^{k\times q}$ ($k<m$) by selecting $k$
representative columns of $\mathbf{A}$ and the corresponding rows of
$\mathbf{B}$ based on their euclidean norm (Drineas et al., 2006a). Then we
approximate
$\mathbf{A}\mathbf{B}\approx\mathbf{A^{\prime}}\mathbf{B^{\prime}}$. As a
result, the number of floating-point operations (FLOPs) and memory access are
both reduced. Based on the idea, previous work successfully accelerates the
dense matrix based operations, such as convolution and linear operations
(Adelman et al., 2021). The approximated operation can plug-and-play replace
the exact operation to improve per-operation efficiency, and thus is
compatible with most of the efficient training methods. Despite the potential,
this perspective however has not been explored for the sparse operations in
GNNs.

The approximation method reduces the computational complexity at the cost of
giving noisy outputs, thus there naturally exists an accuracy-efficiency trade
off. Compared to approximating dense matrix operations, it has two unique
challenges to optimize the trade off for approximated sparse operations.
First, unlike the previous example of approximating linear operation, $k$
cannot directly control the efficiency (FLOPs) for sparse operations. This is
because for dense matrices, each row/column has the same amount of parameters.
Thus the reduction of FLOPs in approximated dense operations is totally
decided by the dimensions of the sub-sampled matrices (i.e., $k$). However, in
sparse operations, each row/column in the sparse adjacency matrix has
different numbers of non-zero entries and the computation is only executed on
non-zero entries (i.e., irregular data format). Thus, the reduction of FLOPs
in the sparse operations is decided by the selection of representative
row/columns. It lacks of the mechanism to directly control the efficiency-
accuracy trade off for each sparse operation. Second, compared to the dense
counterpart, slicing the sparse matrix is much more time-consuming due to this
irregular data format (Fey and Lenssen, 2019), which counteracts the
acceleration from the FLOPs reduction.

To this end, we propose Randomized Sparse Computation, dubbed RSC , the first
approximation framework tailored for efficient GNN training. As we analyzed,
unlike the dense counterpart, we lack the mechanism to directly control the
efficiency-accuracy trade off for individual sparse operations. Instead, we
propose to control the trade off by optimizing the computation resource
allocation at the “global” level. Specifically, to tackle the first challenge,
at the layer-wise level, since a GNN contains several sparse operations and
each operation has a different importance to model accuracy, we propose to
customize the FLOPs of different sparse operations while limiting the total
FLOPs under a certain budget. Specifically, we frame it as a constraint
optimization problem. Then we propose a greedy algorithm to efficiently solve
it. To tackle the Second challenge, at the epoch-wise level, we found that the
selection of representative row/columns tend to remain similar across nearby
iterations. Based on this finding, we develop a caching mechanism to reuse the
previous sampled sparse matrix across nearby iterations to reduce per-epoch
sampling time. Finally, inspired by the recent finding that the final stage of
training usually needs smaller noise to help convergence (Li et al., 2019), we
propose to use approximated sparse operation during most of the training
process, while switch back to the original sparse operation at the final
stage. We find that this switching mechanism significantly reduce the accuracy
drop, at the cost of slightly less acceleration effect. In summary,

  * •

We accelerate the training of GNNs from a brand new perspective, namely,
replacing the expensive sparse operations with their faster approximated
versions.

  * •

Instead of focusing on balancing the efficiency-accuracy trade-off at
operation level, we control the trade off through optimizing resource
allocation at the layer-wise level and epoch-wise level.

  * •

We propose a caching mechanism to reduce the per-epoch cost of sampling the
sparse matrix by reusing previous results.

  * •

We propose a switching mechanism to minimize the accuracy drop by switching
back to the exact operations at the final stage.

  * •

Extensive experiments have demonstrated the effectiveness of the proposed
method. Particularly, RSC can achieve up to $11.6\times$ speedup for a single
sparse operation and a $1.6\times$ end-to-end wall-clock time speedup with
negligible ($\approx 0.3\%$) accuracy drop.

##  2\. Background and Preliminary

###  2.1. Graph Neural Networks

Let ${\mathcal{G}}=({\mathcal{V}},{\mathcal{E}})$ be an undirected graph with
${\mathcal{V}}=(v_{1},\cdots,v_{|{\mathcal{V}}|})$ and
${\mathcal{E}}=(e_{1},\cdots,e_{|{\mathcal{E}}|})$ being the set of nodes and
edges, respectively. Let ${\bm{X}}\in\mathbb{R}^{|{\mathcal{V}}|\times d}$ be
the node feature matrix.
${\bm{A}}\in\mathbb{R}^{|{\mathcal{V}}|\times|{\mathcal{V}}|}$ is the graph
adjacency matrix, where ${\bm{A}}_{i,j}=1$ if $(v_{i},v_{j})\in{\mathcal{E}}$
else ${\bm{A}}_{i,j}=0$.
$\tilde{{\bm{A}}}=\tilde{{\bm{D}}}^{-\frac{1}{2}}({\bm{A}}+{\bm{I}})\tilde{{\bm{D}}}^{-\frac{1}{2}}$
is the normalized adjacency matrix, where $\tilde{{\bm{D}}}$ is the degree
matrix of ${\bm{A}}+{\bm{I}}$. GNNs recursively update the embedding of a node
by aggregating embeddings of its neighbors. For example, the forward pass of
the $l^{\mathrm{th}}$ Graph Convolutional Network (GCN) layer (Kipf and
Welling, 2017) can be defined as:

(1)
${\bm{H}}^{(l+1)}=\textrm{ReLU}(\tilde{{\bm{A}}}{\bm{H}}^{(l)}{\bm{\Theta}}^{(l)}),$

where ${\bm{H}}^{(l)}$ is the node embedding matrix at the $l^{\mathrm{th}}$
layer and ${\bm{H}}^{(0)}={\bm{X}}$. ${\bm{\Theta}}^{(l)}$ is the weight
matrix of the $l^{\mathrm{th}}$ GCN layer.

In practice, $\tilde{{\bm{A}}}$ is often stored in the sparse matrix format,
e.g., compressed sparse row (CSR) (Fey and Lenssen, 2019). From the
implementation aspect, the computation of Equation (1) can be described as:

(2) $\displaystyle{\bm{H}}^{(l+1)}$
$\displaystyle=\textrm{ReLU}\Bigg{(}\texttt{SpMM}\bigg{(}\tilde{\bm{A}},\texttt{MatMul}({\bm{H}}^{(l)},{\bm{\Theta}}^{(l)})\bigg{)}\Bigg{)},$

where $\texttt{SpMM}(\cdot,\cdot)$ is the Sparse-Dense Matrix Multiplication
and $\texttt{MatMul}(\cdot,\cdot)$ is the Dense Matrix Multiplication. Unlike
dense matrices, the elements are randomly distributed in the sparse matrix.
Thus sparse operations, such as SpMM , have many random memory accesses and
are much slower than the dense counter-part (Han et al., 2016; Duan et al.,
2022a). To get a sense of the scale, we shown in Figure 1 that for GCNs, SpMM
may take roughly $70\%\sim 90\%$ of the total training time.

###  2.2. Fast Approximated MatMul with Sampling

Let ${\bm{X}}\in\mathbb{R}^{n\times m}$, ${\bm{Y}}\in\mathbb{R}^{m\times q}$.
The goal is to efficiently estimate the matrix production ${\bm{X}}{\bm{Y}}$.
Truncated Singular Value Decomposition (SVD) outputs provably optimal low-rank
estimation of ${\bm{X}}{\bm{Y}}$ (Adelman et al., 2021). However, SVD is
almost as expensive as the matrix production itself. Instead, the sampling
algorithm is proposed to approximate the matrix product ${\bm{X}}{\bm{Y}}$ by
sampling $k$ columns of ${\bm{X}}$ and corresponding rows of ${\bm{Y}}$ to
form smaller matrices, which are then multiplied as usual (Drineas et al.,
2006a). This algorithm reduces the computational complexity from
$\mathcal{O}(mnq)$ to $\mathcal{O}(knq)$. Specifically,

$\displaystyle{\bm{X}}{\bm{Y}}=\sum_{i=1}^{m}{\bm{X}}_{:,i}{\bm{Y}}_{i,:}$
$\displaystyle\approx\sum_{t=1}^{k}\frac{1}{s_{t}}{\bm{X}}_{:,i_{t}}{\bm{Y}}_{i_{t},:}$
(3) $\displaystyle=\texttt{approx}({\bm{X}}{\bm{Y}}),$

where ${\bm{X}}_{:,i}\in\mathbb{R}^{n\times 1}$ and
${\bm{Y}}_{i,:}\in\mathbb{R}^{1\times q}$ are the $i^{\mathrm{th}}$ column and
row of ${\bm{X}}$ and ${\bm{Y}}$, respectively. In this paper, we call
$({\bm{X}}_{:,i},{\bm{Y}}_{i,:})$ the $i^{\mathrm{th}}$ column-row pair. $k$
is the number of samples ($1\leq k\leq m$). $\\{p_{i}\\}_{i=1}^{m}$ is a
probability distribution over the column-row pairs. $i_{t}\in\\{1,\cdots m\\}$
is the index of the sampled column-row pair at the $t^{\mathrm{th}}$ trial.
$s_{t}$ is the scale factor. Theoretically, (Drineas et al., 2006a) shows that
if we set $s_{t}=\frac{1}{kp_{i_{t}}}$, then we have
$\mathbb{E}[\texttt{approx}({\bm{X}}{\bm{Y}})]={\bm{X}}{\bm{Y}}$. Further, the
expectation of approximation error
$\mathbb{E}[||{\bm{X}}{\bm{Y}}-\texttt{approx}({\bm{X}}{\bm{Y}})||_{F}]$ is
minimized when the sampling probabilities $\\{p_{i}\\}_{i=1}^{m}$ are
proportional to the product of the column-row Euclidean norms (Drineas et al.,
2006a):

(4) $p_{i}=\frac{||{\bm{X}}_{:,i}||_{2}\
||{\bm{Y}}_{i,:}||_{2}}{\sum_{j=1}^{m}||{\bm{X}}_{:,j}||_{2}\
||{\bm{Y}}_{j,:}||_{2}}.$

####  2.2.1. Top-$k$ sampling

The above sampling based approximation method is originally developed for
accelerating the general application of dense matrix multiplication (Drineas
et al., 2006a). Directly applying it to neural network may be sub-optimal
because it does not consider the characteristic of neural networks. (Adelman
et al., 2021) modifies the sampling algorithm by taking the weight
distribution of neural network into account. Based on the empirical
observation that the distribution of weights remains centered around zero
during training (Glorot and Bengio, 2010; Han et al., 2015), (Adelman et al.,
2021) proposes a top-$k$ sampling algorithm as follows:

top-$k$ sampling: Selecting $k$ column-row pairs with the largest
$\frac{||{\bm{X}}_{:,i}||_{2}\
||{\bm{Y}}_{i,:}||_{2}}{\sum_{j=1}^{m}||{\bm{X}}_{:,j}||_{2}\
||{\bm{Y}}_{j,:}||_{2}}$ values deterministically without scaling.

Equivalently, it means $p_{i}$ of column-row pairs with the $k$-largest value
in Equation (4) equals 1, otherwise it equals 0. And $s_{i_{t}}$ is a constant
1. Albeit without the scaling while sampling column-row pairs
deterministically, under on the assumption of zero-centered weight
distribution, (Adelman et al., 2021) theoretically show that top-$k$ sampling
still yield an unbiased estimation of ${\bm{A}}{\bm{B}}$ with minimal
approximation error. Consequently, the top-$k$ sampling algorithm empirically
shows significantly lower accuracy drop when approximating the convolution and
linear operations in the neural networks (Adelman et al., 2021).

Note that the approximation method reduces not only FLOPs, but also the number
of memory accesses since we do not need to read the dropped elements, which is
crucial for accelerating memory-bound operation, such as SpMM (Han et al.,
2016).

In the next section, we explore how to approximate the expensive sparse
operation via the top-$k$ sampling.

##  3\. The Proposed Framework

Figure 2. Overview of RSC . For convenience, ReLU is ignored. RSC only replace
the SpMM in the backward pass with its approximated version using top-$k$
sampling (Section 3.1). $k_{l}$ is the number of samples for top-$k$ sampling
at the $l^{\mathrm{th}}$ layer, which is automatically allocated (Section
3.2). To reduce the overhead of sampling, we also cache the sampled graph and
reuse it across nearby iterations (Section 3.3).

The overview of our proposed framework is shown in Figure 2, where we use the
computation graph of GCN as an example. We first explore which SpMM in the
computation graph can be replaced with its approximated version (Section 3.1).
Then since GNNs have multiple SpMM and each of them may have different
importance to the model performance, we then automatically allocate
computation resource to different SpMM (Section 3.2). Finally, we explore two
simple and effective tricks for improving RSC . The first one is the caching
mechanism to reduce the overhead of sampling sparse matrices (Section 3.3.1),
the second one is the switching mechanism to reduce the accuracy drop (Section
3.3.2).

###  3.1. Where to Apply the Approximation

####  3.1.1. Experimental analysis

Each sparse operation is executed twice at each training step, i.e., one in
the forward pass and the other one in the backward pass. As shown in Figure 2,
here we take SpMM in the $l^{\mathrm{th}}$ GCN layer as an example, the
forward one is
${\bm{H}}^{(l+1)}=\textrm{ReLU}(\texttt{SpMM}(\tilde{{\bm{A}}},{\bm{J}}^{(l)}))$,
where ${\bm{J}}^{(l)}=\texttt{MatMul}({\bm{H}}^{(l)},{\bm{\Theta}}^{(l)})$ is
the intermediate node representations. And the backward one is
$\nabla{\bm{J}}^{(l)}=\texttt{SpMM}(\tilde{{\bm{A}}}^{\top},\nabla{\bm{H}}^{(l+1)})$.
$\nabla{\bm{J}}^{(l)}$ and $\nabla{\bm{H}}^{(l)}$ are the gradient with
respect to ${\bm{J}}^{(l)}$ and ${\bm{H}}^{(l)}$, respectively.

Even though the approximation method itself is statistically unbiased,
replacing the exact sparse operation with their faster approximated versions
still injects noise to the computation graph, and the noise will be propagated
layer-by-layer during the training. As we analyzed above, each SpMM is
executed twice in the training step. Below we first experimentally analyze the
impact of the injected noise in the forward pass and the backward pass. As
shown in Table 1, we apply top-$k$ sampling to approximate the SpMM in the
forward pass, backward pass, or both, respectively.

Table 1. Preliminary results on approximating SpMM via top-$k$ sampling. The model is a two-layer GCN and dataset is Reddit. Here we set the $k$ as $0.1|\mathcal{V}|$ across different layers. Method | Reddit
---|---
without approximation | 95.39±0.04
only forward | 16.45±0.39
only backward | 95.25±0.03
forward and backward | 80.74±1.00

From Table 1, the accuracy drop is negligible if we only replace SpMM in the
backward pass. Notably, if we apply approximation in both the forward and
backward pass, the result is significantly better than only applying top-$k$
sampling in the forward pass. The reason is that when only applying
approximation in the forward pass, some row/columns are not included in the
computation graph, so intuitively these row/columns should be excluded in the
backward pass. “forward and backward” result in Table 1 is built based on this
intuition such that in the backward pass, we use the column-row pairs sampled
in the forward pass to compute the gradient. However, it is still not
comparable to the result of applying approximation only in the backward pass.
Below we mathematically analyze the reason behind the results in Table 1.

####  3.1.2. Theoretical analysis

We first analyze the case of approximating the sparse operations in the
forward pass. Namely, replacing
$\texttt{SpMM}(\tilde{{\bm{A}}},{\bm{J}}^{(l)})$ with
$\texttt{approx}(\tilde{{\bm{A}}}{\bm{J}}^{(l)})$. We note that we have
$\mathbb{E}[f(x)]\neq f(\mathbb{E}[x])$ for any non-linear function
$f(\cdot)$, e.g., $\mathbb{E}[x^{2}]\neq\mathbb{E}^{2}[x]$. Thus, even when
the approximation method gives an unbiased estimation, i.e.,
$\mathbb{E}[\texttt{approx}(\tilde{{\bm{A}}}{\bm{J}}^{(l)})]=\tilde{{\bm{A}}}{\bm{J}}^{(l)}$,
the node embeddings ${\bm{H}}^{(l+1)}$ are still biased since the activation
function is non-linear. To see this,

$\displaystyle\mathbb{E}[{\bm{H}}^{(l+1)}]$
$\displaystyle=\mathbb{E}[\textrm{ReLU}(\texttt{approx}(\tilde{{\bm{A}}}{\bm{J}}^{(l)})])$
$\displaystyle\neq\textrm{ReLU}(\mathbb{E}[\texttt{approx}(\tilde{{\bm{A}}}{\bm{J}}^{(l)})])={\bm{H}}^{(l+1)}.$

Thus, if we apply the approximation for the SpMM in the forward pass, the bias
will be propagated layer-by-layer and cause significantly worse results. For
the case of only approximating the sparse operation in the backward pass, we
have the following proposition:

###### Proposition 0 (Proof in Appendix A).

If the approximation method is itself unbiased, and we only replace the SpMM
in the backward pass with its approximated version, while leaving the forward
one unchanged, then the calculated gradient is provably unbiased.

Due to the page limit, the proof is deferred to the Appendix. The high level
idea is that the gradient of activation function in backward pass is only
related to the pre-activations in the forward pass, and thus is independent to
the approximation error introduced in the backward pass. As suggested by our
theoretical and empirical analysis, as shown in Figure 2, we only approximate
the sparse operations in the backward pass, while leaving all other operations
unchanged. Namely, we replace the exact sparse operation
$\nabla{\bm{J}}^{(l)}=\texttt{SpMM}(\tilde{{\bm{A}}}^{\top},\nabla{\bm{H}}^{(l+1)})$
in the backward pass with
$\texttt{approx}(\tilde{{\bm{A}}}^{\top}\nabla{\bm{H}}^{(l+1)})$ using top-$k$
sampling.

###  3.2. How to Apply the Approximation

Figure 3. For approximated sparse operations, the acceleration is decided by
the selection of column-row pairs. For example, selecting orange column-row
pairs entails $\frac{3}{7}\times$ less FLOPs and memory access compared to the
blue pairs, since these orange low-degree nodes associate with less non-zero
entries in the adjacency matrix.

In the last section, we explored where to place the approximation method.
However, as we mentioned in the introduction, for sparse operations, the
acceleration is decided by the selection of sampled column-row pairs. To see
this, as shown in Figure 3, suppose we use top-$k$ sampling to approximate
$\texttt{SpMM}(\tilde{{\bm{A}}}^{\top},\nabla{\bm{H}})$. Since the
computations are only executed on the non-zero entries, so selecting the
orange pairs (i.e., pair 1 and 3) will result in $\frac{3}{7}\times$ less
computational cost (FLOPs) compared to selecting the blue pair (i.e., pair 0
and 2). For both the orange and blue cases, we have $k=2$. Thus, the number of
samples $k$ cannot directly constrain the FLOPs for each individual operation.
Moreover, a GNN has multiple layers (or operations) and each layer has a
different importance to the model accuracy. To optimize the accuracy-
efficiency trade off, our key idea is to customize the computation resources
(i.e., FLOPs) for each layer by adjusting the number of samples $k_{l}$ in the
$l$-th layer. In this way, we minimize the impact of approximation, while
limiting the overall FLOPs under a certain budget. Based on the idea, we frame
the resource allocation problem as the following constrained optimization
problem:

(5a)
$\displaystyle\min_{\left\\{k_{l}\right\\}}-\sum_{l=1}^{L}\sum_{i\in\mathrm{Top}_{k_{l}}}\|\tilde{\bm{A}}_{:,i}^{\top}\|_{2}\|\nabla{\bm{H}}^{(l+1)}_{i,:}\|_{2},$
(5b) $\displaystyle
s.t.\sum_{l=1}^{L}\sum_{i\in\mathrm{Top}_{k_{l}}}\\#nnz_{i}*d_{l}\leq
C\sum_{i=1}^{L}|\mathcal{E}|d_{l},$

where $C$ is the budget ($0<C<1$) that controls the overall reduced FLOPs.
$k_{l}$ is the number of samples for the top-$k$ sampling at the $l$-th layer.
$d_{l}$ is the hidden dimensions of $l$-th layer, and $\\#nnz_{i}$ is the
number of non-zero entries at the $i$-th column of $\tilde{{\bm{A}}}^{\top}$.
$\mathrm{Top}_{k_{l}}$ is the set of indices associated with the $k_{l}$
largest
$\|\tilde{\bm{A}}_{:,i}^{\top}\|_{2}\|\nabla{\bm{H}}^{(l+1)}_{i,:}\|_{2}$.

Equation (5a) is equivalent to minimizing approximation error
$\mathbb{E}[||\tilde{{\bm{A}}}^{\top}\nabla{\bm{H}}^{(l+1)}-\texttt{approx}(\tilde{{\bm{A}}}^{\top}\nabla{\bm{H}}^{(l+1)})||_{F}]$
summarized over all layers (Adelman et al., 2021). Also, different sparse
operations are weighted summation by the magnitude of gradient
$\|\nabla{\bm{H}}^{(l+1)}\|_{2}$, which implicitly encodes the importance of
different operations.

Equation (5b) is the constraint that controls the overall FLOPs. Specifically,
the FLOPs of SpMM between $\tilde{\bm{A}}$ and the gradient
$\nabla{\bm{H}}\in\mathbb{R}^{N\times d}$ is $\mathcal{O}(|\mathcal{E}|d)$ and
$\sum_{j\in\mathcal{V}}\\#nnz_{j}=|\mathcal{E}|$. Since for sparse operations,
the reduction in the number of memory access and FLOPs are linearly
correlated, Equation (5b) also bounds the number of memory access.

####  3.2.1. Greedy solution

The above combination optimization objective is NP-hard, albeit it can be
solved by dynamic programming. However, the dynamic programming is very slow,
which somehow contradicts with our purpose of being efficient. Thus, we
propose to use a greedy algorithm to solve it. Specifically, it starts with
highest $k_{l}=|\mathcal{V}|$ for all layers. In each move, it chooses a
$k_{l}$ among $\\{k_{l}\\}_{l=1}^{L}$ to reduce by a step size (e.g.,
$0.02|\mathcal{V}|$), such that the increment of errors in Equation (5a) is
minimal. The greedy algorithm will stop when the current total FLOPs fits in
the budget in Equation (5b). This algorithm runs super fast and we found that
it has minimal impact to the efficiency. We provide the pseudo code of our
greedy algorithm in Algorithm 1 of Appendix B.

Figure 4.  For each layer, the selected column-row pairs tend to be very
similar across iterations. Models here are two-layer GCN and GraphSAGE. Here
we show the matching scores (AUC) of top-$k$ indices between every 10 steps.
Figure 5. The process of slicing the sparse matrix in Figure 3 by only
reserving orange columns (in CSR format), which requires to re-build the
Rowptr and Col Array.

###  3.3. When to Apply the Approximation

####  3.3.1. Cache the sampled sparse matrices

Before illustrating the caching mechanism, we first give the details about the
Compressed Sparse Row (CSR) format for representing the sparse matrix here.
CSR stores nonzero values in a matrix and their position in three arrays:
index array Rowptr, column array Col, and value array Val. The elements in
Rowptr act as the starting indices of the elements in Col and Val that
correspond to each row. Specifically, the elements of row $i$ are stored in
indices Rowptr[i] to Rowptr[i+ 1] $-$ 1 of Col and Val . The elements in Col
and Val are the column index and value in that column, respectively. Figure 5
shows the CSR format of the matrix shown in Figure 3. We ignore the Val array
here since the sparse matrix contains only element one.

As shown in Figure 2, RSC reduces the number of FLOPs and memory access via
top-$k$ sampling. From the implementation aspect, executing the top-$k$
sampling contains two steps: First, it decides the indices corresponding to
the top-$k$ largest column row norms in Equation (4). Second, slicing the
matrices according to the indices. In practice, the overhead of calculating
the column-row norms can be ignored. However, unlike dense matrices, slicing
the adjacency matrix is very slow due to its irregular data format. To see
this, suppose the top-$k$ indices of the sparse matrix in Figure 3 correspond
to the orange column-row pairs. We show in Figure 5 the process of slicing the
adjacency matrix in CSR format by reserving only the orange columns (Fey and
Lenssen, 2019). Compared to slicing dense matrices, slicing sparse matrices
requires to re-process the graph to build the new Rowptr and Col, which
introduces significant time overhead, especially for large graphs.

For the full graph training, we use the same adjacency matrix across different
epochs111For sub-graph based training, we can first sample all of the sub-
graphs offline. Then during the training, we apply the caching mechanism to
each sampled graph.. We made a crucial observation that the top-$k$ indices in
the adjacency matrix tend to be the same across iterations. In Figure 4, we
plot the AUC score of top-$k$ indices between every iteration $t$ and
iteration $t+10$ for each layer throughout the whole training process. Here we
note that AUC score is a commonly used ranking measure and a 1.0 AUC score
means the ranking of column-row pairs is identical across iterations. The
results in Figure 4 indicate that the top-$k$ indices won’t change
significantly within a few iterations. Thus, as shown in Figure 2, we propose
to reuse the sampled adjacency matrix for each layer across nearby iterations
for reducing the overhead. In Section 5, we experimentally show that the
caching mechanism does not impact the model performance a lot, but leading to
significant speedup.

####  3.3.2. Switch back to origin sparse operation at the end

When training neural networks, the common practice is to use a large learning
rate for exploration and anneal to a small one for final convergence (Li et
al., 2019). The rationale behind this strategy is that, at the end of the
training process, we need to fine-tune our model with small noise for
convergence. Since our approximation sparse operations will bring extra noise
to the gradient, intuitively, we can switch back to the original sparse
operations to help convergence. More formally, we propose to use approximated
sparse operation during most of the training process, while switching back to
the original sparse operation at the final stage. We experimentally show that
this switching mechanism significantly reduces the accuracy drop, at the cost
of slightly less acceleration effect.

###  3.4. Connections to Graph Data Augmentation

Data augmentation (Liu et al., 2021; Han et al., 2022) is wildly adopted in
the graph learning for improving model generalization, including dropping
nodes (Feng et al., 2020), dropping edges (Rong et al., 2019), and graph mixup
(Han et al., 2022). As shown in Figure 5, the top-$k$ sampling drops the
entire columns in the adjacency matrix, while keeping the number of rows
unchanged.  That means RSC drops all of the out edges for a set of nodes. This
can be viewed as the “structural dropedge” for improving the efficiency. Since
we only apply the top-$k$ sampling in the backward pass and top-$k$ indices
are different for each operation,  RSC essentially forward pass with the whole
graph, backward pass with different subgraphs at each layer. This structural
dropedge and heterogeneous backward propagation introduce the regularization
effect. Thus as shown in the experiment section, RSC may also improve the
model accuracy over the baseline.

##  4\. Related work and Discussion

In this section, we first discuss the related work. Then we discuss the
limitations of our work.

Subgraph-based GNN training. The key idea of this line of work is to improve
the scalability of GNNs by separating the graph into overlapped small batches,
then training models with sampled subgraphs (Hamilton et al., 2017; Huang et
al., 2018; Zou et al., 2019; Chiang et al., 2019; Zeng et al., 2020). However,
this approach reduces the memory footprint but results in extra time cost to
compute the overlapping nodes between batches. Generally, methods in this
category are orthogonal to RSC , and they can be combined, which is shown in
the next section.

Approximated Matrix Multiplication. The known algorithms for approximating
matrix production can be roughly divided into three categories. However, only
a few of them can be used for accelerating GNN training.

  * •

Random walk-based methods (Cohen and Lewis, 1999) performs random walks on a
graph representation of the dense matrices, but is applicable to non-negative
matrices only.

  * •

Butterfly-based methods (Chen et al., 2021; Dao et al., 2022) replace dense
matrices with butterfly matrices. It is not applicable to SpMM in GNNs because
the adjacency matrix often cannot be reduced to a butterfly matrix.

  * •

Column-row sampling methods(Drineas et al., 2006b; Drineas and Kannan, 2001)
subsample the input matrices with important rows and columns, then perform the
production on sampled matrix as usual.

Other randomized GNN training. Dropedge (Rong et al., 2019) randomly drops
edges to avoid the over-smoothing problem. Graph Random Neural Networks
(Grand) (Feng et al., 2020) randomly drop nodes to generate data augmentation
for improving model generalization. Grand+ improves the scalability over Grand
by pre-computing a general propagation matrix and employ it to perform data
augmentation (Feng et al., 2022). As shown in Section 3.4, the key difference
between GRAND(+) and RSC is that RSC does not drop any node. Instead RSC drops
all of the out edges for a set of nodes only during backward pass. Moreover,
the drop pattern are evolving during the training process. This can be viewed
as the “structural dropedge”. However, unlike Dropedge (Rong et al., 2019),
RSC drop the column-row pairs according to the euclidean norm instead of
uniformly dropping.

Limitations of RSC . First, to guarantee the model accuracy, we only replace
the sparse operation in the backward pass. Thus the upper bound of RSC ’s
speedup is limited. However, we note that the backward pass usually is more
time-consuming than the forward pass, which is also empirically shown in Table
2. Second, some GNNs reply on the scatter-and-gather instead of SpMM (and its
variant) to perform the aggregation, such as GAT (Veličković et al., 2017).
They are not covered in this paper. However, scatter-and-gather based GNNs can
also be accelerated by RSC because the column-row sampling is also applicable
to scatter and gather operation. Similarly, the caching and switching
mechanisms are also applicable to them. However, for the resource allocation
Algorithm 1, the scatter and gather operations requires tailored error bound
and the computation cost modeling in Equation (5). We leave it as the future
work.

##  5\. Experiments

We verify the effectiveness of our proposed framework through answering the
following research questions:

  * •

Q1: How effective is RSC in terms of the accuracy with reduced training time?

  * •

Q2: How effective is our proposed allocation strategy compared to the uniform
allocation strategy?

  * •

Q3: What is the layer-wise ratio assigned by RSC ?

  * •

Q4: How effective is the caching and switching mechanism in terms of the
trade-off between efficiency and accuracy?

  * •

Q5: How sensitive is RSC to its key hyperparameters?

Table 2. Comparison on the efficiency at the operation level. fwd is the wall-clock time for a single forward pass (ms), bwd is the wall-clock time for a single backward pass (ms). SpMM_MEAN is one variant of SpMM , which corresponds to the MEAN aggregator used in GraphSAGE (Appendix A.2). RSC can reduce the wall clock time of backward sparse computation by up to $11.6\times$. |  | Reddit | Yelp |  | ogbn-
---
proteins
| ogbn-
---
products
|  | fwd | bwd | fwd | bwd | fwd | bwd | fwd | bwd
SpMM | Baseline | 36.28 | 44.23 | 26.88 | 34.38 | 31.72 | 42.99 | 261.03 | 316.80
+RSC | - | 3.81 (11.6$\times$) | - |  9.86 (3.49$\times$) | - | 14.87 (2.89$\times$) | - | 35.28 (8.98$\times$)
SpMM_MEAN (Appendix A.2) | Baseline | 36.21 | 44.27 | 26.78 | 34.38 | 31.80 | 43.11 | 261.03 | 316.84
+RSC | - |  7.47 (5.92$\times$) | - | 19.62 (1.75$\times$) | - | 5.22 (8.26$\times$) | - | 71.59 (4.43$\times$)
Table 3. Comparison on the test accuracy/F1-micro/AUC and speedup on four
datasets. Bold faces indicate the accuracy drop is negligible ($\approx
0.3\%$) or the result is better compared to the baseline.The hardware here is
a RTX3090 (24GB). All reported results are averaged over ten random trials.

| # nodes
---
# edges
| 230K
---
11.6M
| 717K
---
7.9M
| 132K
---
39.5M
| 2.4M
---
61.9M
Model | Methods | Reddit | Yelp |  | ogbn-
---
proteins
| ogbn-
---
products
Acc. |  | Budget
---
$C$
Speedup | F1-micro |  | Budget
---
$C$
Speedup | AUC |  | Budget
---
$C$
Speedup | Acc. |  | Budget
---
$C$
Speedup
Graph- SAINT | Baseline | 96.40±0.03 | 1 | 1$\times$ | 63.30±0.14 | 1 | 1$\times$ | — | — | — | 79.01±0.21 | 1 | 1$\times$
+RSC | 96.24±0.03 | 0.1 | 1.11$\times$ | 63.34±0.18 | 0.1 | 1.09$\times$ | — | — | — | 78.99±0.32 | 0.3 | 1.04$\times$
GCN | Baseline | 95.33±0.03 | 1 | 1$\times$ | 44.28±1.04 | 1 | 1$\times$ | 71.99±0.66 | 1 | 1$\times$ | 75.74±0.11 | 1 | 1$\times$
+RSC | 95.13±0.05 | 0.1 | 1.47$\times$ | 46.09±0.54 | 0.1 | 1.17$\times$ | 71.60±0.45 | 0.3 | 1.51$\times$ | 75.44±0.21 | 0.3 | 1.35$\times$
GraphSAGE (full batch) | Baseline | 96.61±0.05 | 1 | 1$\times$ | 63.06±0.18 | 1 | 1$\times$ | 76.09±0.77 | 1 | 1$\times$ | 78.73 ± 0.12 | 1 | 1$\times$
+RSC | 96.52±0.04 | 0.1 | 1.32$\times$ | 62.89±0.19 | 0.1 | 1.13$\times$ | 76.30±0.42 | 0.3 | 1.60$\times$ | 78.50± 0.09 | 0.1 | 1.53$\times$
GCNII | Baseline | 96.71±0.07 | 1 | 1$\times$ | 63.45±0.17 | 1 | 1$\times$ | 73.79±1.32 | 1 | 1$\times$ | — | — | —
+RSC | 96.50±0.12 | 0.3 | 1.45$\times$ | 63.57±0.21 | 0.1 | 1.19$\times$ | 75.20±0.54 | 0.5 | 1.41$\times$ | — | — | —

###  5.1. Experimental Settings

Datasets and Baselines. To evaluate RSC , we adopt four common large scale
graph benchmarks from different domains, i.e., Reddit (Hamilton et al., 2017),
Yelp (Zeng et al., 2020), ogbn-proteins (Hu et al., 2020), and ogbn-products
(Hu et al., 2020). We evaluate RSC under both the mini-batch training and
full-batch training settings. For the mini-batch training setting, we
integrate RSC with one of the state-of-the-art sampling methods, GraphSAINT
(Zeng et al., 2020). For the full-batch training setting, we integrate RSC
with three popular models: two commonly used shallow models, namely, GCN (Kipf
and Welling, 2017) and GraphSAGE (Hamilton et al., 2017), and one deep model
GCNII (Chen et al., 2020). To avoid creating confusion, GCN, GraphSAGE, and
GCNII are all trained with the whole graph at each step. For a fair
comparison, we use the MEAN aggregator for GraphSAGE and GraphSAINT throughout
the paper. Details about the hyperparameters and datasets are in Appendix C.

Hyperparameter settings. RSC contains three parts. First, the allocation
strategy. We choose the overall budget $C$ in Equation (5b) from
$\\{0.1,0.3,0.5\\}$. We run the resource allocation strategy every ten steps.
The step size $\alpha$ in Algorithm 1 is set as $0.02|\mathcal{V}|$. Second,
the caching mechanism. According to Figure 4, we sample the adjacency matrix
every 10 steps and reuse the sampled sparse matrix for nearby steps. Third,
the switching mechanism, where we apply RSC for 80% of the total epochs, while
switching back to the original operations for the rest of the 20% epochs.

Evaluation metrics. To evaluate the practical usage of RSC , we report the
wall clock time speedup measured on GPUs. Specifically, the speedup equals
$T_{\text{baseline}}/T_{rsc}$, where $T_{\text{baseline}}$ and $T_{rsc}$ are
the wall clock training time of baseline and RSC , respectively. We note that
the $T_{rsc}$ includes the running time of the greedy algorithm in Section
3.2.1, as well as the effects of caching and switching.

###  5.2. Performance Analysis

####  5.2.1. Accuracy-efficiency trade-off

To answer Q1, we summarize the speedup and the test accuracy/F1-micro/AUC of
different methods in Table 3. Since RSC accelerates the sparse operation in
the backward pass, we also provide the detailed efficiency analysis in Table
2. In summary, we observe:

Figure 6. The Pareto frontier of the accuracy-efficiency trade-off for RSC and
the uniform allocation. Here we disabled the caching and switch mechanism for
a fair comparison. All results are averaged over five random trials.

Figure 7. The allocated layer-wise $k_{l}$ for GCN, GraphSAGE and GCNII on
Reddit, where the overall budget $C$ is set as 0.1. The input of the SpMM in
the first GraphSAGE layer does not require gradient and thus absent in the
Figure (Appendix A.2).

Figure 8. Hyperparameter analysis w.r.t. the budget $C$, the step size
$\alpha$ in Algorithm 1, and when switching back to the original operations.
The model is GraphSAGE and the applied dataset is ogbn-proteins.

❶ At the operation level, RSC can accelerate the sparse operation in the
backward pass by up to $11.6\times$. For end-to-end training, the accuracy
drop of applying RSC over baselines is negligible (0.3%) across different
models and datasets, while achieving up to $1.6\times$ end-to-end wall clock
time speedup. The gap between the operation speedup and the end-to-end speedup
are due to the following two reasons. First, we focus on accelerating the
sparse computations in GNNs, which is the unique bottleneck to GNNs. The other
dense computations can certainly be accelerated by approximation methods, but
this beyond the scope of this paper. Second, according to our theoretical
analysis, we only accelerate the sparse computation in the backward pass
instead of the forward one to guarantee the performance. We note that for
approximation methods that accelerate the training process at operation level,
a $1.2\approx 1.3\times$ wall-clock speedup with negligible accuracy drop can
be regarded as non-trivial (for details, please see Table 1 in (Adelman et
al., 2021)), especially considering that these approximation methods are
orthogonal to most of the existing efficient training methods. For GraphSAINT,
the speedup of RSC is around $1.1\times$, which is smaller than the full graph
training. This is because for subgraph-based training, the equivalent “batch
size” is much smaller than the full graph counterparts. As a result, the GPU
utility is low since it does not assign each processor sufficient amount of
work (You et al., 2019). The speedup for subgraph training is expected to
improve when using larger batch size.

❷ The proper hyperparameter $C$ varies with datasets and models. As shown in
Table 3, the overall budget $C$ varies across dataset and models. For example,
for ogbn-proteins, the overall budget $C$ is 0.3 for GCN and GraphSAGE, while
it is 0.5 for GCNII. For Reddit, the overall budget $C$ is 0.1 for GCN and
GraphSAGE, while it is 0.3 for GCNII. In general, deep GNNs may require a
larger $C$ compared to the shallow GNNs. This is because deeper GNNs are more
sensitive to the approximation errors due to the fact that the approximation
error will compound layer-by-layer during the backward pass (Liu et al.,
2022).

####  5.2.2. The role of resource allocation strategy.

Each sparse operation has a different importance to the model accuracy. Thus
when allocating the same $k_{l}$ to different sparse operations, it may either
harm the model performance with insufficient $k_{l}$ or waste resources on
unimportant operations. To answer Q2, we compare RSC with the uniform
allocation strategy, i.e., setting $k_{l}=C|\mathcal{V}|$ for all sparse
operations in the backward pass. As shown in Figure 6, we plot the Pareto
frontier of the accuracy-efficiency trade-off on Reddit dataset for RSC and
the uniform strategy with different $C$. For a fair comparison, we disabled
the caching and switching mechanism. Due to the page limit, we show the
running time of the greedy algorithm in Table 5 of Appendix B. We conclude
that the overhead of greedy algorithm is negligible compared to the
acceleration effect of RSC . In summary, we observe that:

❸ RSC exhibits a superior trade-off between accuracy and efficiency compared
to the uniform allocation, especially under high speedup regime. Namely,
compared to the uniform allocation, RSC can achieve higher model accuracy
under the same speedup. This can be explained by the fact that each operation
has a different importance to the model performance. RSC can automatically
allocate more resource to important operations under a given total budget.

To answer Q3, we also visualize the allocated $k_{l}$ for each layer across
iterations in Figure 7 (as mentioned in Section 5.1, the resource allocation
algorithm is executed every ten steps). We note that the input of the SpMM in
the first GraphSAGE layer does not require gradient and thus does not appear
in Figure 7, which we will elaborate more details in Appendix A.2.

❹ The $k_{l}$ assigned by RSC evolves along with the training process. From
Figure 7, the allocated $k_{l}$ changes rapidly in the initial phase and
starts to converge at about $100^{\rm{th}}$ epochs. Due to the page limit, we
also visualize the degree of picked nodes in Figure 9 in the Appendix. We
conclude that the preference of different layers also evolves during the
training process.

Table 4. Ablation on the caching and switching mechanism. Experiments are conducted on ogbn-proteins. All results are averaged over five random trials. Ablation on | Caching | Switching | AUC | Speedup
---|---|---|---|---
GCN | ✗ | ✗ |  71.60 ± 0.66 |  1.19$\times$
✗ | ✓ |  72.19 ± 0.79 |  1.14$\times$
✓ | ✗ |  69.80 ± 0.60 |  1.60$\times$
✓ | ✓ |  71.60 ± 0.45 |  1.51$\times$
GraphSAGE | ✗ | ✗ |  75.23 ± 0.79 |  1.37$\times$
✗ | ✓ |  76.39 ± 0.39 |  1.32$\times$
✓ | ✗ |  75.53 ± 0.60 |  1.78$\times$
✓ | ✓ |  76.30 ± 0.42 |  1.60$\times$
GCNII | ✗ | ✗ |  74.07 ± 0.83 |  1.10$\times$
✗ | ✓ |  74.50 ± 0.52 |  1.04$\times$
✓ | ✗ |  72.47 ± 0.75 |  1.46$\times$
✓ | ✓ |  75.20 ± 0.54 |  1.41$\times$

####  5.2.3. The role of caching and switching mechanism.

In section 5.2.2, we have shown the superior results of the proposed resource
allocation strategy. As we mentioned in Section 3.3, we also introduce two
simple tricks to for improving RSC , i.e., the caching and switching
mechanism. To verify the effect of each of them (Q4), we conduct incremental
evaluations on GCN, GraphSAGE and GCNII with ogbn-proteins, which are
summarized in Table 4. The row without caching and switching in Table 4
corresponds to the results with proposed resource allocation strategy. We
observe that:

❺ Switching mechanism significantly improves the model performance, at the
cost of slightly less acceleration effect. Specifically, compared to the
results with only resource allocation strategy, the switching mechanism
improves 0.59%, 1.16%, and 0.43% AUC, at the cost of $0.05\times$,
$0.05\times$, and $0.06\times$ less speedup for GCN, GraphSAGE, and GCNII,
respectively. As we analyzed in Section 3.3.2, the improvement can be
explained by the fact that the final training stage requires smaller gradient
noise to help convergence.

❻ Caching mechanism significantly improves the wall-clock time speedup, at the
cost of worse model performance. Specifically, compared to the results with
only resource allocation strategy, the caching mechanism improves $34.4\%$,
$29.9\%$, and $32.7\%$ speedup, at the cost of $1.8\%$, $-0.3\%$, and $1.6\%$
less AUC for GCN, GraphSAGE, and GCNII, respectively. Although caching
mechanism can significantly reduce the overhead of sampling, the performance
drop is too large ($>1\%$). Intuitively, the accuracy drop of caching also
implies that we could not use a “static” down-sampled graph throughout the
training process.

❼ Surprisingly, jointly applying the caching and switching, the performance
drop can be minimized. Specifically, compared to the results with only
resource allocation strategy, when jointly applying the caching and switching,
it improves $0.0\%$, $1.07\%$, and $1.13\%$ model performance, while achieving
$0.22\times$, $0.23\times$, and $0.2\times$ more speedup simultaneously. In
summary, jointly applying the caching and switching enables a “win-win” in
terms of the model performance and training efficiency.

####  5.2.4. Hyperparameter sensitivity analysis (Q5).

Here we analyze the impacts of the main hyperparameters of RSC : (1) the
budget $C$, which controls the efficiency-accuracy trade off; (2) the step
size $\alpha$ in the greedy Algorithm 1; (3) when switching back to the
original sparse operations. In Figure 8, we vary only one of them with the
others fixed. We conclude (1) larger budget $C$ leads to better accuracy with
smaller speedup, since we are using more computational resources to
approximate the full operation. (2) larger step size $\alpha$ leads to
marginally larger speedup since the greedy algorithm will terminate earlier.
Also the step size $\alpha$ does not affect the model accuracy a lot. In
practice, we set $\alpha=0.02|\mathcal{V}|$. (3) The later we switch back to
the original operation, the larger the accuracy drop and the smaller the
speedup, it is equivalent to using less resources to approximate the full
operation epoch-wisely. Thus, we apply RSC for $80\%$ of the total epochs to
balance the trade-off.

##  6\. Conclusions and Future work

In this paper, we propose RSC , which replaces the sparse computations in GNNs
with their fast approximated versions. RSC can be plugged into most of the
existing training frameworks to improve their efficiency. Future work includes
(1) exploring other types of approximation methods. (2) exploring RSC for GNNs
that relies on scatter-and-gather operations.

## References

  * (1)
  * Adelman et al. (2021) Menachem Adelman, Kfir Levy, Ido Hakimi, and Mark Silberstein. 2021.  Faster neural network training with approximate tensor operations.  _Advances in Neural Information Processing Systems_ 34 (2021), 27877–27889.
  * Chen et al. (2021) Beidi Chen, Tri Dao, Kaizhao Liang, Jiaming Yang, Zhao Song, Atri Rudra, and Christopher Re. 2021.  Pixelated butterfly: Simple and efficient sparse training for neural network models.  _arXiv preprint arXiv:2112.00029_ (2021).
  * Chen et al. (2020) Ming Chen, Zhewei Wei, Zengfeng Huang, Bolin Ding, and Yaliang Li. 2020.  Simple and deep graph convolutional networks. In _International Conference on Machine Learning_. PMLR, 1725–1735.
  * Chiang et al. (2019) Wei-Lin Chiang, Xuanqing Liu, Si Si, Yang Li, Samy Bengio, and Cho-Jui Hsieh. 2019\.  Cluster-gcn: An efficient algorithm for training deep and large graph convolutional networks. In _Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining_. 257–266.
  * Cohen and Lewis (1999) Edith Cohen and David D Lewis. 1999.  Approximating matrix multiplication for pattern recognition tasks.  _Journal of Algorithms_ 30, 2 (1999), 211–252.
  * Cong et al. (2020) Weilin Cong, Rana Forsati, Mahmut Kandemir, and Mehrdad Mahdavi. 2020.  Minimal variance sampling with provable guarantees for fast training of graph neural networks. In _Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining_. 1393–1403.
  * Dao et al. (2022) Tri Dao, Beidi Chen, Nimit S Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré. 2022\.  Monarch: Expressive structured matrices for efficient and accurate training. In _International Conference on Machine Learning_. PMLR, 4690–4721.
  * Drineas and Kannan (2001) Petros Drineas and Ravi Kannan. 2001.  Fast Monte-Carlo algorithms for approximate matrix multiplication. In _Proceedings 42nd IEEE Symposium on Foundations of Computer Science_. IEEE, 452–459.
  * Drineas et al. (2006a) Petros Drineas, Ravi Kannan, and Michael W Mahoney. 2006a.  Fast Monte Carlo algorithms for matrices I: Approximating matrix multiplication.  _SIAM J. Comput._ 36, 1 (2006), 132–157.
  * Drineas et al. (2006b) Petros Drineas, Ravi Kannan, and Michael W Mahoney. 2006b.  Fast Monte Carlo algorithms for matrices I: Approximating matrix multiplication.  _SIAM J. Comput._ 36, 1 (2006), 132–157.
  * Duan et al. (2022a) Keyu Duan, Zirui Liu, Peihao Wang, Wenqing Zheng, Kaixiong Zhou, Tianlong Chen, Xia Hu, and Zhangyang Wang. 2022a.  A Comprehensive Study on Large-Scale Graph Training: Benchmarking and Rethinking.  (2022).
  * Duan et al. (2022b) Keyu Duan, Zirui Liu, Peihao Wang, Wenqing Zheng, Kaixiong Zhou, Tianlong Chen, Xia Hu, and Zhangyang Wang. 2022b.  A Comprehensive Study on Large-Scale Graph Training: Benchmarking and Rethinking. In _Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track_.  https://openreview.net/forum?id=2QrFr_U782Z
  * Feng et al. (2022) Wenzheng Feng, Yuxiao Dong, Tinglin Huang, Ziqi Yin, Xu Cheng, Evgeny Kharlamov, and Jie Tang. 2022.  GRAND+: Scalable Graph Random Neural Networks. In _Proceedings of the ACM Web Conference 2022_. 3248–3258.
  * Feng et al. (2020) Wenzheng Feng, Jie Zhang, Yuxiao Dong, Yu Han, Huanbo Luan, Qian Xu, Qiang Yang, Evgeny Kharlamov, and Jie Tang. 2020.  Graph random neural networks for semi-supervised learning on graphs.  _Advances in neural information processing systems_ 33 (2020), 22092–22103.
  * Fey and Lenssen (2019) Matthias Fey and Jan E. Lenssen. 2019.  Fast Graph Representation Learning with PyTorch Geometric. In _ICLR Workshop on Representation Learning on Graphs and Manifolds_.
  * Glorot and Bengio (2010) Xavier Glorot and Yoshua Bengio. 2010.  Understanding the difficulty of training deep feedforward neural networks. In _Proceedings of the thirteenth international conference on artificial intelligence and statistics_. JMLR Workshop and Conference Proceedings, 249–256.
  * Hamilton et al. (2017) William L Hamilton, Rex Ying, and Jure Leskovec. 2017\.  Inductive representation learning on large graphs. In _Proceedings of the 31st International Conference on Neural Information Processing Systems_. 1025–1035.
  * Han et al. (2016) Song Han, Xingyu Liu, Huizi Mao, Jing Pu, Ardavan Pedram, Mark A Horowitz, and William J Dally. 2016.  EIE: Efficient inference engine on compressed deep neural network.  _ACM SIGARCH Computer Architecture News_ 44, 3 (2016), 243–254.
  * Han et al. (2015) Song Han, Jeff Pool, John Tran, and William Dally. 2015\.  Learning both weights and connections for efficient neural network.  _Advances in neural information processing systems_ 28 (2015).
  * Han et al. (2022) Xiaotian Han, Zhimeng Jiang, Ninghao Liu, and Xia Hu. 2022\.  G-Mixup: Graph Data Augmentation for Graph Classification.  _arXiv preprint arXiv:2202.07179_ (2022).
  * Hu et al. (2020) Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. 2020\.  Open graph benchmark: Datasets for machine learning on graphs.  _arXiv preprint arXiv:2005.00687_ (2020).
  * Huang et al. (2020) Guyue Huang, Guohao Dai, Yu Wang, and Huazhong Yang. 2020\.  Ge-spmm: General-purpose sparse matrix-matrix multiplication on gpus for graph neural networks. In _SC20: International Conference for High Performance Computing, Networking, Storage and Analysis_. IEEE, 1–12.
  * Huang et al. (2018) Wenbing Huang, Tong Zhang, Yu Rong, and Junzhou Huang. 2018\.  Adaptive sampling towards fast graph representation learning. In _Advances in Neural Information Processing Systems_.
  * Jiang et al. (2022) Zhimeng Jiang, Xiaotian Han, Chao Fan, Zirui Liu, Na Zou, Ali Mostafavi, and Xia Hu. 2022.  Fmp: Toward fair graph message passing against topology bias.  _arXiv preprint arXiv:2202.04187_ (2022).
  * Kipf and Welling (2017) Thomas N Kipf and Max Welling. 2017.  Semi-supervised classification with graph convolutional networks. In _International Conference on Learning Representations_.  https://openreview.net/forum?id=SJU4ayYgl
  * Li et al. (2019) Yuanzhi Li, Colin Wei, and Tengyu Ma. 2019.  Towards explaining the regularization effect of initial large learning rate in training neural networks.  _Advances in Neural Information Processing Systems_ 32 (2019).
  * Liu et al. (2021) Zirui Liu, Haifeng Jin, Ting-Hsiang Wang, Kaixiong Zhou, and Xia Hu. 2021.  DivAug: Plug-in Automated Data Augmentation with Explicit Diversity Maximization. In _Proceedings of the IEEE/CVF International Conference on Computer Vision_. 4762–4770.
  * Liu et al. (2022) Zirui Liu, Kaixiong Zhou, Fan Yang, Li Li, Rui Chen, and Xia Hu. 2022.  EXACT: Scalable Graph Neural Networks Training via Extreme Activation Compression. In _International Conference on Learning Representations_.  https://openreview.net/forum?id=vkaMaq95_rX
  * Narayanan et al. (2022) S Deepak Narayanan, Aditya Sinha, Prateek Jain, Purushottam Kar, and SUNDARARAJAN SELLAMANICKAM. 2022\.  IGLU: Efficient GCN Training via Lazy Updates. In _International Conference on Learning Representations_.  https://openreview.net/forum?id=5kq11Tl1z4
  * Rahman et al. (2021) Md Khaledur Rahman, Majedul Haque Sujon, and Ariful Azad. 2021\.  Fusedmm: A unified sddmm-spmm kernel for graph embedding and graph neural networks. In _2021 IEEE International Parallel and Distributed Processing Symposium (IPDPS)_. IEEE, 256–266.
  * Ramezani et al. (2022) Morteza Ramezani, Weilin Cong, Mehrdad Mahdavi, Mahmut Kandemir, and Anand Sivasubramaniam. 2022\.  Learn Locally, Correct Globally: A Distributed Algorithm for Training Graph Neural Networks. In _International Conference on Learning Representations_.  https://openreview.net/forum?id=FndDxSz3LxQ
  * Rong et al. (2019) Yu Rong, Wenbing Huang, Tingyang Xu, and Junzhou Huang. 2019\.  Dropedge: Towards deep graph convolutional networks on node classification.  _arXiv preprint arXiv:1907.10903_ (2019).
  * Veličković et al. (2017) Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. 2017.  Graph attention networks. In _International Conference on Learning Representations_.
  * Wan et al. (2022) Cheng Wan, Youjie Li, Cameron R Wolfe, Anastasios Kyrillidis, Nam Sung Kim, and Yingyan Lin. 2022\.  Pipegcn: Efficient full-graph training of graph convolutional networks with pipelined feature communication.  _arXiv preprint arXiv:2203.10428_ (2022).
  * Wang et al. (2019) Minjie Wang, Da Zheng, Zihao Ye, Quan Gan, Mufei Li, Xiang Song, Jinjing Zhou, Chao Ma, Lingfan Yu, Yu Gai, Tianjun Xiao, Tong He, George Karypis, Jinyang Li, and Zheng Zhang. 2019.  Deep Graph Library: A Graph-Centric, Highly-Performant Package for Graph Neural Networks.  _arXiv preprint arXiv:1909.01315_ (2019).
  * Ying et al. (2018) Rex Ying, Ruining He, Kaifeng Chen, Pong Eksombatchai, William L Hamilton, and Jure Leskovec. 2018\.  Graph convolutional neural networks for web-scale recommender systems. In _Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining_. 974–983.
  * You et al. (2019) Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. 2019\.  Large batch optimization for deep learning: Training bert in 76 minutes.  _arXiv preprint arXiv:1904.00962_ (2019).
  * Zeng et al. (2020) Hanqing Zeng, Hongkuan Zhou, Ajitesh Srivastava, Rajgopal Kannan, and Viktor Prasanna. 2020\.  GraphSAINT: Graph Sampling Based Inductive Learning Method. In _International Conference on Learning Representations_.  https://openreview.net/forum?id=BJe8pkHFwS
  * Zhang et al. (2022) Hengrui Zhang, Zhongming Yu, Guohao Dai, Guyue Huang, Yufei Ding, Yuan Xie, and Yu Wang. 2022.  Understanding gnn computational graph: A coordinated computation, io, and memory perspective.  _Proceedings of Machine Learning and Systems_ 4 (2022), 467–484.
  * Zheng et al. (2020) Da Zheng, Chao Ma, Minjie Wang, Jinjing Zhou, Qidong Su, Xiang Song, Quan Gan, Zheng Zhang, and George Karypis. 2020.  Distdgl: distributed graph neural network training for billion-scale graphs. In _2020 IEEE/ACM 10th Workshop on Irregular Applications: Architectures and Algorithms (IA3)_. IEEE, 36–44.
  * Zou et al. (2019) Difan Zou, Ziniu Hu, Yewen Wang, Song Jiang, Yizhou Sun, and Quanquan Gu. 2019\.  Layer-dependent importance sampling for training deep and large graph convolutional networks.  _arXiv preprint arXiv:1911.07323_ (2019).

##  Appendix A Mathematical Analysis

###  A.1. Proof of Proposition 1

See 3.1

Here we note that in the main text, for the notation convenience, we ignore
the backward pass of ReLU. However, the proof here will consider the non-
linear activation function to prove the unbiasedness. Let
${\bm{H}}^{(l+1)}_{pre}=\texttt{SpMM}(\tilde{{\bm{A}}},{\bm{J}}^{(l)})$ be the
pre-activation. The backward pass of ReLU is:

$\displaystyle\mathbb{E}[\nabla{\bm{H}}^{(l+1)}_{pre}]$
$\displaystyle=\mathbb{E}[\mathbbm{1}_{{\bm{H}}^{(l+1)}_{pre}>0}\odot\nabla{\bm{H}}^{(l+1)}]$
(6)
$\displaystyle=\mathbbm{1}_{{\bm{H}}^{(l+1)}_{pre}>0}\odot\mathbb{E}[\nabla{\bm{H}}^{(l+1)}],$

where $\odot$ is the element-wise product and $\mathbbm{1}$ is the indicator
function. The element-wise product is linear operation and
$\mathbbm{1}_{{\bm{H}}^{(l+1)}_{pre}>0}$ is only related to the pre-activation
in the forward pass, we only apply the approximation during the backward pass
so $\mathbbm{1}_{{\bm{H}}^{(l+1)}_{pre}>0}$ can be extracted from the
expectation. We know that for the last layer, we have
$\mathbb{E}[\nabla{\bm{H}}^{(L)}]={\bm{H}}^{(L)}$ since we do not apply ReLU
at the output layer. We then can prove by induction that
$\mathbb{E}[\nabla{\bm{H}}^{(l+1)}]={\bm{H}}^{(l+1)}$ and
$\mathbb{E}[\nabla{\bm{J}}^{(l)}]=\mathbb{E}[\texttt{approx}(\tilde{{\bm{A}}}^{\top}\nabla{\bm{H}}^{(l+1)}_{pre})]=\nabla{\bm{J}}^{(l)}$
for any layer $l$.

###  A.2. Analysis of MEAN aggregator

For GraphSAGE, one commonly used aggregator is the MEAN aggregator, which can
be expressed as follows:

(7)
${\bm{H}}^{(l+1)}={\bm{W}}_{1}{\bm{H}}^{(l)}+{\bm{W}}_{2}\texttt{SpMM\\_MEAN}({\bm{A}},{\bm{H}}^{(l)}),$

where SpMM_MEAN is one variant of the vanilla SpMM , which replace the reducer
function from sum($\cdot$) to mean($\cdot$). We note that in popular GNN
packages, the MEAN aggregator usually is implemented based on SpMM_MEAN (Fey
and Lenssen, 2019; Wang et al., 2019) to reduce the memory usage. Here we give
an example of SpMM_MEAN to illustrate how it works:

$\texttt{SpMM\\_MEAN}(\begin{bmatrix}1&0\\\ 0&4\\\
5&6\end{bmatrix},\begin{bmatrix}7&8\\\ 9&10\\\
\end{bmatrix})=\small\begin{bmatrix}\frac{1}{2}(1\times 7+0\times
9)&\frac{1}{2}(1\times 8+0\times 10)\\\ \frac{1}{2}(0\times 7+4\times
9)&\frac{1}{2}(0\times 8+4\times 10)\\\ \frac{1}{2}(5\times 7+6\times
9)&\frac{1}{2}(5\times 8+6\times 10)\end{bmatrix},$

Equivalently, the SpMM_MEAN can also be expressed as:

$\texttt{SpMM\\_MEAN}({\bm{A}},{\bm{H}}^{(l)})={\bm{D}}^{-1}{\bm{A}}{\bm{H}}^{(l)},$

where ${\bm{D}}$ is the degree matrix of ${\bm{A}}$. Thus, although we did not
normalize the adjacency matrix in GraphSAGE, when applying the top-$k$
sampling to approximate SpMM_MEAN, the column norm of ${\bm{A}}_{:,j_{i}}$ is
actually $\frac{1}{\sqrt{\mathrm{Deg}_{j_{i}}}}$ due to the normalization.

Also, for GraphSAGE, the inputs to the first SpMM_MEAN operation are
${\bm{A}}$ and ${\bm{X}}$. They do not require gradient since they are not
trainable. Thus, the first SAGE layer is not presented in Figure 9 and Figure
7.

##  Appendix B Pseudo code of the greedy algorithm

In algorithm 1, here we provide the pseudo code of our greedy algorithm for
solving the constrained optimization problem. In Table 5, we show the run time
of the greedy algorithm, which is negligible compared to the acceleration
effect.

Table 5. The running time (second) of the greedy algorithm. | Reddit | Yelp |  | ogbn-
---
proteins
| ogbn-
---
products
GCN | 0.03 | 0.03 | 0.03 | 0.03
GraphSAGE | 0.02 | 0.02 | 0.03 | 0.03
GCNII | 0.05 | 0.05 | 0.06 | -

##  Appendix C Experimental Settings

###  C.1. Software and Hardware Descriptions

All experiments are conducted on a server with four NVIDIA 3090 GPUs, four AMD
EPYC 7282 CPUs, and 252GB host memory. We implement all models based on
Pytorch and Pytorch Geometric. During our experiments, we found that the
version of Pytorch, Pytorch Sparse, and Pytorch Scatter can significantly
impact the running speed of the baseline. Here we list the details of our used
packages in all experiments in Table 6.

Table 6. Package configurations of our experiments. Package | Version
---|---
CUDA | 11.1
pytorch_sparse | 0.6.12
pytorch_scatter | 2.0.8
pytorch_geometric | 1.7.2
pytorch | 1.9.0
OGB | 1.3.2

###  C.2. Statistics of benchmark datasets

The statistics for all used datasets are shown in Table 7. We follow the
standard data splits and all datasets are directly downloaded from Pytorch
Geometric or the protocol of OGB (Hu et al., 2020).

Table 7. Dataset Statistics.

Dataset | Task | Nodes | Edges | Classes | Label Rates
---|---|---|---|---|---
Reddit | multi-class | 232,965 | 11,606,919 | 41 | 65.86%
Yelp | multi-label | 716,847 | 6,977,409 | 100 | 75.00%
ogbn-proteins | binary-Class | 132,534 | 39,561,252 | 2 | 65.00%
ogbn-products | multi-class | 2,449,029 | 61,859,076 | 47 | 8.03%

Input: The gradients of node embeddings
$\\{\nabla{\bm{H}}^{(1)},\cdots\nabla{\bm{H}}^{(L)}\\}$, the adjacency matrix
${\bm{A}}$, the hidden dimension $\\{d_{1},\cdots d_{L}\\}$, the node set
$\mathcal{V}$, the edge set $\mathcal{E}$

Parameters : The step size $\alpha$, the overall budget $C$

Output: The layer-wise $\\{k_{1},\cdots k_{L}\\}$ associated with the top-$k$
sampling

1 $B\leftarrow\sum_{l=1}^{L}|\mathcal{E}|d_{l}$

2 $\forall
i,k_{l}\leftarrow|\mathcal{V}|,\mathrm{Top}_{k_{l}}\leftarrow\\{1,\cdots|\mathcal{V}|\\}$.

3 while  _$B\geq C\sum_{l=1}^{L}|\mathcal{E}|d_{l}$_ do

4 $m\leftarrow\operatorname*{arg\,min}_{l\in\\{1,\cdots
L\\}}(\sum_{i\in\mathrm{Top}_{k_{l}}}\|{\bm{A}}_{:,i}^{\top}\|_{2}\|\nabla{\bm{H}}^{(l+1)}_{i,:}\|_{2}-\sum_{i\in\mathrm{Top}_{k_{l}-\alpha|\mathcal{V}|}}\|{\bm{A}}_{:,i}^{\top}\|_{2}\|\nabla{\bm{H}}^{(l+1)}_{i,:}\|_{2})$

// Choose the layer $m$ to reduce by a step size $\alpha|\mathcal{V}|$, such
that the increment of errors is minimal.

5 $B\leftarrow B-d_{m}\sum_{i\in\mathrm{Top}_{k_{m}}\cap
i\notin\mathrm{Top}_{k_{m}-\alpha|\mathcal{V}|}}\\#nnz_{i}$

// Since we exclude some column-row pairs for layer $m$, here we reduce the
budget $B$ accordingly.

6 $k_{m}\leftarrow k_{m}-\alpha|\mathcal{V}|$

// Update $k_{m}$ accordingly.

7 $\mathrm{Top}_{k_{m}}\leftarrow$ the set of indices $i$ associated with
$k_{m}$ largest
$\|{\bm{A}}_{:,i}^{\top}\|_{2}\|\nabla{\bm{H}}^{(l+1)}_{i,:}\|_{2}$

// Update $\mathrm{Top}_{k_{m}}$ accordingly.

8

9 end while

return _$\\{k_{1},\cdots,k_{L}\\}$_

Algorithm 1 The greedy algorithm

Figure 9. The averaged degrees of nodes picked by top-$k$ sampling along the
whole training process, where the applied dataset is Reddit and overall budget
$C$ is set as 0.1.

###  C.3. Hyperparameter Settings

Regarding Reddit, and Yelp dataset, we follow the hyperparameter
configurations reported in the respective papers as closely as possible.
Regarding ogbn-proteins and ogbn-products dataset, we follow the
hyperparameter configurations and codebases provided on the OGB (Hu et al.,
2020) leader-board. Please refer to the OGB website for more details. The
optimizer is Adam for all these models. All methods terminate after a fixed
number of epochs. We report the test accuracy associated with the highest
validation score. Table 11 summarize the hyperparameter configuration of
GraphSAINT. Table 8, Table 9, and Table 10 summarize the hyperparameter
configuration of full-Batch GCN, full-Batch GraphSAGE, and full-batch GCNII,
respectively.

Table 8. Configuration of Full-Batch GCN.

Dataset | Training | Archtecture
---|---|---
| Learning
---
Rates
Epochs | Dropout | BatchNorm | Layers |  | Hidden
---
Dimension
Reddit | 0.01 | 400 | 0.5 | Yes | 3 | 256
Yelp | 0.01 | 500 | 0.1 | Yes | 3 | 512
| ogbn-
---
proteins
0.01 | 1000 | 0.5 | No | 3 | 256
| ogbn-
---
products
0.001 | 500 | 0.5 | No | 3 | 256

Table 9. Configuration of Full-Batch GraphSAGE.

Dataset | Training | Archtecture
---|---|---
| Learning
---
Rates
Epochs | Dropout | BatchNorm | Layers |  | Hidden
---
Dimension
Reddit | 0.01 | 400 | 0.5 | Yes | 3 | 256
Yelp | 0.01 | 500 | 0.1 | Yes | 3 | 512
| ogbn-
---
proteins
0.01 | 1000 | 0.5 | No | 3 | 256
| ogbn-
---
products
0.001 | 500 | 0.5 | No | 3 | 256

Table 10. Configuration of Full-Batch GCNII.

Dataset | Training | Archtecture
---|---|---
| Learning
---
Rates
Epochs | Dropout | BatchNorm | Layers |  | Hidden
---
Dimension
Reddit | 0.01 | 400 | 0.5 | Yes | 4 | 256
Yelp | 0.01 | 500 | 0.1 | Yes | 4 | 256
| ogbn-
---
proteins
0.01 | 1000 | 0.5 | No | 4 | 256

Table 11. Training configuration of GraphSAINT.

Dataset |  | RandomWalk
---
Sampler
Training | Archtecture
Walk length | Roots |  | Learning
---
Rates
Epochs | Dropout | BatchNorm | Layers |  | Hidden
---
Dimension
Reddit | 4 | 8000 | 0.01 | 40 | 0.1 | Yes | 3 | 128
Yelp | 2 | 8000 | 0.01 | 75 | 0.1 | Yes | 3 | 512
| ogbn-
---
products
3 | 60000 | 0.01 | 20 | 0.5 | No | 3 | 256


# Diverse Multi-Answer Retrieval with Determinantal Point Processes

Poojitha Nandigam∗, Nikhil Rayaprolu∗, Manish Shrivastava
Language Technologies Research Centre (LTRC)
International Institute of Information Technology, Hyderabad, India
<EMAIL_ADDRESS><EMAIL_ADDRESS>
<EMAIL_ADDRESS>
###### Abstract

Often questions provided to open-domain question answering systems are
ambiguous. Traditional QA systems that provide a single answer are incapable
of answering ambiguous questions since the question may be interpreted in
several ways and may have multiple distinct answers. In this paper, we address
multi-answer retrieval which entails retrieving passages that can capture
majority of the diverse answers to the question. We propose a re-ranking based
approach using Determinantal point processes utilizing BERT as kernels. Our
method jointly considers query-passage relevance and passage-passage
correlation to retrieve passages that are both query-relevant and diverse.
Results demonstrate that our re-ranking technique outperforms state-of-the-art
method on the AmbigQA dataset. **footnotetext: Equal contribution

##  1 Introduction

The objective of open-domain question answering is to provide answers to
queries utilising a large collection of documents from the World Wide Web,
Wikipedia etc. More than 50% of questions present in a widely used open-domain
QA dataset (Natural Questions Kwiatkowski et al. (2019)) comprise of ambiguous
questions (Min et al. (2020)). Questions that are ambiguous may be interpreted
in a number of ways and as a result, they need various answers. In this paper,
we concentrate on questions with multiple distinct answers.

Open domain question-answering systems are designed to generate answers from
several data sources. Since similar information can be present across multiple
data sources, it introduces a significant amount of redundancy. Traditional
open-domain QA (Chen et al. (2017)) systems comprise of a Retriever, which
retrieves passages relevant to the question. A passage retriever is primarily
concerned with retrieving passages that are relevant to the query, and it does
not address redundancy in the passages during retrieval. To be able to produce
diverse answers to the question, the passages retrieved must be both relevant
to the question and distinct from one another. After the retrieval stage, we
introduce a novel re-ranking approach to handle redundant passages. As a
result, the re-ranked passages would capture most of the diverse answers to
the question. In this paper, we investigate the multi-answer retrieval task,
which entails retrieving passages that can cover the distinct answers.

Re-ranking methods have been employed previously to improve the question
answering accuracy significantly(Wang et al. (2019);Nogueira and Cho (2019);
Min et al. (2021); Clark and Gardner (2017)). Min et al. (2021) tackles
diverse multi-answer retrieval by proposing a re-ranker based on an auto-
regressive framework in which each passage selected is dependent on the
passages chosen at a previous time step.

Determinantal Point Processes (DPP) (Kulesza and Taskar (2012)) are
probabilistic models that are effective at identifying diverse subsets of
elements from a collection while preserving quality. DPP methods have proven
effective in natural language processing tasks where there is a need for
diverseness. Cho et al. (2019), Li et al. (2019), and Cho et al. (2020),
Sharghi et al. (2017) have used DPPs to perform summarization by choosing
salient but also diverse items to be included in the summaries. In this paper,
we propose an unsupervised re-ranking technique for multi-answer retrieval
utilising Determinantal point processes and BERT to model the kernels.

Our contributions can be summarized as follows: 1) We propose a re-ranking
method based on determinantal point processes that focuses on diverse passage
retrieval.

2) Since our approach is unsupervised, our method does not require a large
amount of data unlike prior re-ranking methodsMin et al. (2021). Instead, we
rely on DPP to identify the most relevant passages to the question that are
distinct from one another.

3) We demonstrate that our technique outperforms the state-of-the-art method
on the AmbigQA dataset using $\mathrm{MRECALL}$ @ $k$ metrics.

##  2 Related Work

Many open domain question answering systems(Chen et al. (2017); Yang et al.
(2019); Izacard and Grave (2021); Guu et al. (2020); Lee et al. (2019)) adopt
the retriever-reader method by retrieving the relevant documents and later
applying neural techniques to predict the answer. The retriever-reader method
was first proposed by Chen et al. (2017). DrQA(Chen et al. (2017)) uses
Wikipedia as knowledge source and employs a sparse retrieval method using TF-
IDF and a recurrent neural network to identify the answer spans. While Yang et
al. (2019) adopts Anserini retriever(Yang et al. (2017)) using BM25 as the
ranking function and BERT model (Devlin et al. (2018)) as the reader. Sparse
retrieval based methods, such as TF-IDF and BM25, face challenges when
retrieving relevant passages that do not match the question’s exact terms.
Dense retrieval-based approaches, on the other hand, overcome this problem by
mapping each word into a vector space in which words with similar meanings
tend to be closer together. ORQA (Lee et al. (2019)) and DPR (Karpukhin et al.
(2020)) employ a question and passage encoder based on BERT and compute a
relevance score. Using this relevance score, the retriever retrieves the most
relevant documents from the corpus.

##  3 Determinantal Point Processes for Re-ranking

Re-ranker acts as a filter to pick a limited number of passages that can be
used as input to generate answers to the questions. We formulate the task of
passage re-ranking as a subset selection problem. Our objective is to choose a
subset of passages ($Y$) of size $k$ from the ground set $\mathcal{Y}$
comprising $\mathsf{N}$ passages that covers all of the answers to a given
question $q$ . DPP models a distribution on all the subsets of the ground set
$\mathcal{Y}$ jointly considering the quality and diversity. A subset $Y$ is
drawn according to the probability distribution $P$.

$P(Y;L)\varpropto det(L_{Y})$ (1) $P(Y;L)=\frac{\det(L_{Y})}{\det(L+I)}$ (2)

where $I$ is the identity matrix,
$L\in\mathbb{R}^{\mathsf{N}\times\mathsf{N}}$ is a positive semi-definite
matrix referred as $L$-ensemble, $\det(.)$ denotes the determinant of a
matrix, and $L_{y}$ is the submatrix of $L$ indexed by items in $Y$. $L$
matrix jointly considers query-passage relevance as well as passage-passage
correlation through eq. 3.

$L_{ij}=Q(i,q)\cdot S(i,j)\cdot Q(j,q)$ (3)

DPP focuses on two measures - quality and similarity ( Fig 1). Quality score
$Q(i,q)$ measures how salient the passage $i$ is and whether it contains an
answer to the question $q$. Similarity score $S(i,j)$ is computed between two
passages $i$ and $j$ to incorporate diversity in the passages. DPP assigns a
probability to a set Y proportional to the determinant of $L$-ensemble which
may be interpreted geometrically as the volume of the parallepiped covered by
the quality and similarity measures (Kulesza and Taskar (2012)). A diverse
passage subset occupies more volume than a subset of similar passages,
therefore DPP assigns higher probability to diverse and relevant passages
rather than the most relevant and similar passages. If passages are relevant
and diverse, then the passages can cover multiple distinct answers to the
question.

Figure 1: An overview of the proposed re-ranking method using DPP. A
similarity score between the passages and a quality score between the question
and passage are computed. These two scores are utilised to construct the DPP
kernel matrix.

###  3.1 BERT for Similarity matrix

To compute the similarity scores, we use a pretrained BERT model (Devlin et
al. (2018); Reimers and Gurevych (2019)) to generate embeddings for every
passage. The model takes the passage as input and produces a 768 dimensional
dense embedding. We use these embeddings to calculate the cosine similarity of
all passages and compute a similarity matrix
$S\in\mathbb{R}^{\mathsf{N}\times\mathsf{N}}$ for the whole passage set. All
the values in the similarity matrix lie in the range of $\lceil 0,1\rceil$. If
passages $i$ and $j$ are similar, the similarity value $S(i,j)$ lies closer to
$1$, if they are distinct, the value lies closer to $0$, and if $i=j$,
$S(i,j)$ becomes equal to 1.

$S(i,j)=cosine\\_sim(BERT_{A}(i),BERT_{A}(j))$ (4)

###  3.2 BERT for Quality matrix

We use a pretrained BERT model trained on MS MARCO (Nguyen et al. (2016)) for
computing the Quality matrix. The model takes in a query and a passage and
generates the quality score. Higher quality score indicates that the passage
is most relevant to the query and therefore most likely to answer the query.
Unlike for computing similarity matrix, we do not perform cosine similarity
over the model’s outputs to produce a score, instead, we use a BERT encoder
that concatenates both query and passage and generates a score. The quality
matrix $Q\in\mathbb{R}^{\mathsf{N}\times\mathsf{N}}$ is computed by performing
the matrix multiplication of the scores ($N\times 1$) with it’s transpose
resulting in $N\times N$ dimensioned vector . These quality scores are then
normalized to lie between $\lceil 0,1\rceil$.

$Q(i,j)=Norm(BERT_{B}([i;j]))$ (5)

###  3.3 Sampling

Traditional DPP sampling algorithms have higher run-time complexity when $L$
matrix is large. We apply an efficient sampling technique - BFGMInference (Li
et al. (2019);Chen et al. (2018)). BFGMInference approximates a greedy
approach to select a passage that maximizes the $det(L_{Y})$ and adds it to
the passage subset.

$\begin{array}[]{c}f(Y)=\log\operatorname{det}\left(L_{Y}\right)\\\
k=\underset{i\in\mathcal{Y}\backslash
Y}{\arg\max}f(Y\cup\\{i\\})-f(Y)\end{array}$ (6)

Models | Top5 | Top 10
---|---|---
| AmbigQA-Dev | AmbigQA-Dev
DPR+(Min et al. (2021)) | 55.2/36.3 | 59.3/39.6
DPR+ \+ Nogueira and Cho (2019) | 63.4/43.1 | 65.8/46.4
JPR(Min et al. (2021)) | 64.8/45.2 | 67.1/48.2
QRR | 62.0/42.3 | 70.8/57.6
DPP-R |  66.9/53.5 | 72.8/58.8
Table 1: Performance of various models on AmbigQA dataset. Each row contains
the $\mathrm{MRECALL}$ @ $k$ metrics for single answer retrieval and multi-
answer retrieval respectively.

##  4 Experiments

In this section, we discuss about the passage retrieval method, the dataset we
used in our experiments, the evaluation metric, and the results of our
experiments.

###  4.1 Passage retrieval

Wikipedia is utilised as the corpus for retrieving passages for the questions.
Each Wikipedia article is broken into multiple passages containing the same
number of words. We retrieve query-relevant passages from Wikipedia using the
Dense Passage Retriever (DPR) (Karpukhin et al. (2020); Lin et al. (2021)).
DPR computes encodings for all passages extracted from the Wikipedia corpus
and builds an index. The inner product of the query and passage encodings is
used to determine the similarity scores between them. Passages with the
highest scores are the ones that are most relevant to the query, and these
passages are subsequently sent into the re-ranker as input.

###  4.2 Dataset

We evaluated our method on an open-domain question-answering dataset AmbigQA
(Min et al. (2020)), which contains multiple-answer questions. The dataset was
created from an anonymised collection of Google search queries submitted by
users seeking information on different subjects. It consists of 14,042
question-answer pairs derived from the Natural questions dataset (Kwiatkowski
et al. (2019)) and is split into train, validation, and test sets. Train set
consists of 10,036 question-answer pairs, validation set consists of 2,002
examples, and test set consists of 4,042 examples.

###  4.3 Evaluation metric

$\mathrm{MRECALL}$ @ $k$ (Min et al. (2021)) is used to evaluate the re-
ranking of passages for questions with diverse answers. As per this metric, if
a query has $n$ answers, the $k$ passages that are retrieved must cover all of
the answers. If $n<=k$, all answers must be covered; if $n>k$, the passages
retrieved must contain at least $k$ answers. A retrieval is deemed successful
if the passages retrieved include all or at least $k$ of the answers to the
query.

###  4.4 Results

We compare our technique to a few additional baselines, all of which were
assessed using the $\mathrm{MRECALL}$ @ $k$ metric on the AmbigQA dataset.

  * •

DPR+ Min et al. (2021) integrates REALM (Guu et al. (2020)) with DPR
(Karpukhin et al. (2020)). As described in Section 4.1, DPR is a dense
retrieval based technique that utilizes the FAISS library to retrieve the
relevant documents. Encoders for the query and passage are initialized using
REALM and the DPR training method is followed.

  * •

DPR+ \+ Nogueira and Cho (2019) employs DPR+ for the first stage of retrieval
and the re-ranking method in Nogueira and Cho (2019) is applied on the
retrieved passages.

  * •

JPR Min et al. (2021) employs DPR+ as the initial ranker and an auto-
regressive framework is adopted as a re-ranker to generate diverse passages.

  * •

Query Relevance Re-ranking(QRR) In this method, we first calculate the quality
scores for each passage (described in section 3.2) and then we sort the
passages based on these scores to pick the top-$k$ passages. Here, similarity
among the passages is not considered.

  * •

DPP-R We employ our method described in section 3 to retrieve highly diverse
and relevant passages.

We calculate the performance of diverse multi passage retrieval using the
$\mathrm{MRECALL}$ @ $k$ measure described in section 4.3. Evaluation on the
AmbigQA dataset demonstrates that our approach outperforms existing re-ranking
techniques. Our technique requires no human annotations for multi passage re-
ranking while outperforming existing methods, as shown in Table 1. DPP is
modelled to select a subset of high-quality and diverse passages, which
contributes to the success of our method for this task. Experiments
demonstrate that the DPP-based technique achieves promising results for
retrieving passages containing diverse answers.

##  5 Discussion

Impact on QA system’s performance:  An Open domain question answering system’s
pipeline consists of three stages. 1) Retrieval 2) Re-ranking followed by 3)
Answer extraction. Improvements in any of these stages significantly improve
the overall system’s ability to answer a question. Nogueira and Cho (2019),
Min et al. (2021) have shown that the use of a re-ranker has led to end-to-end
QA improvements. Based on the results presented in Table 1, the DPP method
enhances re-ranking for both single and multi-answer questions. We believe
that this improvement in re-ranking will also improve the overall performance
of the end-to-end QA system.

Impact of diversity:  DPP-R and JPR retrieve diverse passages utilising DPP
and auto-regressive framework, respectively. Other approaches like QRR,
retrieve just passages that are relevant to the query and do not tackle
passage redundancy. We observe that our approach using DPP performs better
than the QRR method. In order to re-rank, QRR simply considers how relevant a
passage is to the query, and it retrieves the top-k passages with the highest
relevance score for a given query. On the other hand, DPP-R takes into account
how relevant the passage is to the query and also how similar passages are to
each other, in order to eliminate redundant passages leading to diversity in
the retrieved passages. DPP-R and JPR outperform other methods that do not
emphasise diversity in multi-answer retrieval. For single answer retrieval,
DPP-R and JPR have fared better than other methods, with the minor exception
that QRR beats JPR in top-10 re-ranking. This demonstrates that diversity is
an important aspect to consider during the re-ranking stage.

##  6 Conclusion

In this paper, we propose a DPP-based approach to improve the diverseness of
the retrieved passages. We compare our method to the state-of-the-art method
and outperform it by $3\%$ (top $5$), $8\%$ (top $10$) for single-answer
questions, and $18\%$ (top$5$) and $21\%$ (top$10$) for multi-answer retrieval
on AmbigQA dataset.

## References

  * Chen et al. (2017) Danqi Chen, Adam Fisch, Jason Weston, and Antoine Bordes. 2017.  Reading Wikipedia to answer open-domain questions.  In _ACL 2017 - 55th Annual Meeting of the Association for Computational Linguistics, Proceedings of the Conference (Long Papers)_ , volume 1, pages 1870–1879.
  * Chen et al. (2018) Laming Chen, Guoxin Zhang, and Eric Zhou. 2018.  Fast greedy map inference for determinantal point process to improve recommendation diversity.  _Advances in Neural Information Processing Systems_ , 31.
  * Cho et al. (2020) Sangwoo Cho, Logan Lebanoff, Hassan Foroosh, and Fei Liu. 2020.  Improving the similarity measure of determinantal point processes for extractive multi-document summarization.  _ACL 2019 - 57th Annual Meeting of the Association for Computational Linguistics, Proceedings of the Conference_ , pages 1027–1038.
  * Cho et al. (2019) Sangwoo Cho, Chen Li, Dong Yu, Hassan Foroosh, and Fei Liu. 2019.  Multi-document summarization with determinantal point processes and contextualized representations.  _arXiv_.
  * Clark and Gardner (2017) Christopher Clark and Matt Gardner. 2017.  Simple and effective multi-paragraph reading comprehension.  _arXiv preprint arXiv:1710.10723_.
  * Devlin et al. (2018) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2018.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_.
  * Guu et al. (2020) Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Ming Wei Chang. 2020.  REALM: Retrieval-Augmented language model pre-training.  In _37th International Conference on Machine Learning, ICML 2020_ , volume PartF16814, pages 3887–3896.
  * Izacard and Grave (2021) Gautier Izacard and Edouard Grave. 2021.  Leveraging passage retrieval with generative models for open domain question answering.  In _EACL 2021 - 16th Conference of the European Chapter of the Association for Computational Linguistics, Proceedings of the Conference_ , pages 874–880.
  * Karpukhin et al. (2020) Vladimir Karpukhin, Barlas Oğuz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen tau Yih. 2020.  Dense passage retrieval for open-domain question answering.
  * Kulesza and Taskar (2012) Alex Kulesza and Ben Taskar. 2012.  Determinantal point processes for machine learning.
  * Kwiatkowski et al. (2019) Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Matthew Kelcey, Jacob Devlin, Kenton Lee, Kristina N. Toutanova, Llion Jones, Ming-Wei Chang, Andrew Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. 2019.  Natural questions: a benchmark for question answering research.  _Transactions of the Association of Computational Linguistics_.
  * Lee et al. (2019) Kenton Lee, Ming-Wei Chang, and Kristina Toutanova. 2019.  Latent retrieval for weakly supervised open domain question answering.  _arXiv preprint arXiv:1906.00300_.
  * Li et al. (2019) Lei Li, Wei Liu, Marina Litvak, Natalia Vanetik, and Zuying Huang. 2019.  In conclusion not repetition: Comprehensive abstractive summarization with diversified attention based on determinantal point processes.  _CoNLL 2019 - 23rd Conference on Computational Natural Language Learning, Proceedings of the Conference_ , pages 822–832.
  * Lin et al. (2021) Jimmy Lin, Xueguang Ma, Sheng-Chieh Lin, Jheng-Hong Yang, Ronak Pradeep, and Rodrigo Nogueira. 2021.  Pyserini: An easy-to-use python toolkit to support replicable ir research with sparse and dense representations.  _arXiv preprint arXiv:2102.10073_.
  * Min et al. (2021) Sewon Min, Kenton Lee, Ming-Wei Chang, Kristina Toutanova, and Hannaneh Hajishirzi. 2021.  Joint passage ranking for diverse multi-answer retrieval.  _arXiv preprint arXiv:2104.08445_.
  * Min et al. (2020) Sewon Min, Julian Michael, Hannaneh Hajishirzi, and Luke Zettlemoyer. 2020.  Ambigqa: Answering ambiguous open-domain questions.
  * Nguyen et al. (2016) Tri Nguyen, Mir Rosenberg, Xia Song, Jianfeng Gao, Saurabh Tiwary, Rangan Majumder, and Li Deng. 2016.  Ms marco: A human generated machine reading comprehension dataset.  In _CoCo@ NIPS_.
  * Nogueira and Cho (2019) Rodrigo Nogueira and Kyunghyun Cho. 2019.  Passage re-ranking with bert.  _arXiv preprint arXiv:1901.04085_.
  * Reimers and Gurevych (2019) Nils Reimers and Iryna Gurevych. 2019.  Sentence-bert: Sentence embeddings using siamese bert-networks.  In _Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing_. Association for Computational Linguistics.
  * Sharghi et al. (2017) Aidean Sharghi, Jacob S Laurel, and Boqing Gong. 2017.  Query-focused video summarization: Dataset, evaluation, and a memory network based approach.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 4788–4797.
  * Wang et al. (2019) Zhiguo Wang, Patrick Ng, Xiaofei Ma, Ramesh Nallapati, and Bing Xiang. 2019.  Multi-passage bert: A globally normalized bert model for open-domain question answering.  _arXiv preprint arXiv:1908.08167_.
  * Yang et al. (2017) Peilin Yang, Hui Fang, and Jimmy Lin. 2017.  Anserini: Enabling the use of lucene for information retrieval research.  In _Proceedings of the 40th International ACM SIGIR Conference on Research and Development in Information Retrieval_ , pages 1253–1256.
  * Yang et al. (2019) Wei Yang, Yuqing Xie, Aileen Lin, Xingyu Li, Luchen Tan, Kun Xiong, Ming Li, and Jimmy Lin. 2019.  End-to-end open-domain question answering with BERTserini.  In _NAACL HLT 2019 - 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies - Proceedings of the Demonstrations Session_ , pages 72–77.