Technical Papers Archive

High-Performance and Programmable Attentional Graph Neural Networks with Global Tensor Formulations

Authors: Maciej Besta (ETH Zurich); Paweł Renc (AGH-UST, Sano Centre for Computational Medicine); Robert Gerstenberger (ETH Zurich); Paolo Sylos Labini (Free University of Bozen-Bolzano, ETH Zurich); Alexandros Ziogas, Tiancheng Chen, Lukas Gianinazzi, Florian Scheidl, Kalman Szenes, Armon Carigiet, and Patrick Iff (ETH Zurich); Grzegorz Kwasniewski (NextSilicon); Raghavendra Kanakagiri (University of Illinois); Chio Ge and Sammy Jaeger (ETH Zurich); Jarosław Wąs (AGH-UST); Flavio Vella (University of Trento); and Torsten Hoefler (ETH Zurich)

Abstract: Graph attention models (A-GNNs), a type of Graph Neural Networks (GNNs), have been shown to be more powerful than simpler convolutional GNNs (C-GNNs). However, A-GNNs are more complex to program and difficult to scale. To address this, we develop a novel mathematical formulation, based on tensors that group all the feature vectors, targeting both training and inference of A-GNNs The formulation enables straightforward adoption of communication-minimizing routines, it fosters optimizations such as vectorization, and it enables seamless integration with established linear algebra DSLs or libraries such as GraphBLAS. Our implementation uses a data redistribution scheme explicitly developed for sparse-dense tensor operations used heavily in GNNs, and fusing optimizations that further minimize memory usage and communication cost. We ensure theoretical asymptotic reductions in communicated data compared to the established message-passing GNN paradigm. Finally, we provide excellent scalability and speedups of >5x over modern libraries such as Deep Graph Library.




Back to Technical Papers Archive Listing