Loading [MathJax]/extensions/MathMenu.js
TT-GNN: Efficient On-Chip Graph Neural Network Training via Embedding Reformation and Hardware Optimization | IEEE Conference Publication | IEEE Xplore

TT-GNN: Efficient On-Chip Graph Neural Network Training via Embedding Reformation and Hardware Optimization


Abstract:

Training Graph Neural Networks on large graphs is challenging due to the need to store graph data and move them along the memory hierarchy. In this work, we tackle this b...Show More

Abstract:

Training Graph Neural Networks on large graphs is challenging due to the need to store graph data and move them along the memory hierarchy. In this work, we tackle this by effectively compressing graph embedding matrix such that the model training can be fully enabled with on-chip compute and memory resources. Specifically, we leverage the graph homophily property and consider using Tensor-train to represent the graph embedding. This allows nodes with similar neighborhoods to partially share the feature representation.While applying Tensor-train reduces the size of the graph embedding, it imposes several challenges to hardware design. On one hand, utilizing low-rank representation requires the features to be decompressed before being sent to GNN models, which introduces extra computation overhead. On the other hand, the decompressed features might still exceed on-chip memory capacity even with the minibatch setting, causing inefficient off-chip memory access. Thus, we propose the TT-GNN hardware accelerator with a specialized dataflow tailored for on-chip Tensor-train GNN learning. Based on the on-chip memory capacity and training configuration, TT-GNN adaptively breaks down a minibatch into smaller microbatches that can be fitted on-chip. The microbatch composition and scheduling order are designed to maximize data reuse and reduce redundant computations both across and within microbatches. To mitigate TT computation overhead, we further propose a unified algorithm to jointly handle TT decompression during forward propagation and TT gradient derivation during backward propagation. Evaluated on a series of benchmarks, the proposed software-hardware solution is able to outperform existing CPU-GPU training systems on both training performance (1.55~4210×) and energy efficiency (2.83~2254×). We believe TT-GNN introduces a new perspective to address large-scale GNN training and enables possibilities to train GNN models even under a significantly constrained resource bud...
Date of Conference: 28 October 2023 - 01 November 2023
Date Added to IEEE Xplore: 06 February 2024
Print on Demand(PoD) ISBN:979-8-3503-3056-4
Conference Location: Toronto, ON, Canada

1 INTRODUCTION

Originating from spectral graph analysis and fueled by the success of machine learning, graph neural networks (GNNs) have drawn a surge of interest and have been applied to various applications involving non-Euclidean graph-structured data. During the past few years, a wide range of GNN models [3], [10], [12], [38] have been proposed to solve graph-related problems. Exciting progress has been achieved by GNNs in domains such as recommendation systems [41], relation prediction [7], chemistry analysis [45], financial security [49], protein discovery [9], [36], EDA [16], [26], [27] and so on.

References

References is not available for this document.