Nimble GNN Embedding with Tensor-Train Decomposition

TL;DR

This paper introduces a tensor-train decomposition method to significantly compress GNN embedding tables, enabling efficient GPU training and maintaining or improving accuracy on large graph datasets.

Contribution

The paper presents a novel tensor-train based approach for compactly representing GNN embeddings, reducing memory requirements and improving training efficiency on GPUs.

Findings

Embedding size reduced by up to 1,659 times

Achieved comparable or better accuracy without node features

Significant speedups on multi-GPU systems

Abstract

This paper describes a new method for representing embedding tables of graph neural networks (GNNs) more compactly via tensor-train (TT) decomposition. We consider the scenario where (a) the graph data that lack node features, thereby requiring the learning of embeddings during training; and (b) we wish to exploit GPU platforms, where smaller tables are needed to reduce host-to-GPU communication even for large-memory GPUs. The use of TT enables a compact parameterization of the embedding, rendering it small enough to fit entirely on modern GPUs even for massive graphs. When combined with judicious schemes for initialization and hierarchical graph partitioning, this approach can reduce the size of node embedding vectors by 1,659 times to 81,362 times on large publicly available benchmark datasets, achieving comparable or better accuracy and significant speedups on multi-GPU systems. In some cases, our model without explicit node features on input can even match the accuracy of models that use node features.