Optimizing sparse graph neural networks for dense hardware
A computer-implemented method for computing node embeddings of a sparse graph that is an input of a sparse graph neural network is described. Each node embedding corresponds to a respective node of the sparse graph and represents feature information of the respective node and a plurality of neighboring nodes of the respective node. The method includes: receiving an adjacency matrix that represents edges of the sparse graph; receiving a weight matrix representing, for each node of the sparse graph, a level of influence of respective neighboring nodes on the node; initializing, for each node of the sparse graph, a respective node embedding; transforming the adjacency matrix into a low-bandwidth adjacency matrix, and performing the following operations at least once: generating a message propagation matrix as a product of the low-bandwidth adjacency matrix, the node embeddings of the nodes, and the weight matrix, wherein the message propagation matrix represents message propagation among the nodes of the sparse graph, and updating the node embeddings of the sparse graph by processing the message propagation matrix and the node embeddings of the nodes using an encoder neural network of the sparse graph neural network.
1. A computer-implemented method for updating node embeddings of a sparse graph that is an input of a sparse graph neural network, wherein each node embedding corresponds to a respective node of the sparse graph and represents feature information of the respective node and a plurality of neighboring nodes of the respective node, the method comprising:
receiving an adjacency matrix that represents edges of the sparse graph;
receiving a weight matrix representing, for each node of the sparse graph, a level of influence of respective neighboring nodes on the node;
initializing, for each node of the sparse graph, a respective node embedding;
transforming the adjacency matrix into a low-bandwidth adjacency matrix, and
performing the following operations at least once:
generating a message propagation matrix as a product of the low-bandwidth adjacency matrix and the node embeddings of the nodes, wherein the message propagation matrix represents message propagation among the nodes of the sparse graph, and
updating the node embeddings of the sparse graph by processing the message propagation matrix, the node embeddings of the nodes, and the weight matrix using an encoder neural network of the sparse graph neural network.
2. The method of claim 1 , wherein transforming the adjacency matrix into the low-bandwidth adjacency matrix comprises permuting the nodes in the sparse graph to obtain the low-bandwidth adjacency matrix having a bandwidth B, wherein B is a non-negative integer, and wherein all non-zero elements of the low-bandwidth adjacency matrix lie at a distance not greater than B from a main diagonal of the low-bandwidth adjacency matrix.
3. The method of claim 1 , wherein the low-bandwidth adjacency matrix includes diagonal blocks, upper-triangular blocks, lower-triangular blocks, and zero elements outside of the diagonal, upper-triangular, and lower-triangular blocks.
4. The method of claim 3 , wherein permuting the nodes in the sparse graph includes permuting the nodes using a heuristic algorithm.
5. The method of claim 4 , wherein the heuristic algorithm is Reverse Cuthill McKee (RCMK) algorithm.
6. The method of claim 3 , wherein generating the message propagation matrix as the product of the low-bandwidth adjacency matrix and the current node embeddings comprises: loading the low-bandwidth adjacency matrix and the current node embeddings to a dense hardware device for dense matrix computations.
7. The method of claim 6 , wherein generating the message propagation matrix as the product of the low-bandwidth adjacency matrix and the node embeddings comprises: performing a plurality of dense matrix multiplications on the dense hardware device.
8. The method of claim 7 , wherein at least one of the plurality of dense matrix multiplications performed on the dense hardware device is a multiplication of a diagonal block in the low-bandwidth adjacency matrix with a corresponding sub-matrix in the node embeddings.
9. The method of claim 7 , wherein at least one of the plurality of dense matrix multiplications performed on the dense hardware device is a multiplication of an upper-triangular block in the low-bandwidth adjacency matrix with a corresponding sub-matrix in the node embeddings.
10. The method of claim 7 , wherein at least one of the plurality of dense matrix multiplications performed on the dense hardware device is a multiplication of a lower-triangular block in the low-bandwidth adjacency matrix with a corresponding sub-matrix in the node embeddings.
11. The method of claim 1 , where in the encoder neural network is a recurrent neural network.
12. The method of claim 11 , wherein the recurrent neural network is a Gated Recurrent Unit (GRU).
13. The method of claim 1 , further comprises: using the updated node embeddings of the sparse graph to train the sparse graph neural network on a machine learning task.
14. The method of claim 13 , wherein the training of the sparse graph neural network is performed on a dense hardware device.
15. A system comprising one or more computers and one or more storage devices storing instructions that, when executed by the one or more computers, cause the one or more computers to perform operations for updating node embeddings of a sparse graph that is an input of a sparse graph neural network, wherein each node embedding corresponds to a respective node of the sparse graph and represents feature information of the respective node and a plurality of neighboring nodes of the respective node, the operations comprising:
receiving an adjacency matrix that represents edges of the sparse graph;
receiving a weight matrix representing, for each node of the sparse graph, a level of influence of respective neighboring nodes on the node;
initializing, for each node of the sparse graph, a respective node embedding;
transforming the adjacency matrix into a low-bandwidth adjacency matrix, and
performing the following operations at least once:
generating a message propagation matrix as a product of the low-bandwidth adjacency matrix and the node embeddings of the nodes, wherein the message propagation matrix represents message propagation among the nodes of the sparse graph, and
updating the node embeddings of the sparse graph by processing the message propagation matrix, the node embeddings of the nodes, and the weight matrix using an encoder neural network of the sparse graph neural network.
16. One or more non-transitory computer-readable storage media storing instructions that, when executed by one or more computers, cause the one or more computers to perform operations for updating node embeddings of a sparse graph that is an input of a sparse graph neural network, wherein each node embedding corresponds to a respective node of the sparse graph and represents feature information of the respective node and a plurality of neighboring nodes of the respective node, the operations comprising:
receiving an adjacency matrix that represents edges of the sparse graph;
receiving a weight matrix representing, for each node of the sparse graph, a level of influence of respective neighboring nodes on the node;
initializing, for each node of the sparse graph, a respective node embedding;
transforming the adjacency matrix into a low-bandwidth adjacency matrix, and
performing the following operations at least once:
generating a message propagation matrix as a product of the low-bandwidth adjacency matrix and the node embeddings of the nodes, wherein the message propagation matrix represents message propagation among the nodes of the sparse graph, and
updating the node embeddings of the sparse graph by processing the message propagation matrix, the node embeddings of the nodes, and the weight matrix using an encoder neural network of the sparse graph neural network.
17. The one or more non-transitory computer-readable storage media of claim 16 , wherein the operations for transforming the adjacency matrix into the low-bandwidth adjacency matrix comprise: permuting the nodes in the sparse graph to obtain the low-bandwidth adjacency matrix having a bandwidth B, wherein B is a non-negative integer, and wherein all non-zero elements of the low-bandwidth adjacency matrix lie at a distance not greater than B from a main diagonal of the low-bandwidth adjacency matrix.
18. The one or more non-transitory computer-readable storage media of claim 16 , wherein the low-bandwidth adjacency matrix includes diagonal blocks, upper-triangular blocks, lower-triangular blocks, and zero elements outside of the diagonal, upper-triangular, and lower-triangular blocks.
19. The one or more non-transitory computer-readable storage media of claim 16 , wherein the operations for generating the message propagation matrix as the product of the low-bandwidth adjacency matrix and the node embeddings comprises: performing a plurality of dense matrix multiplications on a dense hardware device.
20. The one or more non-transitory computer-readable storage media of claim 19 , wherein at least one of the plurality of dense matrix multiplications performed on the dense hardware device is a multiplication of a diagonal block in the low-bandwidth adjacency matrix with a corresponding sub-matrix in the node embeddings.
21. The one or more non-transitory computer-readable storage media of claim 19 , wherein at least one of the plurality of dense matrix multiplications performed on the dense hardware device is a multiplication of an upper-triangular block in the low-bandwidth adjacency matrix with a corresponding sub-matrix in the node embeddings.
22. The one or more non-transitory computer-readable storage media of claim 19 , wherein at least one of the plurality of dense matrix multiplications performed on the dense hardware device is a multiplication of a lower-triangular block in the low-bandwidth adjacency matrix with a corresponding sub-matrix in the node embeddings.
23. The one or more non-transitory computer-readable storage media of claim 16 , wherein the operations further comprise: using the updated node embeddings of the sparse graph to train the sparse graph neural network on a machine learning task.
24. The one or more non-transitory computer-readable storage media of claim 23 , wherein the training of the sparse graph neural network is performed on a dense hardware device.