Optimizing sparse graph neural networks for dense hardware

Abstract

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.