GNNAutoScale: Scalable and Expressive Graph Neural Networks via Historical Embeddings

Matthias Fey, Jan E. Lenssen, Frank Weichert, Jure Leskovec
Proceedings of the 38th International Conference on Machine Learning, PMLR 139:3294-3304, 2021.

Abstract

We present GNNAutoScale (GAS), a framework for scaling arbitrary message-passing GNNs to large graphs. GAS prunes entire sub-trees of the computation graph by utilizing historical embeddings from prior training iterations, leading to constant GPU memory consumption in respect to input node size without dropping any data. While existing solutions weaken the expressive power of message passing due to sub-sampling of edges or non-trainable propagations, our approach is provably able to maintain the expressive power of the original GNN. We achieve this by providing approximation error bounds of historical embeddings and show how to tighten them in practice. Empirically, we show that the practical realization of our framework, PyGAS, an easy-to-use extension for PyTorch Geometric, is both fast and memory-efficient, learns expressive node representations, closely resembles the performance of their non-scaling counterparts, and reaches state-of-the-art performance on large-scale graphs.

Cite this Paper

BibTeX
@InProceedings{pmlr-v139-fey21a, title = {GNNAutoScale: Scalable and Expressive Graph Neural Networks via Historical Embeddings}, author = {Fey, Matthias and Lenssen, Jan E. and Weichert, Frank and Leskovec, Jure}, booktitle = {Proceedings of the 38th International Conference on Machine Learning}, pages = {3294--3304}, year = {2021}, editor = {Meila, Marina and Zhang, Tong}, volume = {139}, series = {Proceedings of Machine Learning Research}, month = {18--24 Jul}, publisher = {PMLR}, pdf = {http://proceedings.mlr.press/v139/fey21a/fey21a.pdf}, url = {https://proceedings.mlr.press/v139/fey21a.html}, abstract = {We present GNNAutoScale (GAS), a framework for scaling arbitrary message-passing GNNs to large graphs. GAS prunes entire sub-trees of the computation graph by utilizing historical embeddings from prior training iterations, leading to constant GPU memory consumption in respect to input node size without dropping any data. While existing solutions weaken the expressive power of message passing due to sub-sampling of edges or non-trainable propagations, our approach is provably able to maintain the expressive power of the original GNN. We achieve this by providing approximation error bounds of historical embeddings and show how to tighten them in practice. Empirically, we show that the practical realization of our framework, PyGAS, an easy-to-use extension for PyTorch Geometric, is both fast and memory-efficient, learns expressive node representations, closely resembles the performance of their non-scaling counterparts, and reaches state-of-the-art performance on large-scale graphs.} }
Endnote
%0 Conference Paper %T GNNAutoScale: Scalable and Expressive Graph Neural Networks via Historical Embeddings %A Matthias Fey %A Jan E. Lenssen %A Frank Weichert %A Jure Leskovec %B Proceedings of the 38th International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2021 %E Marina Meila %E Tong Zhang %F pmlr-v139-fey21a %I PMLR %P 3294--3304 %U https://proceedings.mlr.press/v139/fey21a.html %V 139 %X We present GNNAutoScale (GAS), a framework for scaling arbitrary message-passing GNNs to large graphs. GAS prunes entire sub-trees of the computation graph by utilizing historical embeddings from prior training iterations, leading to constant GPU memory consumption in respect to input node size without dropping any data. While existing solutions weaken the expressive power of message passing due to sub-sampling of edges or non-trainable propagations, our approach is provably able to maintain the expressive power of the original GNN. We achieve this by providing approximation error bounds of historical embeddings and show how to tighten them in practice. Empirically, we show that the practical realization of our framework, PyGAS, an easy-to-use extension for PyTorch Geometric, is both fast and memory-efficient, learns expressive node representations, closely resembles the performance of their non-scaling counterparts, and reaches state-of-the-art performance on large-scale graphs.
APA
Fey, M., Lenssen, J.E., Weichert, F. & Leskovec, J.. (2021). GNNAutoScale: Scalable and Expressive Graph Neural Networks via Historical Embeddings. Proceedings of the 38th International Conference on Machine Learning, in Proceedings of Machine Learning Research 139:3294-3304 Available from https://proceedings.mlr.press/v139/fey21a.html.

Related Material